Update handler.py

handler.py

@@ -10,8 +10,8 @@ from transformers import DonutProcessor, VisionEncoderDecoderModel
 class EndpointHandler:
     def __init__(self, path=""):
         # load model and processor from path
-        self.processor = DonutProcessor.from_pretrained("debu-das/donut_receipt_v2.29")
-        self.model = VisionEncoderDecoderModel.from_pretrained("debu-das/donut_receipt_v2.29")
+        self.processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+        self.model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
@@ -22,7 +22,7 @@ class EndpointHandler:
         return self.process_document(image)
 
 
-    def process_document(self, image):
+    def process_document(self, image:Image) -> dict[str, Any]:
         # prepare encoder inputs
         pixel_values = self.processor(image, return_tensors="pt").pixel_values
 
@@ -49,4 +49,4 @@ class EndpointHandler:
         sequence = sequence.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
         sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
 
-        return self.processor.token2json(sequence)
+        return self.processor.token2json(sequence)