Update optimum_encoder.py

optimum_encoder.py

@@ -19,7 +19,6 @@ class OptimumEncoder(BaseEncoder):
     _tokenizer: Any = PrivateAttr()
     _model: Any = PrivateAttr()
     _torch: Any = PrivateAttr()
-    _dim: int = 1024
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -45,7 +44,7 @@ class OptimumEncoder(BaseEncoder):
                 "`pip install semantic-router[local]`"
             )
         try:
-            from transformers import AutoTokenizer, AutoConfig
+            from transformers import AutoTokenizer
         except ImportError:
             raise ImportError(
                 "Please install transformers to use OptimumEncoder. "
@@ -59,12 +58,6 @@ class OptimumEncoder(BaseEncoder):
             self.name,
             **self.tokenizer_kwargs,
         )
-
-        config = AutoConfig.from_pretrained(
-            self.name
-        )
-
-        self._dim = config['hidden_size']
         
         provider_options = {
             "trt_engine_cache_enable": True,
@@ -116,7 +109,6 @@ class OptimumEncoder(BaseEncoder):
         batch_size: int = 32,
         normalize_embeddings: bool = True,
         pooling_strategy: str = "mean",
-        matryoshka_dim: int = 1024,
         convert_to_numpy: bool = False
     ) -> List[List[float]] | List[np.ndarray]:
         all_embeddings = []
@@ -142,15 +134,12 @@ class OptimumEncoder(BaseEncoder):
                 raise ValueError(
                     "Invalid pooling_strategy. Please use 'mean' or 'max'."
                 )
-            print(f'Embeddings {embeddings}')
+
             if normalize_embeddings:
                 if convert_to_numpy:
                     embeddings = normalize(embeddings[:, 0]).astype(np.float32)
                 else:
                     embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1).detach().cpu().tolist()
-
-            if self._dim > matryoshka_dim:
-                embeddings = embeddings[:, :matryoshka_dim]
                 
             all_embeddings.extend(embeddings)