path="."

handler.py

@@ -1,12 +1,14 @@
-from transformers import pipeline
+from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
 from typing import Dict, List, Any
 from tokenizers.decoders import WordPiece
 
 
 class EndpointHandler:
-    def __init__(self):
-        self.model = pipeline('ner', model='dicta-il/dictabert-ner', aggregation_strategy='simple')
-        self.model.tokenizer.backend_tokenizer.decoder = WordPiece()
+    def __init__(self, path="."):
+        model = AutoModelForTokenClassification.from_pretrained(path)
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        self.pipeline = pipeline('ner', model=model, tokenizer=tokenizer, aggregation_strategy='simple')
+        self.pipeline.tokenizer.backend_tokenizer.decoder = WordPiece()
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -16,4 +18,4 @@ class EndpointHandler:
         Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        return self.model(data['inputs'])
+        return self.pipeline(data['inputs'])