Upload bert_layers.py

bert_layers.py

@@ -610,7 +610,7 @@ class BertForMaskedLM(BertPreTrainedModel):
                 'If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for '
                 'bi-directional self-attention.')
 
-        self.bert = BertModel(config, add_pooling_layer=False)
+        self.bert = BertModel(config, add_pooling_layer=True)
         self.cls = BertOnlyMLMHead(config,
                                    self.bert.embeddings.word_embeddings.weight)
 
@@ -705,18 +705,13 @@ class BertForMaskedLM(BertPreTrainedModel):
                 return_dict=return_dict,
                 masked_tokens_mask=masked_tokens_mask,
             )
-            
+
         if torch.isnan(outputs[0]).any():
             print("NaNs in outputs.")
             raise ValueError()
         
-        #print("MLM Outputs")
-        #print(outputs[0].shape)
-
-        pooled_output = outputs[0]
-
-        last_hidden_state_formatted = outputs[0][:,0,:].view(-1, self.config.hidden_size)
-        return {"sentence_embedding": last_hidden_state_formatted}
+        pooled_output = outputs[1]
+        return {"sentence_embedding": pooled_output}
 
     def prepare_inputs_for_generation(self, input_ids: torch.Tensor,
                                       attention_mask: torch.Tensor,