fix load model

app/ner/models/base_model.py

@@ -1,16 +0,0 @@
-from datetime import datetime
-from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForTokenClassification
-from optimum.pipelines import pipeline
-start = datetime.now()
-
-
-
-tokenizer = AutoTokenizer.from_pretrained("model_ner")
-model = ORTModelForTokenClassification.from_pretrained("model_ner", provider="CPUExecutionProvider")
-pipe = pipeline("token-classification", model=model, tokenizer=tokenizer, accelerator="ort", device=-1)
-
-result = pipe("Tôi tên là Trần Văn Đại, địa chỉ 12 Phan đình phùng")
-end = datetime.now()
-print(result)
-print("Time taken: ", (end - start).total_seconds())

app/ner/services/ner.py

@@ -8,11 +8,13 @@ class NER:
         self.pipeline = None
 
     def load_model(self):
-        from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
+        from transformers import AutoTokenizer
+        from optimum.onnxruntime import ORTModelForTokenClassification
+        from optimum.pipelines import pipeline
 
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir, local_files_only=True)
-        self.model = AutoModelForTokenClassification.from_pretrained(self.model_dir, local_files_only=True)
-        self.pipeline = pipeline(settings.TASK_NAME, model=self.model, tokenizer=self.tokenizer)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
+        self.model = ORTModelForTokenClassification.from_pretrained(self.model_dir)
+        self.pipeline = pipeline(task=settings.TASK_NAME, model=self.model, tokenizer=self.tokenizer, accelerator=settings.ACCELERATOR, device=settings.DEVICE)
         
     def predict(self, text: str, entity_tag: str = None):
         if self.pipeline is None:

core/conf.py

@@ -32,8 +32,10 @@ class Settings(BaseSettings):
     # DATABASE_PASSWORD: str
     
     # MODEl NER
-    NER_MODEL_DIR: str = "/app/ner/models/ner"
-    TASK_NAME: str = 'ner'
+    NER_MODEL_DIR: str = 'app/ner/models/ner'
+    TASK_NAME: str = 'token-classification'
+    ACCELERATOR: str = 'ort'
+    DEVICE: str = 'cpu'
 
     # FastAPI
     FASTAPI_API_V1_PATH: str = '/api/v1'