Update app.py

app.py

@@ -103,104 +103,59 @@ async def get_docs(request: Request):
 from fastapi import FastAPI
 from fastapi.responses import RedirectResponse
 import gradio as gr
-from transformers import pipeline
-import pdfplumber
-import docx
-from pptx import Presentation
+from transformers import VilBertForQuestionAnswering, ViltProcessor
 from PIL import Image
-import pytesseract
-import easyocr
-import os
-from io import BytesIO
+import torch
 
-# Initialize FastAPI app
+# Initialize FastAPI
 app = FastAPI()
 
-# Load models
-qa_pipeline = pipeline("question-answering", model="bert-large-uncased-whole-word-masking-finetuned-squad")
-image_qa_pipeline = pipeline("vqa", model="Salesforce/blip-vqa-base")
-reader = easyocr.Reader(['en'])
+# Load VilBERT model and processor
+model = VilBertForQuestionAnswering.from_pretrained("facebook/vilbert-vqa")
+processor = ViltProcessor.from_pretrained("facebook/vilbert-vqa")
 
-# File parsing functions with error handling
-def extract_text_from_pdf(file):
-    try:
-        with pdfplumber.open(file) as pdf:
-            return "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
-    except Exception as e:
-        return f"Error reading PDF: {str(e)}"
-
-def extract_text_from_docx(file):
-    try:
-        doc = docx.Document(file)
-        return "\n".join(para.text for para in doc.paragraphs)
-    except Exception as e:
-        return f"Error reading DOCX: {str(e)}"
-
-def extract_text_from_pptx(file):
-    try:
-        prs = Presentation(file)
-        return "\n".join(shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text"))
-    except Exception as e:
-        return f"Error reading PPTX: {str(e)}"
-
-def extract_text_from_image(file):
-    try:
-        # Use easyocr for better image text extraction
-        return easyocr.Reader(['en']).readtext(file)
-    except Exception as e:
-        return f"Error reading image: {str(e)}"
-
-# Main QA logic for documents and images
-def answer_question(question, file):
-    file_ext = os.path.splitext(file.name)[-1].lower()
-
-    if file_ext == ".pdf":
-        context = extract_text_from_pdf(file)
-    elif file_ext == ".docx":
-        context = extract_text_from_docx(file)
-    elif file_ext == ".pptx":
-        context = extract_text_from_pptx(file)
-    elif file_ext in [".png", ".jpg", ".jpeg", ".bmp"]:
-        context = extract_text_from_image(file)
-    else:
-        return "❌ Unsupported file format."
+# Function to handle image question answering
+def answer_question_from_image(image, question):
+    if image is None or question.strip() == "":
+        return "Please upload an image and enter a question."
+    
+    # Process input
+    inputs = processor(images=image, text=question, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted_idx = outputs.logits.argmax(-1).item()
+    
+    # For VilBERT VQA, class index maps to predefined answers (like "yes", "no", etc.)
+    # You'd need the VQA label mapping to decode this properly
+    # For now, just return the index
+    return f"Predicted answer ID: {predicted_idx}"
 
-    if not context.strip():
-        return "⚠️ No readable text found in the document."
+# Create Image QA interface
+img_interface = gr.Interface(
+    fn=answer_question_from_image,
+    inputs=[gr.Image(label="Upload Image"), gr.Textbox(label="Ask a Question")],
+    outputs="text",
+    title="AI Image Question Answering"
+)
 
-    result = qa_pipeline(question=question, context=context)
-    return result["answer"]
+# Dummy doc QA interface (replace with your own implementation)
+def dummy_doc_qa(doc, question):
+    return "This is a placeholder for Document QA."
 
-# Create Gradio interfaces for both document and image QA
 doc_interface = gr.Interface(
-    fn=answer_question,
-    inputs=[
-        gr.Textbox(label="Ask a question"),
-        gr.File(label="Upload a document (PDF, DOCX, PPTX)")
-    ],
-    outputs=gr.Textbox(label="Answer"),
-    title="Document Question Answering",
-    description="Upload a document and ask a question. Get answers from the document content.",
-)
-
-img_interface = gr.Interface(
-    fn=answer_question,
-    inputs=[
-        gr.Textbox(label="Ask a question"),
-        gr.File(label="Upload an image (PNG, JPG, etc.)")
-    ],
-    outputs=gr.Textbox(label="Answer"),
-    title="Image Question Answering",
-    description="Upload an image and ask a question. Get answers from the text extracted from the image.",
+    fn=dummy_doc_qa,
+    inputs=[gr.File(label="Upload Document"), gr.Textbox(label="Ask a Question")],
+    outputs="text",
+    title="Document Question Answering"
 )
 
-# Create a Tabbed Interface to switch between document and image QA
+# Combine into a tabbed interface
 demo = gr.TabbedInterface([doc_interface, img_interface], ["Document QA", "Image QA"])
 
-# Mount Gradio app in FastAPI
+# Mount Gradio inside FastAPI at root "/"
 app = gr.mount_gradio_app(app, demo, path="/")
 
-# Redirect to Gradio interface
+# Redirect root URL to Gradio UI
 @app.get("/")
 def home():
     return RedirectResponse(url="/")