Update optimum_encoder.py

optimum_encoder.py

@@ -19,6 +19,7 @@ class OptimumEncoder(BaseEncoder):
     _tokenizer: Any = PrivateAttr()
     _model: Any = PrivateAttr()
     _torch: Any = PrivateAttr()
+    _dim: int = 1024
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -44,7 +45,7 @@ class OptimumEncoder(BaseEncoder):
                 "`pip install semantic-router[local]`"
             )
         try:
-            from transformers import AutoTokenizer
+            from transformers import AutoTokenizer, AutoConfig
         except ImportError:
             raise ImportError(
                 "Please install transformers to use OptimumEncoder. "
@@ -59,6 +60,12 @@ class OptimumEncoder(BaseEncoder):
             **self.tokenizer_kwargs,
         )
 
+        config = AutoConfig.from_pretrained(
+            self.name
+        )
+
+        self._dim = config['hidden_size']
+        
         provider_options = {
             "trt_engine_cache_enable": True,
             "trt_engine_cache_path": os.getenv('HF_HOME'),
@@ -69,6 +76,7 @@ class OptimumEncoder(BaseEncoder):
         session_options.log_severity_level = 0
 
         ort_model = ORTModelForFeatureExtraction.from_pretrained(
+            config=config,
             model_id=self.name, 
             file_name='model_fp16.onnx',
             subfolder='onnx',
@@ -141,9 +149,7 @@ class OptimumEncoder(BaseEncoder):
                 else:
                     embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1).detach().cpu().tolist()
 
-            original_dimensions = embeddings.size()
-
-            if original_dimensions > matryoshka_dim:
+            if self._dim > matryoshka_dim:
                 embeddings = embeddings[:, :matryoshka_dim]
                 
             all_embeddings.extend(embeddings)