app.py

"""import gradio as gr
import uvicorn
import numpy as np
import fitz  # PyMuPDF
import tika
import torch
from fastapi import FastAPI
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

# Initialize Tika for DOCX & PPTX parsing
tika.initVM()

# Initialize FastAPI
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")

ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx"}

# ✅ Function to Validate File Type
def validate_file_type(file):
    if isinstance(file, str):  # Text-based input (NamedString)
        return None
    if hasattr(file, "name"):
        ext = file.name.split(".")[-1].lower()
        if ext not in ALLOWED_EXTENSIONS:
            return f"❌ Unsupported file format: {ext}"
        return None
    return "❌ Invalid file format!"

# ✅ Extract Text from PDF
def extract_text_from_pdf(pdf_bytes):
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    return "\n".join([page.get_text() for page in doc])

# ✅ Extract Text from DOCX & PPTX using Tika
def extract_text_with_tika(file_bytes):
    return parser.from_buffer(file_bytes)["content"]

# ✅ Extract Text from Excel
def extract_text_from_excel(file_bytes):
    wb = load_workbook(BytesIO(file_bytes), data_only=True)
    text = []
    for sheet in wb.worksheets:
        for row in sheet.iter_rows(values_only=True):
            text.append(" ".join(str(cell) for cell in row if cell))
    return "\n".join(text)

# ✅ Truncate Long Text for Model
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):
    # Image Processing (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

    file_ext = file.name.split(".")[-1].lower() if hasattr(file, "name") else None
    file_bytes = file.read() if hasattr(file, "read") else None
    if not file_bytes:
        return "❌ Could not read file content!"

    # Extract Text from Supported Documents
    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)
    else:
        return "❌ Unsupported file format!"

    if not text:
        return "⚠️ No text extracted from the document."

    truncated_text = truncate_text(text)
    response = qa_pipeline(f"Question: {question}\nContext: {truncated_text}")

    return response[0]["generated_text"]

# ✅ Gradio Interface (Unified for Images & Documents)
with gr.Blocks() as demo:
    gr.Markdown("## 📄 AI-Powered Document & Image QA")

    with gr.Row():
        file_input = gr.File(label="Upload Document / Image")
        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("/")
def home():
    return RedirectResponse(url="/")

# ✅ Run FastAPI + Gradio
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)
""" 
import gradio as gr
import uvicorn
import numpy as np
import pymupdf 
import tika
import torch
from fastapi import FastAPI
from transformers import pipeline, AutoTokenizer
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() 

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"}

def validate_file_type(file):
    if file is None:
        return "❌ No file uploaded!"
    if isinstance(file, str):  
        return None
    if hasattr(file, "name"):
        ext = file.name.split(".")[-1].lower()
        if ext not in ALLOWED_EXTENSIONS:
            return f"❌ Unsupported file format: {ext}"
        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")
        return "\n".join([page.get_text("text") for page in doc])
    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)
        return parsed.get("content", "⚠️ No text found.").strip()
    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)
        text = []
        for sheet in wb.worksheets:
            for row in sheet.iter_rows(values_only=True):
                text.append(" ".join(str(cell) for cell in row if cell))
        return "\n".join(text) if text else "⚠️ No text found."
    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))
        text = []
        for slide in ppt.slides:
            for shape in slide.shapes:
                if hasattr(shape, "text"):
                    text.append(shape.text)
        return "\n".join(text) if text else "⚠️ No text found."
    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])

# ✅ 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

        if isinstance(file, str) and os.path.exists(file):
            file_ext = file.split(".")[-1].lower()
            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_bytes = file.read()
        else:
            return "❌ Unexpected file type received!"

        # ✅ 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:
            return f"⚠️ No text extracted. Error: {text}"

        truncated_text = truncate_text(text)
        response = qa_pipeline(f"Question: {question}\nContext: {truncated_text}")

        return response[0]["generated_text"]

    except Exception as e:
        return f"❌ Processing Error: {str(e)}"

# ✅ Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 📄 AI-Powered Document & Image QA")
    with gr.Row():
        file_input = gr.File(label="Upload Document / Image")
        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("/")
def home():
    return RedirectResponse(url="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)