Update handler.py

handler.py

@@ -1,5 +1,6 @@
 import torch
 import io
+import re
 
 from typing import Any, Dict
 from PIL import Image
@@ -18,20 +19,34 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
         image = inputs["image"]
         image = Image.open(io.BytesIO(eval(image)))
-        text = inputs["text"]
-    
-        # preprocess
-        encoding = self.processor(image, return_tensors="pt")
-        outputs = self.model(**encoding)
-        # postprocess the prediction
-        logits = outputs.logits
-        best_idx = logits.argmax(-1).item()
-        best_answer = self.model.config.id2label[best_idx]
-        probabilities = torch.softmax(logits, dim=-1)[0]
-        id2label = self.model.config.id2label
-        answers = []
-        for idx, prob in enumerate(probabilities):
-            answer = id2label[idx]
-            answer_score = float(prob)
-            answers.append({"answer": answer, "answer_score": answer_score})
-        return {"best_answer": best_answer, "answers": answers}
+        return self.process_document(image)
+
+
+    def process_document(self, image):
+        # prepare encoder inputs
+        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+
+        # prepare decoder inputs
+        task_prompt = "<s_cord-v2>"
+        decoder_input_ids = self.processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
+                
+        # generate answer
+        outputs = self.model.generate(
+            pixel_values.to(self.device),
+            decoder_input_ids=decoder_input_ids.to(self.device),
+            max_length=self.model.decoder.config.max_position_embeddings,
+            early_stopping=True,
+            pad_token_id=self.processor.tokenizer.pad_token_id,
+            eos_token_id=self.processor.tokenizer.eos_token_id,
+            use_cache=True,
+            num_beams=1,
+            bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+        )
+
+        # postprocess
+        sequence = self.processor.batch_decode(outputs.sequences)[0]
+        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)