Repository: huggingface/transformers

Issue:
SqueezeBert does not appear to properly generate text
## Environment info
Google Colab 
Using CPU with High Ram

### Who can help
 @sgugger @forresti @LysandreJik 

## Information
Model I am using: Squeezebert-uncased, squeezebert-mnli, etc.

The problem arises when using:

Trying to generate the likely output of the input sequence and predicting masked tokens.

## To reproduce

```

from torch import nn
from transformers import AutoModelForMaskedLM, AutoTokenizer
model = AutoModelForMaskedLM.from_pretrained('squeezebert/squeezebert-mnli')
tokenizer = AutoTokenizer.from_pretrained('squeezebert/squeezebert-mnli')
#model.tie_weights()
input_txt = ["[MASK] was an American [MASK]  and lawyer who served as the 16th president  of the United States from 1861 to 1865. [MASK] led the nation through the American Civil War, the country's greatest [MASK], [MASK], and [MASK] crisis. ", \
             "George [MASK], who served as the first  president of the United States from [MASK] to 1797, was an American political leader, [MASK] [MASK], statesman, and Founding Father. Previously, he led Patriot forces to [MASK] in the nation's War for Independence. ", \
             "[MASK], the first African-American [MASK] of the [MASK] [MASK], is an American politician and attorney who served as the 44th [MASK] of the United States from [MASK] to 2017.  [MASK] was a member of the [MASK] [MASK]. "]
#input_txt = 
input_txt= [i.replace("[MASK]", tokenizer.mask_token) for i in input_txt] #
inputs = tokenizer(input_txt, return_tensors='pt', add_special_tokens=True, padding=True)
inputs['output_attentions'] = True
inputs['output_hidden_states'] = True
inputs['return_dict'] = True
outputs = model(**inputs)
if True:
  predictions = outputs.logits
  for pred in predictions:
    print ("**")
    sorted_preds, sorted_idx = pred.sort(dim=-1, descending=True)
    for k in range(2):
        predicted_index = [sorted_idx[i, k].item() for i in range(0,len(predictions[0]))]
        predicted_token = ' '.join([tokenizer.convert_ids_to_tokens([predicted_index[x]])[0] for x in range(1,len(predictions[0]))]).replace('Ġ', ' ').replace('  ', ' ').replace('##', '')
        print(predicted_token)

```
## Expected behavior

I expected at least the input to be echoed out, with the slots filling with Lincoln, Washington and Obama. This works for bert, distlbert, roberta, etc. 

## Actual output

Some weights of the model checkpoint at squeezebert/squeezebert-mnli were not used when initializing SqueezeBertForMaskedLM: ['classifier.weight', 'classifier.bias']
- This IS expected if you are initializing SqueezeBertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).
- This IS NOT expected if you are initializing SqueezeBertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of SqueezeBertForMaskedLM were not initialized from the model checkpoint at squeezebert/squeezebert-mnli and are newly initialized: ['lm_head.weight', 'lm_head.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
odict_keys(['logits', 'hidden_states', 'attentions'])
**
tone lani rce soto olar rce ux wer lani bal bal vus novice rce rce rce lani owe frey owe gent naire tres che lani lani nae ui territories accusing oaks accusing ois lor resulting resulting rce lor rce rendering rce rce tres assist ois accusing rendering warns accusing gent
culture bowls hectares awan rce bal ade wd an rce mole hoe yde lani lani lani rce tres resulted bal resulted resulting tone consequently bowls fellow wo ois crafts oaks withdrew nations wu resulting fellow rce resulting verses motivated lori motivated motivated gent vus naire dealt warns gent warns tres
**
culture sas hari lani rce gaa lani novice rce rce rce rce tres nae jan thal rce rce rce awan olar v8 rce olar example rce select rce rce hore rden resulting lori resulting drive led bon peoples jal gau nae hoe lies lies lies lies lins lies resulting tone
continuum tone repeat gaa lani wo rce coven lani lani lani lani gle aw aw awan sco lani yde rce yde olar ux rce rce trait xie xie cao particular elder lani lani naturally blend lie aman commando folding rendering helps ois lete wi lins lins hoe independence sons tones
**
tone acts attribute trait pour pour trait % sities ub azi % acts lani rce awan act cao yde wd hoe hoe hoe hoe % vos vos rce hort hoe sept jan vers naire hum candle therefore lists chen hoe lie side mut hen mor lungs zoo lie side side
hum fever acts pour shropshire cz % sities isson penalties lie sities act acts bble pour yde ave shropshire yde lto ango ango pour lden rce hoe gil hoe tres aw nae dha therefore bisexual therefore lb mates rden too zoo forum naire dealt lag mole mess pore forum ior



Discussion:
Hello! First of all, you're using the `squeezebert-mnli` checkpoint, which is a checkpoint that was fine-tuned on the MNLI dataset. It cannot be used to do masked language modeling.

I believe you should be using the `squeezebert-uncased` checkpoint instead.

However, even when using that checkpoint with the MLM pipeline I cannot obtain sensible results. Maybe @forresti can chime in and let us know if something's up!


Thanks @LysandreJik . I used both squeezebert-mnli and squeezebert-uncased (not shown). Same type of results. Thanks for checking. @forresti  any thoughts? Is there something wrong with the squeezbert tokenizer? 
@ontocord Sorry for the slow reply. I will dig into this on Thursday this week.

Resolving patch:
diff --git a/src/transformers/modeling_squeezebert.py b/src/transformers/modeling_squeezebert.py
--- a/src/transformers/modeling_squeezebert.py
+++ b/src/transformers/modeling_squeezebert.py
@@ -373,6 +373,53 @@ def forward(self, hidden_states):
         return pooled_output
 
 
+class SqueezeBertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class SqueezeBertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = SqueezeBertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+class SqueezeBertOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = SqueezeBertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
 class SqueezeBertPreTrainedModel(PreTrainedModel):
     """
     An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
@@ -594,16 +641,19 @@ def forward(
 
 @add_start_docstrings("""SqueezeBERT Model with a `language modeling` head on top. """, SQUEEZEBERT_START_DOCSTRING)
 class SqueezeBertForMaskedLM(SqueezeBertPreTrainedModel):
+
+    authorized_missing_keys = [r"predictions.decoder.bias"]
+
     def __init__(self, config):
         super().__init__(config)
 
         self.transformer = SqueezeBertModel(config)
-        self.lm_head = nn.Linear(config.embedding_size, config.vocab_size)
+        self.cls = SqueezeBertOnlyMLMHead(config)
 
         self.init_weights()
 
     def get_output_embeddings(self):
-        return self.lm_head
+        return self.cls.predictions.decoder
 
     @add_start_docstrings_to_model_forward(SQUEEZEBERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
     @add_code_sample_docstrings(
@@ -646,7 +696,7 @@ def forward(
         )
 
         sequence_output = outputs[0]
-        prediction_scores = self.lm_head(sequence_output)
+        prediction_scores = self.cls(sequence_output)
 
         masked_lm_loss = None
         if labels is not None:


Test patch: