Update main.py

main.py

@@ -6,28 +6,21 @@ from PIL import Image
 from io import BytesIO
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
+import torch
+import re
 
-app = FastAPI()
-
-# Use a pipeline as a high-level helper
-nlp_qa = pipeline("document-question-answering", model="impira/layoutlm-invoices")
-# Use a pipeline as a high-level helper
-nlp_ner = pipeline('question-answering', model='deepset/roberta-base-squad2', tokenizer='deepset/roberta-base-squad2')
-
+from transformers import DonutProcessor, VisionEncoderDecoderModel
 
-description = """
-## Image-based Document QA
-This API performs document question answering using a LayoutLM-based model.
+app = FastAPI()
 
-### Endpoints:
-- **POST /uploadfile/:** Upload an image file to extract text and answer provided questions.
-- **POST /pdfUpload/:** Provide a file to extract text and answer provided questions.
-"""
+processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
+model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
 
-app = FastAPI(docs_url="/", description=description)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-@app.post("/uploadfile/", description=description)
-async def perform_document_qa(
+@app.post("/donutQA/")
+async def donut_question_answering(
     file: UploadFile = File(...),
     questions: str = Form(...),
 ):
@@ -38,25 +31,57 @@ async def perform_document_qa(
         # Open the image using PIL
         image = Image.open(BytesIO(contents))
 
-        # Perform document question answering for each question using LayoutLM-based model
-        answers_dict = {}
-        for question in questions.split(','):
-            result = nlp_qa(
-                image,
-                question.strip()
-            )
-
-            # Access the 'answer' key from the first item in the result list
-            answer = result[0]['answer']
+        # Split the questions into a list
+        question_list = questions.split(',')
 
-            # Format the question as a string without extra characters
-            formatted_question = question.strip("[]")
+        # Process document with Donut model for each question
+        answers = process_document(image, question_list)
 
-            answers_dict[formatted_question] = answer
+        # Return a dictionary with questions and corresponding answers
+        result_dict = dict(zip(question_list, answers))
+        return result_dict
 
-        return answers_dict
     except Exception as e:
-        return JSONResponse(content=f"Error processing file: {str(e)}", status_code=500)
+        return {"error": f"Error processing file: {str(e)}"}
+
+def process_document(image, questions):
+    # prepare encoder inputs
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    
+    # prepare decoder inputs
+    task_prompt = "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
+    
+    # Initialize a list to store answers for each question
+    answers = []
+
+    # Process each question
+    for question in questions:
+        prompt = task_prompt.replace("{user_input}", question)
+        decoder_input_ids = processor.tokenizer(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
+
+        # Append the answer to the list
+        answers.append(processor.token2json(sequence))
+
+    return answers
 
 @app.post("/pdfQA/", description=description)
 async def pdf_question_answering(