Update app.py

app.py

@@ -1,8 +1,104 @@
 import gradio as gr
+import torch, os, json, requests
+from PIL import Image
+from transformers import DonutProcessor, VisionEncoderDecoderModel, VisionEncoderDecoderConfig
 
+def load_image_from_URL(url):
+    res = requests.get(url)
+
+    if res.status_code == 200:
+        img = Image.open(requests.get(url, stream = True).raw)
+
+        if img.mode == "RGBA":
+            img = img.convert("RGB")
+
+        return img
+
+    return None
+
+class OCRVQAModel(torch.nn.Module):
+    def add_tokens(self, list_of_tokens):
+        self.added_tokens.update(list_of_tokens)
+        newly_added_num = self.processor.tokenizer.add_tokens(list_of_tokens)
+        
+        if newly_added_num > 0:
+            self.donut.decoder.resize_token_embeddings(len(self.processor.tokenizer))
+    
+    def __init__(self, config):
+        super().__init__()
+        
+        self.model_name_or_path = config['donut']
+        self.processor_name_or_path = config['processor']
+        self.config_name_or_path = config['config']
+                
+        self.donut_config = VisionEncoderDecoderConfig.from_pretrained(self.config_name_or_path)
+        self.donut_config.encoder.image_size = [800, 600]
+        self.donut_config.decoder.max_length = 64
+        
+        self.processor = DonutProcessor.from_pretrained(self.processor_name_or_path)
+        self.donut = VisionEncoderDecoderModel.from_pretrained(self.model_name_or_path, config = self.donut_config)
+        
+        self.added_tokens = set([])
+        self.setup()
+        
+    def setup(self):
+        self.add_tokens(["<yes/>", "<no/>"])
+        self.processor.feature_extractor.size = self.donut_config.encoder.image_size[::-1]
+        self.processor.feature_extractor.do_align_long_axis = False
+            
+    def inference(self, image_src, prompt, device):
+        if os.path.exists(image_src): 
+            image = Image.open(image_src)
+        else: 
+            image = load_image_from_URL(image_src)
+
+        if not Image:
+            return {
+                'question': prompt,
+                'answer': 'Some error occurred during inference time.'
+            }
+
+        self.donut.eval()
+        with torch.no_grad():
+            image_ids = self.processor(image, return_tensors="pt").pixel_values.to(device)
+
+            question = f'<s_docvqa><s_question>{prompt}</s_question><s_answer>'
+
+            embedded_question = self.processor.tokenizer(
+                question,
+                add_special_tokens = False, 
+                return_tensors = "pt"
+            )["input_ids"].to(device)
+
+            outputs = self.donut.generate(
+                image_ids,
+                decoder_input_ids=embedded_question,
+                max_length = self.donut.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
+            )
+            
+            return self.processor.token2json(self.processor.batch_decode(outputs.sequences)[0])
+
+model = OCRVQAModel(
+    'ndtran/donut_ocr-vqa-200k',
+    'ndtran/donut_ocr-vqa-200k'
+)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
 
 def get_answer(image, question) -> str:
-    return "I don't know"
+    global model, device
+    result = model.inference(image, question, device)
+    return result.get('answer', 'I don\'t know :<')
 
 with gr.Blocks() as demo:
     with gr.Row():