Update app.py

app.py

@@ -1,35 +1,28 @@
 
 import re
 import gradio as gr
+
 import torch
 from transformers import DonutProcessor, VisionEncoderDecoderModel
-import transformers
-from PIL import Image
-import random
-import numpy as np
-
-# hidde logs
-transformers.logging.disable_default_handler()
-
-
-# Load our model from Hugging Face
-processor = DonutProcessor.from_pretrained("Iqra56/Donut_Updated")
-model = VisionEncoderDecoderModel.from_pretrained("Iqra56/Donut_Updated")
 
-# Move model to GPU
+#processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+#model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+#processor = DonutProcessor.from_pretrained("Iqra56/ENGLISHDONUT")
+#model = VisionEncoderDecoderModel.from_pretrained("Iqra56/ENGLISHDONUT")
+processor = DonutProcessor.from_pretrained("Iqra56/DONUTWOKEYS")
+model = VisionEncoderDecoderModel.from_pretrained("Iqra56/DONUTWOKEYS")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-# Load random document image from the test set
-test_sample = processed_dataset["test"][random.randint(1,7)]
-
-def run_prediction(sample, model=model, processor=processor):
-    # prepare inputs
-    pixel_values = torch.tensor(test_sample["pixel_values"]).unsqueeze(0)
+def process_document(image):
+    # prepare encoder inputs
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    
+    # prepare decoder inputs
     task_prompt = "<s>"
     decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
-
-    # run inference
+          
+    # generate answer
     outputs = model.generate(
         pixel_values.to(device),
         decoder_input_ids=decoder_input_ids.to(device),
@@ -50,18 +43,18 @@ def run_prediction(sample, model=model, processor=processor):
     
     return processor.token2json(sequence)
 
-description = "Gradio Demo for Donut, an instance of `VisionEncoderDecoderModel` fine-tuned on DocVQA (document visual question answering). To use it, simply upload your image and type a question and click 'submit', or click one of the examples to load them. Read more at the links below."
+description = "Gradio Demo for Donut, an instance of `VisionEncoderDecoderModel` fine-tuned on CORD (document parsing). To use it, simply upload your image and click 'submit', or click one of the examples to load them. Read more at the links below."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.15664' target='_blank'>Donut: OCR-free Document Understanding Transformer</a> | <a href='https://github.com/clovaai/donut' target='_blank'>Github Repo</a></p>"
 
 demo = gr.Interface(
     fn=process_document,
-    inputs=["image", "text"],
+    inputs="image",
     outputs="json",
-    title="Demo: Donut 🍩 for DocVQA",
+    title="Demo: Donut 🍩 for Document Parsing",
     description=description,
     article=article,
     enable_queue=True,
-    examples=[["example_1.png", "When is the coffee break?"], ["example_2.jpeg", "What's the population of Stoddard?"]],
+    examples=[[""], [""], [""]],
     cache_examples=False)
 
 demo.launch()