Language Model to get only the score of predefined tokens
# 🚀 Feature request

Hey, 

In Masked Language Model such as [TFDistilBertForMaskedLM](https://huggingface.co/transformers/model_doc/distilbert.html#tfdistilbertformaskedlm). The example explains well how to get the scores for [MASK] token. Nevertheless I did not find anything related to scale it up to 50k texts and especially get the results for a list of predefined tokens (let's say 10 tokens like in [MNLI examples for ZSL](https://huggingface.co/joeddav/xlm-roberta-large-xnli)).

The feature would be to add a _only_ids_ parameter in the `call` function of `TFDistilBertForMaskedLM` to get at the end for each sentence only the scores of the predefines tokens we are interested in. 

Maybe I did not search at the right place, however I managed to find something by digging in source code of HuggingFace. 

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->

## Motivation

Well my motivation came from two things: 

I was trying to implement my own pipeline of the [PET Algorithm from Schik et al.](https://arxiv.org/pdf/2001.07676.pdf). And you need to perform weak labeling on a dataset of ~10k texts a few times (as much time as you have PVP couples). 

So first, I tried to code with the [TFDistilBertForMaskedLM](https://huggingface.co/transformers/model_doc/distilbert.html#tfdistilbertformaskedlm) class in order to do label the data. But came the second motivation: 

```python
ResourceExhaustedError: OOM when allocating tensor
```

This errors occured on colab and also in [Lambda GPU](https://lambdalabs.com) (gpu.1x.rtx6000) whenever I tried to do it for more than 1k sentences. I dug a bit and the errors seemed to come from the `call` function of the class `TFDistilBertLMHead`, because it manipulates vectors of shape `[bs, dim, config.voc_size]`. 

As I did not find anything on the issues or StackOverflow, I decided to get into the code of `modeling_tf_distilbert.py` to try to change a bit the `call` functions and try my code on [Lambda GPU](https://lambdalabs.com) (gpu.1x.rtx6000) to look for improvements (I used the test set of yahoo question dataset).

## Your contribution

I changed a couple of things (not complicated things) in the classes and compared it (results below).

First change (easy) - I would call it CS for cutting out the sequence:

What we do is before applying the `vocab_projector` method we reduce the focus only on the mask token. It enables to manipulates tensors of shape (bs, 1, voc_size) instead of tensors of shape (bs, max_length, voc_size). 

```python
class CB_TFDistilBertForMaskedLM(TFDistilBertForMaskedLM):

    def call(
        self,
        inputs=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        labels=None,
        training=False,
        mask_only=False #If you just focus on the mask token id change the False by it (103 here)
    ):
        r"""
        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
        """
        
        return_dict = return_dict if return_dict is not None else self.distilbert.return_dict
        if isinstance(inputs, (tuple, list)):
            labels = inputs[7] if len(inputs) > 7 else labels
            if len(inputs) > 7:
                inputs = inputs[:7]
        elif isinstance(inputs, (dict, BatchEncoding)):
            labels = inputs.pop("labels", labels)

        distilbert_output = self.distilbert(
            inputs,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        hidden_states = distilbert_output[0]  # (bs, seq_length, dim)
        prediction_logits = self.vocab_transform(hidden_states)  # (bs, seq_length, dim)
        prediction_logits = self.act(prediction_logits)  # (bs, seq_length, dim)
        prediction_logits = self.vocab_layer_norm(prediction_logits)  # (bs, seq_length, dim)

        #transform from (bs, seq_length, dim) to (bs, 1, dim) where 1 is the MASK
        if mask_only:
            indexes = tf.where(inputs['input_ids'] == mask_only).numpy()
            prediction_logits = get_mask_rep(prediction_logits, indexes)
            
        prediction_logits = self.vocab_projector(prediction_logits, ids_only=ids_only)

        loss = None if labels is None else self.compute_loss(labels, prediction_logits)

        if not return_dict:
            output = (prediction_logits,) + distilbert_output[1:]
            return ((loss,) + output) if loss is not None else output

        return TFMaskedLMOutput(
            loss=loss,
            logits=prediction_logits,
            hidden_states=distilbert_output.hidden_states,
            attentions=distilbert_output.attentions
        )
    
def get_mask_rep(score, indexes):
    """
    Objective: get the score only for the indexes of interest.
    
    Inputs:
        - score, np.array or tf.tensor: the scores for all vocabulary
        - indexes, np.array: the array indicating [MASK] positionning
    Outputs:
        - final_scores, np.array: the final scores but only for the tokens of interest
    """
    n = score.shape[0]
    
    embedded_mask = dict(zip(indexes[:, 0], indexes[:, 1]))
    
    final_scores = []
    
    for index in range(score.shape[0]):
        final_scores.append(score[index, embedded_mask.get(index), :].numpy().reshape(1, -1))
        
    return np.array(final_scores)
```

Second change - I would call it RVCS for reducing the vocabulary tensor and cuting the sequence:

The idea is to reduce the tensor of `_linear` method in the _TFEmbeddings_ to manipulates tensors of size (dim, n_tokens_of_interests) instead of tensors of (dim, vocab_size).

```python
class CB_TFEmbeddings(TFEmbeddings):

    def call(self, input_ids=None, position_ids=None, inputs_embeds=None, 
             mode="embedding", training=False, ids_only=False):
        """Get token embeddings of inputs.
        Args:
            inputs: list of two int64 tensors with shape [batch_size, length]: (input_ids, position_ids)
            mode: string, a valid value is one of "embedding" and "linear".
        Returns:
            outputs: (1) If mode == "embedding", output embedding tensor, float32 with
                shape [batch_size, length, embedding_size]; (2) mode == "linear", output
                linear tensor, float32 with shape [batch_size, length, vocab_size].
        Raises:
            ValueError: if mode is not valid.

        Shared weights logic adapted from
            https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24
        """
        if mode == "embedding":
            return self._embedding(input_ids, position_ids, inputs_embeds, training=training)
        elif mode == "linear":
            return self._linear(input_ids, ids_only)
        else:
            raise ValueError("mode {} is not valid.".format(mode))

    def _linear(self, inputs, ids_only):
        """Computes logits by running inputs through a linear layer.
        Args:
            inputs: A float32 tensor with shape [batch_size, length, hidden_size]
        Returns:
            float32 tensor with shape [batch_size, length, vocab_size].
        """

        batch_size = shape_list(inputs)[0]
        length = shape_list(inputs)[1]
        
        x = tf.reshape(inputs, [-1, self.dim])
        
        if ids_only:#reduce the word_embeddings vector to the token ids of interest
            vocab_size = len(ids_only)
            word_embeddings = tf.convert_to_tensor(self.word_embeddings.numpy()[ids_only, :])
        else:
            word_embeddings = self.word_embeddings
            vocab_size = self.vocab_size
        
        logits = tf.matmul(x, word_embeddings, transpose_b=True)

        return tf.reshape(logits, [batch_size, length, vocab_size])
        
        
class CB_TFDistilBertMainLayer(TFDistilBertMainLayer):
    config_class = DistilBertConfig

    def __init__(self, config, **kwargs):
        super().__init__(config, **kwargs)
        self.num_hidden_layers = config.num_hidden_layers
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.return_dict = config.use_return_dict

        self.embeddings = CB_TFEmbeddings(config, name="embeddings")  # Embeddings
        self.transformer = TFTransformer(config, name="transformer")  # Encoder
        

class CB_TFDistilBertLMHead(TFDistilBertLMHead):

    def call(self, hidden_states, ids_only):
        hidden_states = self.input_embeddings(hidden_states, mode="linear", ids_only=ids_only)
        if ids_only: #if we focus only on some tokens, reduce the bias to those dimensions
            other_bias = tf.convert_to_tensor(self.bias.numpy()[ids_only])
            bias = other_bias
        else:
            bias = self.bias
        hidden_states = hidden_states + bias
        return hidden_states


class CB_TFDistilBertForMaskedLM(TFDistilBertForMaskedLM):
    
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.vocab_size = config.vocab_size

        self.distilbert = CB_TFDistilBertMainLayer(config, name="distilbert")
        self.vocab_transform = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="vocab_transform"
        )
        self.act = get_tf_activation("gelu")
        self.vocab_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name="vocab_layer_norm")
        self.vocab_projector = CB_TFDistilBertLMHead(config, self.distilbert.embeddings, name="vocab_projector")
        

    def call(
        self,
        inputs=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        labels=None,
        training=False,
        mask_only=False #If you just focus on the mask token id change the False by it (103 here)
        only_ids=False #Otherwise the tokens of interests as a list
    ):
        r"""
        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
        """
        
        return_dict = return_dict if return_dict is not None else self.distilbert.return_dict
        if isinstance(inputs, (tuple, list)):
            labels = inputs[7] if len(inputs) > 7 else labels
            if len(inputs) > 7:
                inputs = inputs[:7]
        elif isinstance(inputs, (dict, BatchEncoding)):
            labels = inputs.pop("labels", labels)

        distilbert_output = self.distilbert(
            inputs,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        hidden_states = distilbert_output[0]  # (bs, seq_length, dim)
        prediction_logits = self.vocab_transform(hidden_states)  # (bs, seq_length, dim)
        prediction_logits = self.act(prediction_logits)  # (bs, seq_length, dim)
        prediction_logits = self.vocab_layer_norm(prediction_logits)  # (bs, seq_length, dim)

        #transform from (bs, seq_length, dim) to (bs, 1, dim) where 1 is the MASK
        if mask_only:
            indexes = tf.where(inputs['input_ids'] == mask_only).numpy()
            prediction_logits = get_mask_rep(prediction_logits, indexes)
            
        prediction_logits = self.vocab_projector(prediction_logits, ids_only=ids_only)

        loss = None if labels is None else self.compute_loss(labels, prediction_logits)

        if not return_dict:
            output = (prediction_logits,) + distilbert_output[1:]
            return ((loss,) + output) if loss is not None else output

        return TFMaskedLMOutput(
            loss=loss,
            logits=prediction_logits,
            hidden_states=distilbert_output.hidden_states,
            attentions=distilbert_output.attentions
        )
    
def get_mask_rep(score, indexes):
    """
    Objective: get the score only for the indexes of interest.
    
    Inputs:
        - score, np.array or tf.tensor: the scores for all vocabulary
        - indexes, np.array: the array indicating [MASK] positionning
    Outputs:
        - final_scores, np.array: the final scores but only for the tokens of interest
    """
    n = score.shape[0]
    
    embedded_mask = dict(zip(indexes[:, 0], indexes[:, 1]))
    
    final_scores = []
    
    for index in range(score.shape[0]):
        final_scores.append(score[index, embedded_mask.get(index), :].numpy().reshape(1, -1))
        
    return np.array(final_scores)
```

Thus after that I uploaded a notebook on Lambda and looked at the speed of each version of code and especially how many sentences could I process with the same GPU (gpu.1x.rtx6000):

First line is the approximative maximum number of sentences I could process at once. Next lines are the time to process _n_ sentences by batches of 64. If there is a cross it because we got the MOO error before reaching the total number of sentences. 

|  Method |  Normal |     CS     |   RVCS   |
|:--------:|:--------:|:--------:|:--------:|
|max batch| 450 | 4500 | 4500  |
|900 stcs| 1.21 | 1.20 | 4.5 |
|30k stcs| x | 42.2 | 134.4 |
|40k tokens| x | 53.5 | 180.7 |
|50k tokens| x | x | 228.4 |
|60k tokens| x | x | 272.7 |

Well it is not much, but at least we can use the GPU for 40k sentences with a little change in the tensors shape and with decent performances. If we add the changes of reducing the tensors in the `_linear` method we reduce by 3 the performance but we can perform the process on more sentence with the same GPU. 

Code for the experiment:

```python
import time

patterns = create_patterns('{} [SEP] it talks about [MASK])

verbalizers = ['Society', 'Science', 'Health', 
               'Education', 'Computer', 'Sports', 
               'Business', 'Entertainment',
               'Relationship', 'Politics']

verb_ids = get_token_id(verbalizers, tokenizer)

model = TFDistilBertForMaskedLM.from_pretrained("distilbert-base-multilingual-cased")
#model = CB_TFDistilBertForMaskedLM.from_pretrained("distilbert-base-multilingual-cased")

step=64
st = time.time()
for i in range(0, n, step):
    texts = apply_pattern(sentences[:step], patterns[1])

    inputs = tokenizer.batch_encode_plus(texts, return_tensors='tf', 
                                 max_length=256, truncation=True, padding=True)

    score = model(inputs, return_dict=True)#, mask_only=103, ids_only=verb_ids)
print(time.time() - st)
```
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
