Update modeling_sim.py

modeling_sim.py

@@ -11,8 +11,7 @@ class SimModel(MobileBertPreTrainedModel):
         self.post_init()
 
     def forward(self, input_ids, attention_mask, token_type_ids, return_dict):
-        print(input_ids, attention_mask, token_type_ids)
-        print(self.word_embeddings(input_ids).size())
-        import pdb
-        pdb.set_trace()
-        return (self.word_embeddings(input_ids), None, None)
+        embeddings = self.word_embeddings(input_ids)
+        masked_embeddings = embeddings * attention_mask[:, :, None]
+        mean_pooled_embeddings = masked_embeddings.sum(dim=1) / attention_mask[:, :, None].sum(dim=1)
+        return (embeddings, mean_pooled_embeddings, embeddings)