Update main.py

main.py

@@ -1,19 +1,29 @@
 import fitz
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import JSONResponse
+from transformers import pipeline
 from PIL import Image
 from io import BytesIO
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
-from transformers import pipeline, DistilBertTokenizer, DistilBertForQuestionAnswering
 
 app = FastAPI()
 
+# Set up CORS middleware
+origins = ["*"]  # or specify your list of allowed origins
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 nlp_qa = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
 
 description = """
 ## Image-based Document QA
-This API performs document question answering using a DistilBERT-based model.
+This API performs document question answering using a LayoutLMv2-based model.
 
 ### Endpoints:
 - **POST /uploadfile/:** Upload an image file to extract text and answer provided questions.
@@ -34,16 +44,16 @@ async def perform_document_qa(
         # Open the image using PIL
         image = Image.open(BytesIO(contents))
 
-        # Perform document question answering for each question using DistilBERT-based model
+        # Perform document question answering for each question using LayoutLMv2-based model
         answers_dict = {}
         for question in questions.split(','):
             result = nlp_qa(
-                question.strip(),
-                image
+                image,
+                question.strip()
             )
 
-            # Access the 'answer' key from the result
-            answer = result['answer']
+            # Access the 'answer' key from the first item in the result list
+            answer = result[0]['answer']
 
             # Format the question as a string without extra characters
             formatted_question = question.strip("[]")