Merge pull request #23 from meedan/devops-471-add-model-dir

DEVOPS-471 add paths to allow for using new model directories
meedan · Aug 5, 2023 · ebb299b · ebb299b
2 parents 50cdfbe + 33452ae
commit ebb299b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lib/model/generic_transformer.py b/lib/model/generic_transformer.py
@@ -1,3 +1,4 @@
+import os
 from typing import Union, Dict, List
 from sentence_transformers import SentenceTransformer
 
@@ -11,7 +12,7 @@ def __init__(self, model_name: str):
         """
         self.model = None
         if model_name:
-            self.model = SentenceTransformer(model_name)
+            self.model = SentenceTransformer(model_name, cache_folder=os.getenv("MODEL_DIR", "./models"))
 
     def respond(self, docs: Union[List[schemas.Message], schemas.Message]) -> List[schemas.TextOutput]:
         """