Update modeling_cge2.py

modeling_cge2.py

@@ -493,7 +493,12 @@ class CgeForEmbedding(CGEModel):
             all_embeddings = all_embeddings[0]
         return all_embeddings
 
-
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
+        outputs = self.plm_model(input_ids, attention_mask, output_hidden_states=True)
+        hidden_states = outputs.hidden_states[self.keep_max_layer]
+        embeddings = self.get_sentence_embedding(self.embedding_method, hidden_states, 'disc', attention_mask)
+        return {"sentence_embeddings": embeddings}
+        
     def encode(self, sentences, batch_size=16, convert_to_numpy=False,
             convert_to_tensor=True, show_progress_bar=True, max_seq_length=None, **kwargs):
         if max_seq_length is None: