Update optimum_encoder.py

optimum_encoder.py

@@ -115,9 +115,8 @@ class OptimumEncoder(BaseModel, Embeddings):
         docs: List[str],
         batch_size: int = 32,
         normalize_embeddings: bool = True,
-        pooling_strategy: str = "mean",
-        convert_to_numpy: bool = False
-    ) -> List[List[float]] | List[Dict[str, np.ndarray]]:
+        pooling_strategy: str = "mean"
+    ) -> List[List[float]]:
         all_embeddings = []
         for i in tqdm(range(0, len(docs), batch_size)):
             batch_docs = docs[i : i + batch_size]
@@ -129,28 +128,23 @@ class OptimumEncoder(BaseModel, Embeddings):
             with self._torch.no_grad():
                 model_output = self._model(**encoded_input)
 
-                if pooling_strategy == "mean":
-                    embeddings = self._mean_pooling(
-                        model_output, encoded_input["attention_mask"]
-                    )
-                elif pooling_strategy == "max":
-                    embeddings = self._max_pooling(
-                        model_output, encoded_input["attention_mask"]
-                    )
-                else:
-                    raise ValueError(
-                        "Invalid pooling_strategy. Please use 'mean' or 'max'."
-                    )
-                    
-                if normalize_embeddings:
-                    embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1)
-                    
-                if convert_to_numpy:
-                    embeddings = {'text': embeddings.cpu().detach().numpy()}
-                else:
-                    embeddings = embeddings.tolist()
+            if pooling_strategy == "mean":
+                embeddings = self._mean_pooling(
+                    model_output, encoded_input["attention_mask"]
+                )
+            elif pooling_strategy == "max":
+                embeddings = self._max_pooling(
+                    model_output, encoded_input["attention_mask"]
+                )
+            else:
+                raise ValueError(
+                    "Invalid pooling_strategy. Please use 'mean' or 'max'."
+                )
+
+            if normalize_embeddings:
+                embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1)
                 
-                all_embeddings.extend(embeddings)
+            all_embeddings.extend(embeddings.tolist())
 
         return all_embeddings
 
@@ -159,7 +153,7 @@ class OptimumEncoder(BaseModel, Embeddings):
         docs: str,
         normalize_embeddings: bool = True,
         pooling_strategy: str = "mean"
-    ) -> np.ndarray:
+    ) -> List[float]:
         encoded_input = self._tokenizer(
             docs, padding=True, truncation=True, return_tensors="pt"
         ).to(self.device)
@@ -167,23 +161,23 @@ class OptimumEncoder(BaseModel, Embeddings):
         with self._torch.no_grad():
             model_output = self._model(**encoded_input)
 
-            if pooling_strategy == "mean":
-                embeddings = self._mean_pooling(
-                    model_output, encoded_input["attention_mask"]
-                )
-            elif pooling_strategy == "max":
-                embeddings = self._max_pooling(
-                    model_output, encoded_input["attention_mask"]
-                )
-            else:
-                raise ValueError(
-                    "Invalid pooling_strategy. Please use 'mean' or 'max'."
-                )
-                
-            if normalize_embeddings:
-                embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1)
-        
-        return embeddings.cpu().detach().numpy()
+        if pooling_strategy == "mean":
+            embeddings = self._mean_pooling(
+                model_output, encoded_input["attention_mask"]
+            )
+        elif pooling_strategy == "max":
+            embeddings = self._max_pooling(
+                model_output, encoded_input["attention_mask"]
+            )
+        else:
+            raise ValueError(
+                "Invalid pooling_strategy. Please use 'mean' or 'max'."
+            )
+
+        if normalize_embeddings:
+            embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        print(embeddings)        
+        return embeddings.tolist()
 
     def _mean_pooling(self, model_output, attention_mask):
         token_embeddings = model_output[0]