Update app.py

app.py

@@ -11,10 +11,10 @@ from PIL import Image
 
 logging.basicConfig(level=logging.ERROR)
 device = -1  # CPU-only
-print("⚠️ CPU-only. Expect ~20–30s for 300,000 chars.")
+print("⚠️ Optimized CPU-only version.")
 
 try:
-    summarizer = pipeline("summarization", model="t5-small", device=device, torch_dtype=torch.float32)
+    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
 except Exception as e:
     print(f"❌ Model loading failed: {str(e)}")
     exit(1)
@@ -30,7 +30,6 @@ def visualize_chunk_status(chunk_data):
     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)
@@ -38,18 +37,23 @@ def visualize_chunk_status(chunk_data):
     return Image.open(buf)
 
 def create_summary_flowchart(summaries):
-    fig, ax = plt.subplots(figsize=(6, len(summaries) * 0.8 + 1))
+    filtered = [s for s in summaries if s.startswith("**Chunk") and "Skipped" not in s and "Error" not in s]
+    if not filtered:
+        return None
+
+    fig, ax = plt.subplots(figsize=(6, len(filtered) * 0.8 + 1))
     ax.axis('off')
-    
-    ypos = list(range(len(summaries) * 2, 0, -2))
+
+    ypos = list(range(len(filtered) * 2, 0, -2))
     boxprops = dict(boxstyle="round,pad=0.5", facecolor="lightblue", edgecolor="black")
 
-    for i, (y, summary) in enumerate(zip(ypos, summaries)):
-        summary_text = summary.split("**Chunk")[1] if summary.startswith("**Chunk") else summary
-        summary_text = summary_text.strip().replace("**:", ":")[:120] + ("..." if len(summary) > 120 else "")
-        ax.text(0.5, y, summary_text, ha='center', va='center', bbox=boxprops, fontsize=9, wrap=True)
+    for i, (y, summary) in enumerate(zip(ypos, filtered)):
+        summary_text = summary.split("**Chunk")[1]
+        summary_text = summary_text.replace("**:", ":").split("\n", 1)[-1].strip()[:120]
+        ax.text(0.5, y, summary_text + ("..." if len(summary_text) > 100 else ""),
+                ha='center', va='center', bbox=boxprops, fontsize=9)
 
-        if i < len(summaries) - 1:
+        if i < len(filtered) - 1:
             ax.annotate('', xy=(0.5, y - 1), xytext=(0.5, y - 0.2),
                         arrowprops=dict(arrowstyle="->", lw=1.5))
 
@@ -63,6 +67,7 @@ def create_summary_flowchart(summaries):
 def summarize_file(file_bytes):
     start = time.time()
     chunk_info = []
+    summaries = []
 
     try:
         doc = fitz.open(stream=file_bytes, filetype="pdf")
@@ -77,24 +82,17 @@ def summarize_file(file_bytes):
     if not text.strip():
         return "❌ No text found", None, None
 
-    text = text[:300000]
-    chunks = [text[i:i+2000] for i in range(0, len(text), 2000)]
-    summaries = []
-
+    chunks = [text[i:i+1500] for i in range(0, min(len(text), 30000), 1500)]  # max 20 chunks
     for i, chunk in enumerate(chunks):
         chunk_start = time.time()
         chunk_result = {'chunk': i + 1, 'status': '', 'time': 0}
 
-        if time.time() - start > 20:
-            summaries.append("⚠️ Stopped early")
-            break
-
         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'
         else:
             try:
-                summary = summarizer(chunk, max_length=60, min_length=10, do_sample=False)[0]['summary_text']
+                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'
             except Exception as e:
@@ -104,7 +102,7 @@ def summarize_file(file_bytes):
         chunk_result['time'] = time.time() - chunk_start
         chunk_info.append(chunk_result)
 
-    final_summary = f"**Chars**: {len(text)}\n**Time**: {time.time()-start:.2f}s\n\n" + "\n\n".join(summaries)
+    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
@@ -113,12 +111,12 @@ demo = gr.Interface(
     fn=summarize_file,
     inputs=gr.File(label="📄 Upload PDF", type="binary"),
     outputs=[
-        gr.Textbox(label="📝 Summarized Output"),
+        gr.Textbox(label="📝 Summary"),
         gr.Image(label="📊 Chunk Status", type="pil"),
-        gr.Image(label="🔁 Flowchart Summary", type="pil")
+        gr.Image(label="🔁 Flow Summary", type="pil")
     ],
-    title="AI-Powered PDF Summarizer",
-    description="Summarizes long PDFs (up to 300,000 characters) and visualizes chunk processing + flow of content."
+    title="📘 PDF Summarizer with Visual Flow",
+    description="Summarizes up to 30,000 characters from a PDF. Includes chunk status and flowchart visualizations."
 )
 
 if __name__ == "__main__":
@@ -128,3 +126,4 @@ if __name__ == "__main__":
         print(f"❌ Gradio launch failed: {str(e)}")
 
 
+