Spaces:
Running
Running
| import gradio as gr | |
| import PyPDF2 | |
| import io | |
| from transformers import pipeline, AutoTokenizer | |
| import torch | |
| import re | |
| from typing import List, Tuple | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| class PDFSummarizer: | |
| def __init__(self): | |
| # Use a much faster, lighter model for summarization | |
| self.model_name = "sshleifer/distilbart-cnn-12-6" # Much faster than BART-large | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {self.device}") | |
| try: | |
| # Initialize the summarization pipeline with optimizations | |
| self.summarizer = pipeline( | |
| "summarization", | |
| model=self.model_name, | |
| device=0 if self.device == "cuda" else -1, | |
| framework="pt", | |
| model_kwargs={"torch_dtype": torch.float16 if self.device == "cuda" else torch.float32} | |
| ) | |
| # Initialize tokenizer for length calculations | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) | |
| print("Model loaded successfully") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| # Fallback to an even faster model | |
| self.model_name = "facebook/bart-large-cnn" | |
| self.summarizer = pipeline("summarization", model=self.model_name, device=-1) | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) | |
| print("Fallback model loaded") | |
| def extract_text_from_pdf(self, pdf_file) -> str: | |
| """Extract text content from PDF file""" | |
| try: | |
| pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_file)) | |
| text = "" | |
| for page_num, page in enumerate(pdf_reader.pages): | |
| page_text = page.extract_text() | |
| if page_text.strip(): | |
| text += f"\n--- Page {page_num + 1} ---\n" | |
| text += page_text | |
| return text.strip() | |
| except Exception as e: | |
| raise Exception(f"Error extracting text from PDF: {str(e)}") | |
| def clean_text(self, text: str) -> str: | |
| """Clean and preprocess text""" | |
| # Remove extra whitespaces and newlines | |
| text = re.sub(r'\s+', ' ', text) | |
| # Remove special characters but keep punctuation | |
| text = re.sub(r'[^\w\s.,!?;:()\-"]', ' ', text) | |
| # Remove page markers | |
| text = re.sub(r'--- Page \d+ ---', '', text) | |
| return text.strip() | |
| def chunk_text(self, text: str, max_chunk_length: int = 512) -> List[str]: | |
| """Split text into smaller, more manageable chunks for faster processing""" | |
| sentences = text.split('. ') | |
| chunks = [] | |
| current_chunk = "" | |
| for sentence in sentences: | |
| # Check if adding this sentence would exceed the limit | |
| potential_chunk = current_chunk + sentence + ". " | |
| # Use faster length estimation | |
| if len(potential_chunk.split()) <= max_chunk_length: | |
| current_chunk = potential_chunk | |
| else: | |
| if current_chunk: | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = sentence + ". " | |
| if current_chunk: | |
| chunks.append(current_chunk.strip()) | |
| # Limit number of chunks for speed | |
| return chunks[:5] # Process max 5 chunks for speed | |
| def summarize_chunk(self, chunk: str, max_length: int = 100, min_length: int = 30) -> str: | |
| """Summarize a single chunk of text with speed optimizations""" | |
| try: | |
| # Speed optimizations | |
| summary = self.summarizer( | |
| chunk, | |
| max_length=max_length, | |
| min_length=min_length, | |
| do_sample=False, | |
| truncation=True, | |
| early_stopping=True, | |
| num_beams=2 # Reduced from default 4 for speed | |
| ) | |
| return summary[0]['summary_text'] | |
| except Exception as e: | |
| return f"Error summarizing chunk: {str(e)}" | |
| def process_pdf(self, pdf_file, summary_type: str) -> Tuple[str, str, str]: | |
| """Main function to process PDF and generate summary""" | |
| try: | |
| # Extract text from PDF | |
| raw_text = self.extract_text_from_pdf(pdf_file) | |
| if not raw_text.strip(): | |
| return "❌ Error: No text could be extracted from the PDF.", "", "" | |
| # Clean the text | |
| cleaned_text = self.clean_text(raw_text) | |
| # Calculate text statistics | |
| word_count = len(cleaned_text.split()) | |
| char_count = len(cleaned_text) | |
| if word_count < 50: | |
| return "❌ Error: PDF contains too little text to summarize.", "", "" | |
| # Chunk the text for processing | |
| chunks = self.chunk_text(cleaned_text) | |
| # Determine summary parameters based on type (optimized for speed) | |
| if summary_type == "Brief (Quick)": | |
| max_len, min_len = 60, 20 | |
| elif summary_type == "Detailed": | |
| max_len, min_len = 100, 40 | |
| else: # Comprehensive | |
| max_len, min_len = 150, 60 | |
| # Summarize each chunk (with progress tracking) | |
| chunk_summaries = [] | |
| for i, chunk in enumerate(chunks): | |
| print(f"Processing chunk {i+1}/{len(chunks)}") | |
| summary = self.summarize_chunk(chunk, max_len, min_len) | |
| chunk_summaries.append(summary) | |
| # Combine summaries | |
| combined_summary = " ".join(chunk_summaries) | |
| # Skip final summarization for speed if we have few chunks | |
| if len(chunks) <= 2: | |
| final_summary = combined_summary | |
| else: | |
| # Quick final summary for multiple chunks | |
| final_summary = self.summarize_chunk( | |
| combined_summary, | |
| max_length=min(200, max_len * 1.5), | |
| min_length=min_len | |
| ) | |
| # Create statistics | |
| summary_stats = f""" | |
| 📊 **Document Statistics:** | |
| - Original word count: {word_count:,} | |
| - Original character count: {char_count:,} | |
| - Pages processed: {len(chunks)} | |
| - Summary word count: {len(final_summary.split()):,} | |
| - Compression ratio: {word_count / len(final_summary.split()):.1f}:1 | |
| """ | |
| return final_summary, summary_stats, "✅ Summary generated successfully!" | |
| except Exception as e: | |
| return f"❌ Error processing PDF: {str(e)}", "", "" | |
| # Initialize the summarizer | |
| pdf_summarizer = PDFSummarizer() | |
| def summarize_pdf_interface(pdf_file, summary_type): | |
| """Gradio interface function""" | |
| if pdf_file is None: | |
| return "❌ Please upload a PDF file.", "", "" | |
| try: | |
| # Read the uploaded file - pdf_file is already the file path | |
| with open(pdf_file, 'rb') as f: | |
| pdf_content = f.read() | |
| # Process the PDF | |
| summary, stats, status = pdf_summarizer.process_pdf(pdf_content, summary_type) | |
| return summary, stats, status | |
| except Exception as e: | |
| return f"❌ Error: {str(e)}", "", "" | |
| # Create Gradio interface | |
| def create_interface(): | |
| with gr.Blocks( | |
| title="📄 AI PDF Summarizer", | |
| theme=gr.themes.Soft(), | |
| css=""" | |
| .gradio-container { | |
| max-width: 1200px !important; | |
| } | |
| .summary-box { | |
| border-left: 4px solid #2196F3; | |
| padding: 16px; | |
| background-color: #f8f9fa; | |
| } | |
| """ | |
| ) as interface: | |
| gr.Markdown(""" | |
| # ✨ Your AI-Powered PDF Assistant | |
| ### Transform lengthy documents into clear, concise summaries in seconds. | |
| Upload any PDF file below and let our intelligent model do the work for you. Perfect for students, researchers, and professionals who need to quickly grasp the core ideas of any document. | |
| --- | |
| ## Key Features | |
| - ⚡ **Lightning-Fast Summaries**: Powered by a state-of-the-art AI model. | |
| - 🧠 **Intelligent Text Chunking**: Handles long documents with ease by breaking them down into manageable pieces. | |
| - 📊 **Detailed Document Stats**: Get instant insights on word count, compression ratio, and more. | |
| - 🎯 **Customizable Summary Length**: Choose the level of detail that works best for you. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| pdf_input = gr.File( | |
| label="📁 Upload PDF File", | |
| file_types=[".pdf"], | |
| type="filepath" | |
| ) | |
| summary_type = gr.Radio( | |
| choices=["Brief (Quick)", "Detailed", "Comprehensive"], | |
| value="Detailed", | |
| label="📏 Summary Length", | |
| info="Choose how detailed you want the summary to be" | |
| ) | |
| summarize_btn = gr.Button( | |
| "🚀 Generate Summary", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| status_output = gr.Textbox( | |
| label="📋 Status", | |
| interactive=False, | |
| max_lines=2 | |
| ) | |
| with gr.Column(scale=2): | |
| summary_output = gr.Textbox( | |
| label="📝 Generated Summary", | |
| lines=15, | |
| max_lines=20, | |
| interactive=False, | |
| elem_classes=["summary-box"] | |
| ) | |
| stats_output = gr.Markdown( | |
| label="📊 Document Statistics", | |
| value="Upload a PDF to see statistics" | |
| ) | |
| # Examples section | |
| gr.Markdown(""" | |
| --- | |
| ## 💡 Tips for Optimal Results: | |
| - **Quality Matters**: Ensure your PDF has selectable text (not just scanned images). | |
| - **Length**: Works best with documents between 500-10,000 words. | |
| - **Language**: Optimized for English content. | |
| - **Format**: Clean, well-formatted PDFs produce the best summaries. | |
| ## 🔧 How It Works: | |
| This application uses a fine-tuned **BART** model for summarization. The process involves: | |
| 1. Extracting text from your PDF. | |
| 2. Cleaning and preprocessing the text. | |
| 3. Intelligently chunking the document to handle long texts. | |
| 4. Generating a summary for each chunk and then combining them for a final, coherent result. | |
| """) | |
| # Connect the button to the function | |
| summarize_btn.click( | |
| fn=summarize_pdf_interface, | |
| inputs=[pdf_input, summary_type], | |
| outputs=[summary_output, stats_output, status_output] | |
| ) | |
| # Auto-process when file is uploaded | |
| pdf_input.change( | |
| fn=summarize_pdf_interface, | |
| inputs=[pdf_input, summary_type], | |
| outputs=[summary_output, stats_output, status_output] | |
| ) | |
| return interface | |
| # Launch the application | |
| if __name__ == "__main__": | |
| interface = create_interface() | |
| interface.launch() | |