Update app.py

app.py

@@ -99,31 +99,42 @@ def m2(que, image):
     return processor3.batch_decode(generated_ids, skip_special_tokens=True)
 
 def m3(que, image):
-    processor3 = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
-    model3 = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
+    # processor3 = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
+    # model3 = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
 
-    model3.to(device)
+    # model3.to(device)
 
-    prompt = "<s_docvqa><s_question>{que}</s_question><s_answer>"
-    decoder_input_ids = processor3.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids
+    # prompt = "<s_docvqa><s_question>{que}</s_question><s_answer>"
+    # decoder_input_ids = processor3.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids
 
-    pixel_values = processor3(image, return_tensors="pt").pixel_values
+    # pixel_values = processor3(image, return_tensors="pt").pixel_values
 
-    outputs = model3.generate(
-        pixel_values.to(device),
-        decoder_input_ids=decoder_input_ids.to(device),
-        max_length=model3.decoder.config.max_position_embeddings,
-        pad_token_id=processor3.tokenizer.pad_token_id,
-        eos_token_id=processor3.tokenizer.eos_token_id,
-        use_cache=True,
-        bad_words_ids=[[processor3.tokenizer.unk_token_id]],
-        return_dict_in_generate=True,
-    )
+    # outputs = model3.generate(
+    #     pixel_values.to(device),
+    #     decoder_input_ids=decoder_input_ids.to(device),
+    #     max_length=model3.decoder.config.max_position_embeddings,
+    #     pad_token_id=processor3.tokenizer.pad_token_id,
+    #     eos_token_id=processor3.tokenizer.eos_token_id,
+    #     use_cache=True,
+    #     bad_words_ids=[[processor3.tokenizer.unk_token_id]],
+    #     return_dict_in_generate=True,
+    # )
+
+    # sequence = processor3.batch_decode(outputs.sequences)[0]
+    # sequence = sequence.replace(processor3.tokenizer.eos_token, "").replace(processor3.tokenizer.pad_token, "")
+    # sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
+    # return processor3.token2json(sequence)['answer']
+
+    
+    processor3 = AutoProcessor.from_pretrained("google/pix2struct-docvqa-large")
+    model3 = AutoModelForSeq2SeqLM.from_pretrained("google/pix2struct-docvqa-large")
+
+    inputs = processor3(images=image, text=que, return_tensors="pt")
+
+    predictions = model3.generate(**inputs)
+    return processor3.decode(predictions[0], skip_special_tokens=True)
 
-    sequence = processor3.batch_decode(outputs.sequences)[0]
-    sequence = sequence.replace(processor3.tokenizer.eos_token, "").replace(processor3.tokenizer.pad_token, "")
-    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
-    return processor3.token2json(sequence)['answer']
+    
 
 def m4(que, image):
     processor3 = Pix2StructProcessor.from_pretrained('google/matcha-plotqa-v2')