Upload handler.py with huggingface_hub

handler.py

@@ -1,14 +1,12 @@
-from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
 from typing import Dict, List, Any
-from tokenizers.decoders import WordPiece
+from transformers import AutoModel, AutoTokenizer
 
 
 class EndpointHandler:
     def __init__(self, path="."):
-        model = AutoModelForTokenClassification.from_pretrained(path)
-        tokenizer = AutoTokenizer.from_pretrained(path)
-        self.pipeline = pipeline('token-classification', model=model, tokenizer=tokenizer, aggregation_strategy='simple')
-        self.pipeline.tokenizer.backend_tokenizer.decoder = WordPiece()
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModel.from_pretrained(path, trust_remote_code=True)
+        self.model.eval()
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -18,4 +16,6 @@ class EndpointHandler:
         Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        return self.pipeline(data['inputs'])
+        # return self.pipeline(data['inputs'])
+        return self.model.predict([data['inputs']], self.tokenizer, output_style='json')
+