Upload model

modeling_gzipembed.py

@@ -39,6 +39,19 @@ class GZIPEmbeddingModel(PreTrainedModel):
             x = x.to(self.reduction_head.dtype).to(self.reduction_head.device)
             return self.reduction_head(x)
         return x if not return_tensor else torch.tensor(x)
+    
+    def encode(self, sentences, batch_size=32, **kwargs):
+        """
+        Returns a list of embeddings for the given sentences.
+        Args:
+            sentences (`List[str]`): List of sentences to encode
+            batch_size (`int`): Batch size for the encoding
+
+        Returns:
+            `List[np.ndarray]` or `List[tensor]`: List of embeddings for the given sentences
+        """
+        x = self.forward(sentences, num_procs=batch_size, return_tensor=False)
+        return [torch.tensor(i) for i in x]
 
     def normalize(self, x):
         x = ''.join([char for char in x.lower() if char in "abcdefghijklmnopqrstuvwxyz "])