Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import fitz
+import fitz  # PyMuPDF
 import torch
 from transformers import pipeline
 import time, logging, re
@@ -8,33 +8,29 @@ matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 import io
 from PIL import Image
-import nltk
-
-# Download punkt tokenizer if not already
-nltk.download('punkt', quiet=True)
-from nltk.tokenize import sent_tokenize
 
+# Logging setup
 logging.basicConfig(level=logging.ERROR)
-device = -1  # CPU-only
-print("⚠️ Optimized CPU-only version.")
+print("⚠️ Running in CPU-only mode.")
 
+# Load summarization model
 try:
-    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
+    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=-1)
 except Exception as e:
-    print(f"❌ Model loading failed: {str(e)}")
-    exit(1)
+    raise RuntimeError(f"Model loading failed: {str(e)}")
 
 def visualize_chunk_status(chunk_data):
-    status_colors = {'summarized': 'green', 'skipped': 'orange', 'error': 'red'}
+    colors = {'summarized': 'green', 'skipped': 'orange', 'error': 'red'}
     labels = [f"C{i['chunk']}" for i in chunk_data]
-    colors = [status_colors.get(i['status'], 'gray') for i in chunk_data]
+    bar_colors = [colors.get(i['status'], 'gray') for i in chunk_data]
     times = [i.get('time', 0.1) for i in chunk_data]
 
     fig, ax = plt.subplots(figsize=(10, 2.5))
-    ax.barh(labels, times, color=colors)
+    ax.barh(labels, times, color=bar_colors)
     ax.set_xlabel("Time (s)")
     ax.set_title("📊 Chunk Processing Status")
     plt.tight_layout()
+
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
@@ -42,34 +38,21 @@ def visualize_chunk_status(chunk_data):
     return Image.open(buf)
 
 def create_summary_flowchart(summaries):
-    # Filter only successful summaries
-    filtered = []
-    for s in summaries:
-        if s.startswith("**Chunk") and "Skipped" not in s and "Error" not in s:
-            parts = s.split("**:", 1)
-            if len(parts) > 1:
-                filtered.append(parts[1].strip())
-
-    if not filtered:
+    blocks = [s for s in summaries if s.startswith("**Chunk") and "Skipped" not in s and "Error" not in s]
+    if not blocks:
         return None
 
-    fig_height = max(2, len(filtered) * 1.5)
-    fig, ax = plt.subplots(figsize=(6, fig_height))
+    fig, ax = plt.subplots(figsize=(6, len(blocks) * 0.8 + 1))
     ax.axis('off')
-
-    ypos = list(range(len(filtered) * 2, 0, -2))
+    ypos = list(range(len(blocks) * 2, 0, -2))
     boxprops = dict(boxstyle="round,pad=0.5", facecolor="lightblue", edgecolor="black")
 
-    for i, (y, summary) in enumerate(zip(ypos, filtered)):
-        summary_text = summary.replace("\n", " ").strip()[:120]
-        if len(summary_text) == 120:
-            summary_text += "..."
-
-        ax.text(0.5, y, summary_text, ha='center', va='center', bbox=boxprops, fontsize=9)
-
-        if i < len(filtered) - 1:
-            ax.annotate('', xy=(0.5, y - 1.5), xytext=(0.5, y - 0.5),
-                        arrowprops=dict(arrowstyle="->", lw=1.5))
+    for i, (y, summary) in enumerate(zip(ypos, blocks)):
+        text = summary.split("**Chunk")[-1].replace("**:", ":").split("\n", 1)[-1].strip()[:120]
+        ax.text(0.5, y, text + ("..." if len(text) > 100 else ""),
+                ha='center', va='center', bbox=boxprops, fontsize=9)
+        if i < len(blocks) - 1:
+            ax.annotate('', xy=(0.5, y - 1), xytext=(0.5, y - 0.2), arrowprops=dict(arrowstyle="->", lw=1.5))
 
     buf = io.BytesIO()
     plt.tight_layout()
@@ -78,31 +61,13 @@ def create_summary_flowchart(summaries):
     plt.close(fig)
     return Image.open(buf)
 
-def split_text_into_chunks(text, max_tokens=1500):
-    sentences = sent_tokenize(text)
-    chunks = []
-    current_chunk = ""
-
-    for sentence in sentences:
-        if len(current_chunk) + len(sentence) <= max_tokens:
-            current_chunk += " " + sentence
-        else:
-            chunks.append(current_chunk.strip())
-            current_chunk = sentence
-
-    if current_chunk:
-        chunks.append(current_chunk.strip())
-
-    return chunks[:20]  # Limit to 20 chunks max
-
 def summarize_file(file_bytes):
     start = time.time()
-    chunk_info = []
-    summaries = []
+    chunk_info, summaries = [], []
 
     try:
         doc = fitz.open(stream=file_bytes, filetype="pdf")
-        text = "".join(page.get_text("text") for page in doc)
+        text = " ".join(page.get_text("text") for page in doc)
         text = re.sub(r"\$\s*([^$]+)\s*\$", r"\1", text)
         text = re.sub(r"\\cap", "intersection", text)
         text = re.sub(r"\s+", " ", text).strip()
@@ -110,34 +75,31 @@ def summarize_file(file_bytes):
     except Exception as e:
         return f"❌ Text extraction failed: {str(e)}", None, None
 
-    if not text.strip():
-        return "❌ No text found", None, None
-
-    chunks = split_text_into_chunks(text)
+    if not text:
+        return "❌ No text found in PDF.", None, None
 
+    chunks = [text[i:i+1500] for i in range(0, min(len(text), 30000), 1500)]
     for i, chunk in enumerate(chunks):
-        chunk_start = time.time()
-        chunk_result = {'chunk': i + 1, 'status': '', 'time': 0}
+        t0 = time.time()
+        info = {'chunk': i + 1, 'status': '', 'time': 0}
 
         if sum(1 for c in chunk if not c.isalnum()) / len(chunk) > 0.5:
             summaries.append(f"**Chunk {i+1}**: Skipped (equation-heavy)")
-            chunk_result['status'] = 'skipped'
+            info['status'] = 'skipped'
         else:
             try:
                 summary = summarizer(chunk, max_length=80, min_length=15, do_sample=False)[0]['summary_text']
                 summaries.append(f"**Chunk {i+1}**:\n{summary}")
-                chunk_result['status'] = 'summarized'
+                info['status'] = 'summarized'
             except Exception as e:
                 summaries.append(f"**Chunk {i+1}**: ❌ Error: {str(e)}")
-                chunk_result['status'] = 'error'
+                info['status'] = 'error'
 
-        chunk_result['time'] = time.time() - chunk_start
-        chunk_info.append(chunk_result)
+        info['time'] = time.time() - t0
+        chunk_info.append(info)
 
     final_summary = f"**Processed chunks**: {len(chunks)}\n**Time**: {time.time() - start:.2f}s\n\n" + "\n\n".join(summaries)
-    process_img = visualize_chunk_status(chunk_info)
-    flow_img = create_summary_flowchart(summaries)
-    return final_summary, process_img, flow_img
+    return final_summary, visualize_chunk_status(chunk_info), create_summary_flowchart(summaries)
 
 demo = gr.Interface(
     fn=summarize_file,
@@ -152,10 +114,8 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    try:
-        demo.launch(share=False, server_port=7860)
-    except Exception as e:
-        print(f"❌ Gradio launch failed: {str(e)}")
+    demo.launch(server_port=7860, share=False)
+