Update handler.py

handler.py

@@ -1,3 +1,4 @@
+from typing import Dict, List, Any
 from transformers import LayoutLMForTokenClassification, LayoutLMv2Processor, PegasusForConditionalGeneration, AlbertTokenizer
 import torch
 
@@ -6,11 +7,10 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class EndpointHandler():
     def __init__(self, path=""):
         # load model and processor from path
-        model_id = "popaqy/pegasus-base-qag-bg-finetuned-spelling6-bg"
-        self.model = PegasusForConditionalGeneration.from_pretrained(model_id).to(device)
-        self.tokenizer = AlbertTokenizer.from_pretrained(model_id)
+        self.model = PegasusForConditionalGeneration.from_pretrained(path).to(device)
+        self.tokenizer = AlbertTokenizer.from_pretrained(path)
 
-    def __call__(self, data):
+    def __call__(self, data: Dict[str, str]) -> Dict[str, List[str]]:
 
         # process input
         wrongSentence = data.pop("inputs", data)
@@ -22,4 +22,4 @@ class EndpointHandler():
         # run prediction
         output = self.model.generate(input['input_ids'])
 
-        return self.tokenizer.batch_decode(output, skip_special_tokens=True)
+        return {"predictions": self.tokenizer.batch_decode(output, skip_special_tokens=True)}