fix model dir

app/ner/services/ner.py

@@ -1,8 +1,8 @@
 from core.conf import settings
 
 class NER:
-    def __init__(self, model_name: str = settings.NER_MODEL_NAME):
-        self.model_name = model_name
+    def __init__(self, model_dir: str = settings.NER_MODEL_DIR):
+        self.model_dir = model_dir
         self.model = None
         self.tokenizer = None
         self.pipeline = None
@@ -10,8 +10,8 @@ class NER:
     def load_model(self):
         from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        self.model = AutoModelForTokenClassification.from_pretrained(self.model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
+        self.model = AutoModelForTokenClassification.from_pretrained(self.model_dir)
         self.pipeline = pipeline(settings.TASK_NAME, model=self.model, tokenizer=self.tokenizer)
         
     def predict(self, text: str, entity_tag: str = None):

core/conf.py

@@ -32,7 +32,7 @@ class Settings(BaseSettings):
     # DATABASE_PASSWORD: str
     
     # MODEl NER
-    NER_MODEL_NAME: str = 'ner'
+    NER_MODEL_DIR: str = '/app/ner/models/ner'
     TASK_NAME: str = 'ner'
 
     # FastAPI