app.py

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import fitz  # PyMuPDF
import docx
import pptx
import openpyxl
import re
import nltk
from nltk.tokenize import sent_tokenize
import torch
from fastapi import FastAPI
from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
from gtts import gTTS
import tempfile
import os
import easyocr
from fpdf import FPDF
import datetime
from concurrent.futures import ThreadPoolExecutor
import hashlib

nltk.download('punkt', quiet=True)

app = FastAPI()

MODEL_NAME = "facebook/bart-large-cnn"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
model.eval()
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=-1, batch_size=4)

reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
executor = ThreadPoolExecutor()

summary_cache = {}

def clean_text(text: str) -> str:
    text = re.sub(r'\s+', ' ', text)
    text = re.sub(r'\u2022\s*|\d\.\s+', '', text)
    text = re.sub(r'\[.*?\]|\(.*?\)', '', text)
    text = re.sub(r'\bPage\s*\d+\b', '', text, flags=re.IGNORECASE)
    return text.strip()

def extract_text(file_path: str, file_extension: str):
    try:
        if file_extension == "pdf":
            with fitz.open(file_path) as doc:
                text = "\n".join(page.get_text("text") for page in doc)
                if len(text.strip()) < 50:
                    images = [page.get_pixmap() for page in doc]
                    temp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
                    images[0].save(temp_img.name)
                    ocr_result = reader.readtext(temp_img.name, detail=0)
                    os.unlink(temp_img.name)
                    text = "\n".join(ocr_result) if ocr_result else text
                return clean_text(text), ""

        elif file_extension == "docx":
            doc = docx.Document(file_path)
            return clean_text("\n".join(p.text for p in doc.paragraphs)), ""

        elif file_extension == "pptx":
            prs = pptx.Presentation(file_path)
            text = [shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")]
            return clean_text("\n".join(text)), ""

        elif file_extension == "xlsx":
            wb = openpyxl.load_workbook(file_path, read_only=True)
            text = [" ".join(str(cell) for cell in row if cell) for sheet in wb.sheetnames for row in wb[sheet].iter_rows(values_only=True)]
            return clean_text("\n".join(text)), ""

        elif file_extension in ["jpg", "jpeg", "png"]:
            ocr_result = reader.readtext(file_path, detail=0)
            return clean_text("\n".join(ocr_result)), ""

        return "", "Unsupported file format"
    except Exception as e:
        return "", f"Error reading {file_extension.upper()} file: {str(e)}"

def chunk_text(text: str, max_tokens: int = 950):
    try:
        sentences = sent_tokenize(text)
    except:
        words = text.split()
        sentences = [' '.join(words[i:i+20]) for i in range(0, len(words), 20)]

    chunks = []
    current_chunk = ""
    for sentence in sentences:
        token_length = len(tokenizer.encode(current_chunk + " " + sentence))
        if token_length <= max_tokens:
            current_chunk += " " + sentence
        else:
            chunks.append(current_chunk.strip())
            current_chunk = sentence

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks

def generate_summary(text: str, length: str = "medium") -> str:
    cache_key = hashlib.md5((text + length).encode()).hexdigest()
    if cache_key in summary_cache:
        return summary_cache[cache_key]

    length_params = {
        "short": {"max_length": 80, "min_length": 30},
        "medium": {"max_length": 200, "min_length": 80},
        "long": {"max_length": 300, "min_length": 210}
    }
    chunks = chunk_text(text)
    try:
        summaries = summarizer(
            chunks,
            max_length=length_params[length]["max_length"],
            min_length=length_params[length]["min_length"],
            do_sample=False,
            truncation=True,
            no_repeat_ngram_size=2,
            num_beams=2,
            early_stopping=True
        )
        summary_texts = [s['summary_text'] for s in summaries]
    except Exception as e:
        summary_texts = [f"[Batch error: {str(e)}]"]

    final_summary = " ".join(summary_texts)
    final_summary = ". ".join(s.strip().capitalize() for s in final_summary.split(". ") if s.strip())
    final_summary = final_summary if len(final_summary) > 25 else "Summary too short - document may be too brief"

    summary_cache[cache_key] = final_summary
    return final_summary

def text_to_speech(text: str):
    try:
        tts = gTTS(text)
        temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        tts.save(temp_audio.name)
        return temp_audio.name
    except Exception as e:
        print(f"Error in text-to-speech: {e}")
        return ""

def create_pdf(summary: str, original_filename: str):
    try:
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", size=12)
        pdf.set_font("Arial", 'B', 16)
        pdf.cell(200, 10, txt="Document Summary", ln=1, align='C')
        pdf.set_font("Arial", size=12)
        pdf.cell(200, 10, txt=f"Original file: {original_filename}", ln=1)
        pdf.cell(200, 10, txt=f"Generated on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=1)
        pdf.ln(10)
        pdf.multi_cell(0, 10, txt=summary)
        temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
        pdf.output(temp_pdf.name)
        return temp_pdf.name
    except Exception as e:
        print(f"Error creating PDF: {e}")
        return ""

def summarize_document(file, summary_length: str, enable_tts: bool = True):
    if file is None:
        return "Please upload a document first", "", None, None
    file_path = file.name
    file_extension = file_path.split(".")[-1].lower()
    original_filename = os.path.basename(file_path)
    text, error = extract_text(file_path, file_extension)
    if error:
        return error, "", None, None
    if not text or len(text.split()) < 30:
        return "Document is too short or contains too little text to summarize", "", None, None
    try:
        summary = generate_summary(text, summary_length)
        audio_path = text_to_speech(summary) if enable_tts else None
        pdf_path = create_pdf(summary, original_filename) if summary else None
        return summary, "", audio_path, pdf_path
    except Exception as e:
        return f"Summarization error: {str(e)}", "", None, None

with gr.Blocks(title="Document Summarizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📄 Advanced Document Summarizer")
    gr.Markdown("Upload a document to generate a summary with audio and optional PDF download")

    with gr.Row():
        with gr.Column():
            file_input = gr.File(
                label="Upload Document",
                file_types=[".pdf", ".docx", ".pptx", ".xlsx", ".jpg", ".jpeg", ".png"],
                type="filepath"
            )
            length_radio = gr.Radio(
                ["short", "medium", "long"],
                value="medium",
                label="Summary Length"
            )
            submit_btn = gr.Button("Generate Summary", variant="primary")

        with gr.Column():
            output = gr.Textbox(label="Summary", lines=10)
            audio_output = gr.Audio(label="Audio Summary")
            pdf_download = gr.File(label="Download Summary as PDF", visible=False)

    def summarize_and_return_ui(file, summary_length):
        summary, _, audio_path, pdf_path = summarize_document(file, summary_length)
        return (
            summary,
            audio_path,
            gr.File(visible=pdf_path is not None, value=pdf_path)
        )

    submit_btn.click(
        fn=summarize_and_return_ui,
        inputs=[file_input, length_radio],
        outputs=[output, audio_output, pdf_download]
    )

@app.get("/files/{file_name}")
async def get_file(file_name: str):
    file_path = os.path.join(tempfile.gettempdir(), file_name)
    if os.path.exists(file_path):
        return FileResponse(file_path)
    return JSONResponse({"error": "File not found"}, status_code=404)

app = gr.mount_gradio_app(app, demo, path="/")

@app.get("/")
def redirect_to_interface():
    return RedirectResponse(url="/")