Update optimum_encoder.py

optimum_encoder.py

@@ -1,6 +1,7 @@
 import requests
 import time
 import os
+import numpy as np
 from typing import Any, List, Optional, Dict
 from pydantic.v1 import PrivateAttr
 
@@ -101,7 +102,8 @@ class OptimumEncoder(BaseEncoder):
         batch_size: int = 32,
         normalize_embeddings: bool = True,
         pooling_strategy: str = "mean",
-        matryoshka_dim: int = 512
+        matryoshka_dim: int = 512,
+        convert_to_numpy: bool = True
     ) -> List[List[float]]:
         all_embeddings = []
         for i in range(0, len(docs), batch_size):
@@ -137,6 +139,9 @@ class OptimumEncoder(BaseEncoder):
 
             embeddings = embeddings.tolist()
             all_embeddings.extend(embeddings)
+
+        if convert_to_numpy:
+            return np.asarray([emb.numpy() for emb in all_embeddings])
         return all_embeddings
 
     def _mean_pooling(self, model_output, attention_mask):