Update app.py

app.py

@@ -1,100 +1,46 @@
-import numpy as np
-import streamlit as st
-from transformers import BartTokenizer, BartForConditionalGeneration
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.cluster import KMeans
-from nltk.tokenize import sent_tokenize
-import nltk
-import PyPDF2
-from io import BytesIO
-
-# --------------------------- DOWNLOAD NLTK DATA ---------------------------
-@st.cache_resource
-def download_nltk_data():
-    try:
-        nltk.data.find('tokenizers/punkt')
-    except LookupError:
-        nltk.download('punkt', quiet=True)
-        nltk.download('punkt_tab', quiet=True)
-
-download_nltk_data()
-
-# --------------------------- LOAD MODEL ---------------------------
-@st.cache_resource
-def load_model():
-    tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
-    model = BartForConditionalGeneration.from_pretrained('facebook/bart-large-cnn')
-    return tokenizer, model
-
-tokenizer, model = load_model()
-
-# --------------------------- HELPER FUNCTIONS ---------------------------
-
-def extract_text_from_pdf(file) -> str:
-    """Extract text from uploaded PDF file."""
-    pdf_reader = PyPDF2.PdfReader(BytesIO(file.read()))
-    text = ""
-    for page in pdf_reader.pages:
-        text += page.extract_text() or ""
-    return text.strip()
-
-def summarize(text, max_length=150, min_length=50):
-    """Summarize a given text using BART."""
-    inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors='pt')
-    summary_ids = model.generate(
+def summarize_large_text(text, max_length=150, min_length=50):
+    """Summarize long text by splitting into chunks and combining summaries."""
+    chunks = chunk_text(text, max_words=800)
+    summaries = []
+    for chunk in chunks:
+        inputs = tokenizer([chunk], max_length=1024, truncation=True, return_tensors='pt')
+        summary_ids = model.generate(
+            inputs['input_ids'],
+            num_beams=4,
+            max_length=max_length,
+            min_length=min_length,
+            early_stopping=True
+        )
+        chunk_summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+        summaries.append(chunk_summary)
+
+    # Combine all chunk summaries and re-summarize them
+    combined_summary_text = " ".join(summaries)
+    inputs = tokenizer([combined_summary_text], max_length=1024, truncation=True, return_tensors='pt')
+    final_ids = model.generate(
         inputs['input_ids'],
         num_beams=4,
-        max_length=max_length,
-        min_length=min_length,
+        max_length=200,
+        min_length=80,
         early_stopping=True
     )
-    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    final_summary = tokenizer.decode(final_ids[0], skip_special_tokens=True)
+    return final_summary
 
-def cluster_documents(documents, n_clusters=3):
-    """Cluster similar documents using TF-IDF + KMeans."""
-    vectorizer = TfidfVectorizer(stop_words='english')
-    X = vectorizer.fit_transform(documents)
-    kmeans = KMeans(n_clusters=n_clusters, random_state=42).fit(X)
-    return kmeans.labels_
-
-def chunk_text(text, max_words=1000):
-    """Split long text into smaller chunks for summarization."""
-    if len(text.split()) <= max_words:
-        return [text]
-    
-    sentences = sent_tokenize(text)
-    chunks, current_chunk, current_word_count = [], [], 0
-
-    for sentence in sentences:
-        sentence_words = len(sentence.split())
-        if current_word_count + sentence_words <= max_words:
-            current_chunk.append(sentence)
-            current_word_count += sentence_words
-        else:
-            chunks.append(" ".join(current_chunk))
-            current_chunk = [sentence]
-            current_word_count = sentence_words
-
-    if current_chunk:
-        chunks.append(" ".join(current_chunk))
-
-    return chunks
 
 def multi_document_summarize(documents):
-    """Summarize multiple related documents using clustering + BART."""
+    """Summarize multiple related documents using clustering + improved BART summarization."""
     results = {
         'individual_summaries': [],
         'cluster_summaries': [],
         'final_summary': None
     }
 
-    
-    # 1️⃣ Individual summaries
+    # 1️⃣ Individual summaries (using improved function)
     for doc in documents:
-        chunks = chunk_text(doc)
-        doc_summary = " ".join([summarize(chunk) for chunk in chunks])
+        doc_summary = summarize_large_text(doc)
         results['individual_summaries'].append(doc_summary)
-    
+
     # 2️⃣ Clustering (if >1 doc)
     if len(documents) > 1:
         n_clusters = min(3, len(documents))
@@ -103,8 +49,7 @@ def multi_document_summarize(documents):
         for cluster_id in np.unique(clusters):
             cluster_docs = [doc for doc, c in zip(documents, clusters) if c == cluster_id]
             combined_text = " ".join(cluster_docs)
-            chunks = chunk_text(combined_text)
-            cluster_summary = " ".join([summarize(chunk) for chunk in chunks])
+            cluster_summary = summarize_large_text(combined_text)
             results['cluster_summaries'].append({
                 'doc_indices': [i for i, c in enumerate(clusters) if c == cluster_id],
                 'summary': cluster_summary
@@ -112,13 +57,12 @@ def multi_document_summarize(documents):
         
         # 3️⃣ Final overall summary
         all_summaries = results['individual_summaries'] + [cs['summary'] for cs in results['cluster_summaries']]
-        results['final_summary'] = summarize(" ".join(all_summaries), max_length=200, min_length=100)
+        results['final_summary'] = summarize_large_text(" ".join(all_summaries))
     else:
         results['final_summary'] = results['individual_summaries'][0]
     
     return results
-    #download
-    st.download_button("Download Summary", results['final_summary'], file_name="summary.txt")
+
 
 
  # --------------------------- STREAMLIT UI ---------------------------