Updated app.py to take pdfs

app.py

@@ -1,55 +1,15 @@
-import re
 import gradio as gr
+from helpers import majority_vote_dicts
+from model import process_document, process_pdf
 
-import torch
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-
-model_name = "harish3110/donut-quandri-all-data"
-
-processor = DonutProcessor.from_pretrained(model_name)
-model = VisionEncoderDecoderModel.from_pretrained(model_name)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-
-def process_document(image):
-    # prepare encoder inputs
-    pixel_values = processor(image, return_tensors="pt").pixel_values
-    
-    # prepare decoder inputs
-    task_prompt = "<s_cord-v2>"
-    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
-          
-    # generate answer
-    outputs = model.generate(
-        pixel_values.to(device),
-        decoder_input_ids=decoder_input_ids.to(device),
-        max_length=model.decoder.config.max_position_embeddings,
-        early_stopping=True,
-        pad_token_id=processor.tokenizer.pad_token_id,
-        eos_token_id=processor.tokenizer.eos_token_id,
-        use_cache=True,
-        num_beams=1,
-        bad_words_ids=[[processor.tokenizer.unk_token_id]],
-        return_dict_in_generate=True,
-    )
-    
-    # postprocess
-    sequence = processor.batch_decode(outputs.sequences)[0]
-    sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
-    
-    return processor.token2json(sequence)
-
+input_interface = gr.inputs.File(label="Upload a PDF file")
 description = "Gradio Demo for Donut trained on Quandri internal dataset."
-
-demo = gr.Interface(
-    fn=process_document,
-    inputs="image",
-    outputs="json",
-    title="Demo: Donut 🍩 for Document Parsing on Quandri Dataset",
-    description=description,
-    examples=[["aviva.jpeg"], ["wawanesa.jpeg"], ["sgi.jpeg"], ["travelers.jpeg"]],
-    cache_examples=False)
-
-demo.launch(auth=("quandri", "quandri1234"))
+gradio_interface = gr.Interface(fn=process_pdf, 
+                                inputs=input_interface, 
+                                outputs= 'json',  
+                                title="Demo: Donut 🍩 for Document Parsing on Quandri Dataset",
+                                description=description,
+                                examples=[["data/aviva.pdf"], ["data/wawanesa.pdf"], ["data/sgi.pdf"]],
+                                )
+
+gradio_interface.launch()