Update optimum_encoder.py

optimum_encoder.py

@@ -11,6 +11,7 @@ from langchain_core.pydantic_v1 import BaseModel, Extra, root_validator
 class OptimumEncoder(BaseModel, Embeddings):
     name: str = "mixedbread-ai/mxbai-embed-large-v1"
     device: Optional[str] = None
+    cache_dir: str = None
     _tokenizer: Any
     _model: Any
     _torch: Any
@@ -24,6 +25,7 @@ class OptimumEncoder(BaseModel, Embeddings):
     def validate_environment(cls, values: Dict) -> Dict:
         name = values.get('name')
         device = values.get('device')
+        cache_dir = values.get('cache_dir')
 
         try:
             import onnxruntime as ort
@@ -66,6 +68,7 @@ class OptimumEncoder(BaseModel, Embeddings):
             file_name='model_fp16.onnx',
             subfolder='onnx',
             provider='CUDAExecutionProvider',
+            cache_dir=cache_dir,
             use_io_binding=True,
             #provider_options=provider_options,
             session_options=session_options