Update app.py

app.py

@@ -122,38 +122,32 @@ if __name__ == "__main__":
 import gradio as gr
 import uvicorn
 import numpy as np
-import pymupdf 
+import pymupdf  # PyMuPDF
 import tika
 import torch
 from fastapi import FastAPI
-from transformers import pipeline, AutoTokenizer
+from transformers import pipeline
 from PIL import Image
 from io import BytesIO
 from starlette.responses import RedirectResponse
 from tika import parser
 from openpyxl import load_workbook
 from pptx import Presentation
-import easyocr
 import os
 
-tika.initVM() 
-
+tika.initVM()
 app = FastAPI()
 
-# Load models
 device = "cuda" if torch.cuda.is_available() else "cpu"
 qa_pipeline = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=device)
 image_captioning_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-reader = easyocr.Reader(["en"])  
-
-ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx", "png", "jpg", "jpeg"}
+ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx"}
 
 def validate_file_type(file):
-    if file is None:
+    if not file:
         return "❌ No file uploaded!"
-    if isinstance(file, str):  
+    if isinstance(file, str) and os.path.exists(file):
         return None
     if hasattr(file, "name"):
         ext = file.name.split(".")[-1].lower()
@@ -162,7 +156,6 @@ def validate_file_type(file):
         return None
     return "❌ Invalid file format!"
 
-# ✅ Extract Text from PDF
 def extract_text_from_pdf(file_bytes):
     try:
         doc = pymupdf.open(stream=file_bytes, filetype="pdf")
@@ -170,7 +163,6 @@ def extract_text_from_pdf(file_bytes):
     except Exception as e:
         return f"❌ PDF Error: {str(e)}"
 
-# ✅ Extract Text from DOCX & PPTX using Tika
 def extract_text_with_tika(file_bytes):
     try:
         parsed = parser.from_buffer(file_bytes)
@@ -178,7 +170,6 @@ def extract_text_with_tika(file_bytes):
     except Exception as e:
         return f"❌ Tika Error: {str(e)}"
 
-# ✅ Extract Text from Excel
 def extract_text_from_excel(file_bytes):
     try:
         wb = load_workbook(BytesIO(file_bytes), data_only=True)
@@ -190,7 +181,6 @@ def extract_text_from_excel(file_bytes):
     except Exception as e:
         return f"❌ Excel Error: {str(e)}"
 
-# ✅ Extract Text from PPTX
 def extract_text_from_pptx(file_bytes):
     try:
         ppt = Presentation(BytesIO(file_bytes))
@@ -203,41 +193,15 @@ def extract_text_from_pptx(file_bytes):
     except Exception as e:
         return f"❌ PPTX Error: {str(e)}"
 
-# ✅ Extract Text from Image using OCR
-def extract_text_from_image(image_file):
-    try:
-        image = Image.open(image_file).convert("RGB")
-        np_image = np.array(image)
-        
-        if np_image.std() < 10:  # Low contrast check
-            return "⚠️ No meaningful content detected in the image."
-        
-        result = reader.readtext(np_image)
-        return " ".join([res[1] for res in result]) if result else "⚠️ No text found."
-    except Exception as e:
-        return f"❌ Image OCR Error: {str(e)}"
-
-# ✅ Truncate Long Text for Model
-def truncate_text(text, max_tokens=450):
-    tokens = tokenizer.tokenize(text)
-    return tokenizer.convert_tokens_to_string(tokens[:max_tokens])
+def truncate_text(text, max_length=2048):
+    return text[:max_length] if len(text) > max_length else text
 
-# ✅ Answer Questions from Image or Document
 def answer_question(file, question: str):
     try:
-        # ✅ Handle Image Files (Gradio sends images as NumPy arrays)
-        if isinstance(file, np.ndarray):  
-            image = Image.fromarray(file)
-            caption = image_captioning_pipeline(image)[0]['generated_text']
-            response = qa_pipeline(f"Question: {question}\nContext: {caption}")
-            return response[0]["generated_text"]
-
-        # ✅ Validate File
         validation_error = validate_file_type(file)
         if validation_error:
             return validation_error
 
-        # ✅ Read File Bytes
         file_bytes = None
         file_ext = None
 
@@ -246,24 +210,21 @@ def answer_question(file, question: str):
             with open(file, "rb") as f:
                 file_bytes = f.read()
         elif hasattr(file, "read"):
-            file_ext = file.name.split(".")[-1].lower() if hasattr(file, "name") else None
+            file_ext = file.name.split(".")[-1].lower()
             file_bytes = file.read()
         else:
-            return "❌ Unexpected file type received!"
+            return f"❌ Unexpected file type received! Type: {type(file)}"
 
-        # ✅ Extract Text Based on File Type
         if file_ext == "pdf":
             text = extract_text_from_pdf(file_bytes)
         elif file_ext in ["docx", "pptx"]:
             text = extract_text_with_tika(file_bytes)
         elif file_ext == "xlsx":
             text = extract_text_from_excel(file_bytes)
-        elif file_ext in ["png", "jpg", "jpeg"]:
-            text = extract_text_from_image(BytesIO(file_bytes))
         else:
             return f"❌ Unsupported file format: {file_ext}"
 
-        if not text or "⚠️" in text:
+        if not text or "❌" in text:
             return f"⚠️ No text extracted. Error: {text}"
 
         truncated_text = truncate_text(text)
@@ -274,17 +235,15 @@ def answer_question(file, question: str):
     except Exception as e:
         return f"❌ Processing Error: {str(e)}"
 
-# ✅ Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("## 📄 AI-Powered Document & Image QA")
+    gr.Markdown("## 📄 AI-Powered Document QA")
     with gr.Row():
-        file_input = gr.File(label="Upload Document / Image")
+        file_input = gr.File(label="Upload Document")
         question_input = gr.Textbox(label="Ask a Question", placeholder="What is this document about?")
     answer_output = gr.Textbox(label="Answer")
     submit_btn = gr.Button("Get Answer")
     submit_btn.click(answer_question, inputs=[file_input, question_input], outputs=answer_output)
 
-# ✅ Mount Gradio with FastAPI
 app = gr.mount_gradio_app(app, demo, path="/")
 
 @app.get("/")