Spaces:
Running
Running
| 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] | |
| ) | |
| 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="/") | |
| def redirect_to_interface(): | |
| return RedirectResponse(url="/") |