Update app.py

app.py

@@ -101,13 +101,13 @@ async def get_docs(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 """
 from fastapi import FastAPI
+from fastapi.responses import RedirectResponse
 import gradio as gr
 from transformers import pipeline
 import pdfplumber, docx
 from pptx import Presentation
 from PIL import Image
 import pytesseract
-import fitz
 import easyocr
 import os
 
@@ -119,7 +119,7 @@ qa_pipeline = pipeline("question-answering", model="microsoft/phi-2", tokenizer=
 image_qa_pipeline = pipeline("vqa", model="Salesforce/blip-vqa-base")
 reader = easyocr.Reader(['en'])
 
-# File parsing
+# File parsing functions
 def extract_text_from_pdf(file):
     with pdfplumber.open(file) as pdf:
         return "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
@@ -136,10 +136,10 @@ def extract_text_from_image(file):
     image = Image.open(file)
     return pytesseract.image_to_string(image)
 
-# Main QA logic
+# 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":
@@ -157,21 +157,36 @@ def answer_question(question, file):
     result = qa_pipeline(question=question, context=context)
     return result["answer"]
 
-# Gradio interface
-gr_interface = gr.Interface(
+# 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 a document or image")
+        gr.File(label="Upload an image (PNG, JPG, etc.)")
     ],
     outputs=gr.Textbox(label="Answer"),
-    title="AI Question Answering (Text & Image)",
-    description="Upload a file (PDF, DOCX, PPTX, Image) and ask a question. Get instant answers from document content.",
+    title="Image Question Answering",
+    description="Upload an image and ask a question. Get answers from the text extracted from the image.",
 )
 
+# Create a Tabbed Interface to switch between document and image QA
+demo = gr.TabbedInterface([doc_interface, img_interface], ["Document QA", "Image QA"])
+
 # Mount Gradio app in FastAPI
-@app.get("/")
-def redirect_root():
-    return {"message": "Visit /gradio for the interface."}
+app = gr.mount_gradio_app(app, demo, path="/")
 
-app = gr.mount_gradio_app(app, gr_interface, path="/gradio")
+# Redirect to Gradio interface
+@app.get("/")
+def home():
+    return RedirectResponse(url="/")