Update app.py

app.py

@@ -9,17 +9,24 @@ import matplotlib.pyplot as plt
 import io
 from PIL import Image
 from concurrent.futures import ThreadPoolExecutor
+import numpy as np
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Set device (CPU or GPU)
+# Set device and optimize for speed
 device = 0 if torch.cuda.is_available() else -1
 logger.info(f"🔧 Using {'GPU' if device == 0 else 'CPU'}")
 
-# Load model
+# Load a lighter model for faster inference
 try:
-    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
+    summarizer = pipeline(
+        "summarization",
+        model="t5-small",  # Lightweight model
+        device=device,
+        framework="pt",
+        torch_dtype=torch.float16 if device == 0 else torch.float32  # Half-precision on GPU
+    )
 except Exception as e:
     logger.error(f"❌ Model loading failed: {str(e)}")
     exit(1)
@@ -30,46 +37,36 @@ def visualize_chunk_status(chunk_data):
     colors = [status_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)
+    fig, ax = plt.subplots(figsize=(8, 2))  # Smaller figure size
+    ax.barh(labels, times, color=colors, height=0.4)  # Reduced bar height
     ax.set_xlabel("Time (s)")
-    ax.set_title("📊 Chunk Processing Status")
-    plt.tight_layout()
+    ax.set_title("Chunk Status")
+    plt.tight_layout(pad=0.5)  # Minimal padding
     buf = io.BytesIO()
-    plt.savefig(buf, format='png')
+    plt.savefig(buf, format='png', dpi=100)  # Lower DPI for speed
     plt.close(fig)
     buf.seek(0)
     return Image.open(buf)
 
 def create_summary_flowchart(summaries):
-    filtered = [
-        s for s in summaries 
-        if s.startswith("**Chunk") and "Skipped" not in s and "Error" not in s
-    ]
+    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_height = max(2, len(filtered) * 0.8 + 1)
-    fig, ax = plt.subplots(figsize=(6, fig_height))
+    fig_height = max(1.5, len(filtered) * 0.5)  # Reduced height scaling
+    fig, ax = plt.subplots(figsize=(5, fig_height))
     ax.axis('off')
 
-    ypos = list(range(len(filtered) * 2, 0, -2))
-    boxprops = dict(boxstyle="round,pad=0.5", facecolor="lightblue", edgecolor="black")
-
+    ypos = np.arange(len(filtered) * 1.5, 0, -1.5)  # Tighter spacing
     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()
-        if len(summary_text) > 120:
-            summary_text = summary_text[:120] + "..."
-        ax.text(0.5, y, summary_text, ha='center', va='center', bbox=boxprops, fontsize=9)
-
+        summary_text = summary.split("**Chunk")[1].split("\n", 1)[-1].strip()[:100] + ("..." if len(summary.split("**Chunk")[1].split("\n", 1)[-1].strip()) > 100 else "")
+        ax.text(0.5, y, summary_text, ha='center', va='center', fontsize=8, bbox=dict(facecolor="lightblue", edgecolor="black", pad=0.2))
         if i < len(filtered) - 1:
-            ax.annotate('', xy=(0.5, y - 1.2), xytext=(0.5, y - 0.3),
-                        arrowprops=dict(arrowstyle="->", lw=1.5))
+            ax.arrow(0.5, y - 0.2, 0, -1.1, head_width=0.02, head_length=0.1, fc='black', ec='black')
 
-    plt.tight_layout()
+    plt.tight_layout(pad=0.1)
     buf = io.BytesIO()
-    fig.savefig(buf, format='png', bbox_inches='tight')
+    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
     plt.close(fig)
     buf.seek(0)
     return Image.open(buf)
@@ -78,16 +75,16 @@ def process_chunk(i, chunk):
     chunk_result = {'chunk': i + 1, 'status': '', 'time': 0}
     start_time = time.time()
 
-    if sum(1 for c in chunk if not c.isalnum()) / len(chunk) > 0.5:
-        result = f"**Chunk {i+1}**: Skipped (equation-heavy)"
+    if not chunk.strip() or sum(1 for c in chunk if not c.isalnum()) / len(chunk) > 0.5:
+        result = f"**Chunk {i+1}**: Skipped (empty or equation-heavy)"
         chunk_result['status'] = 'skipped'
     else:
         try:
-            summary = summarizer(chunk, max_length=80, min_length=15, do_sample=False)[0]['summary_text']
+            summary = summarizer(chunk, max_length=60, min_length=10, do_sample=False)[0]['summary_text']
             result = f"**Chunk {i+1}**:\n{summary}"
             chunk_result['status'] = 'summarized'
         except Exception as e:
-            result = f"**Chunk {i+1}**: ❌ Error: {str(e)}"
+            result = f"**Chunk {i+1}**: Error: {str(e)}"
             chunk_result['status'] = 'error'
 
     chunk_result['time'] = time.time() - start_time
@@ -98,66 +95,91 @@ def summarize_file(file_bytes):
     summaries = []
     chunk_info = []
 
+    # Stream text extraction
     try:
         doc = fitz.open(stream=file_bytes, filetype="pdf")
-        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()
-        text = "".join(c for c in text if ord(c) < 128)
+        text = ""
+        for page in doc:
+            page_text = page.get_text("text")
+            if not page_text.strip():
+                continue
+            text += page_text
+            if len(text) > 30000:  # Early cutoff
+                text = text[:30000]
+                break
+        doc.close()
+        
+        # Fast text cleaning
+        text = re.sub(r"\$\s*[^$]+\s*\$|\\cap|\s+", lambda m: "intersection" if m.group(0) == "\\cap" else " ", text)
+        text = "".join(c for c in text if ord(c) < 128)[:30000]
     except Exception as e:
-        return f"❌ Text extraction failed: {str(e)}", None, None
+        return f"Text extraction failed: {str(e)}", None, None
 
     if not text.strip():
-        return "❌ No text found", None, None
-
-    chunks = [text[i:i+1500] for i in range(0, min(len(text), 30000), 1500)]
-
-    with ThreadPoolExecutor(max_workers=4) as executor:
+        return "No text found", None, None
+
+    # Smaller, sentence-aware chunks
+    chunks = []
+    current_chunk = ""
+    for sentence in re.split(r'(?<=[.!?])\s+', text):
+        if len(current_chunk) + len(sentence) <= 1000:
+            current_chunk += sentence
+        else:
+            if current_chunk:
+                chunks.append(current_chunk)
+            current_chunk = sentence
+        if len(chunks) >= 30:  # Limit chunks
+            break
+    if current_chunk:
+        chunks.append(current_chunk)
+
+    # Dynamic worker count based on CPU/GPU
+    max_workers = min(8, max(2, torch.cuda.device_count() * 4 if device == 0 else 4))
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
         results = list(executor.map(lambda ic: process_chunk(*ic), enumerate(chunks)))
 
     for summary, info in results:
         summaries.append(summary)
         chunk_info.append(info)
 
-    final_summary = f"**Processed chunks**: {len(chunks)}\n**Time**: {time.time() - start:.2f}s\n\n" + "\n\n".join(summaries)
+    final_summary = f"**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
 
 demo = gr.Interface(
     fn=summarize_file,
-    inputs=gr.File(label="📄 Upload PDF", type="binary"),
+    inputs=gr.File(label="Upload PDF", type="binary"),
     outputs=[
-        gr.Textbox(label="📝 Summary", lines=20),
-        gr.Image(label="📊 Chunk Status", type="pil"),
-        gr.Image(label="🔁 Flow Summary", type="pil")
+        gr.Textbox(label="Summary", lines=15),
+        gr.Image(label="Chunk Status", type="pil"),
+        gr.Image(label="Flow Summary", type="pil")
     ],
-    title="📘 PDF Summarizer with Visual Flow",
-    description="Summarizes up to 30,000 characters from a PDF. Includes chunk status and flowchart visualizations."
+    title="PDF Summarizer",
+    description="Summarizes PDFs up to 30,000 characters with visualizations."
 )
 
 if __name__ == "__main__":
     try:
-        logger.info("Starting Gradio application on http://127.0.0.1:7860")
+        logger.info("Starting Gradio on http://127.0.0.1:7860")
         demo.launch(
             share=False,
             server_name="127.0.0.1",
             server_port=7860,
-            debug=True
+            debug=False  # Disable debug for speed
         )
     except Exception as e:
-        logger.error(f"Gradio launch failed on port 7860: {str(e)}")
-        logger.info("Trying alternative port 7861...")
+        logger.error(f"Failed on port 7860: {str(e)}")
+        logger.info("Trying port 7861...")
         try:
             demo.launch(
                 share=False,
                 server_name="127.0.0.1",
                 server_port=7861,
-                debug=True
+                debug=False
             )
         except Exception as e2:
-            logger.error(f"Gradio launch failed on port 7861: {str(e2)}")
+            logger.error(f"Failed on port 7861: {str(e2)}")
             raise