Update app.py

app.py

@@ -1,14 +1,16 @@
-import sys, os, json, shutil, re, time, gc, hashlib
+import sys, os, json, shutil, re, time, gc
 import pandas as pd
 from datetime import datetime
 from typing import List, Tuple, Dict, Union
 import gradio as gr
+from concurrent.futures import ThreadPoolExecutor
 
 # Constants
 MAX_MODEL_TOKENS = 131072
 MAX_NEW_TOKENS = 4096
-MAX_CHUNK_TOKENS = 8192  # IMPORTANT: Split input into 8k tokens chunks
+MAX_CHUNK_TOKENS = 8192
 PROMPT_OVERHEAD = 300
+BATCH_SIZE = 2  # NEW: batch 2 prompts together for faster processing
 
 # Paths
 persistent_dir = "/data/hf_cache"
@@ -41,20 +43,17 @@ def clean_response(text: str) -> str:
 def extract_text_from_excel(path: str) -> str:
     all_text = []
     xls = pd.ExcelFile(path)
-
     for sheet_name in xls.sheet_names:
         try:
             df = xls.parse(sheet_name).astype(str).fillna("")
         except Exception:
             continue
-
         for idx, row in df.iterrows():
-            non_empty = [cell.strip() for cell in row if cell.strip() != ""]
+            non_empty = [cell.strip() for cell in row if cell.strip()]
             if len(non_empty) >= 2:
                 text_line = " | ".join(non_empty)
                 if len(text_line) > 15:
                     all_text.append(f"[{sheet_name}] {text_line}")
-
     return "\n".join(all_text)
 
 def split_text(text: str, max_tokens=MAX_CHUNK_TOKENS) -> List[str]:
@@ -73,6 +72,9 @@ def split_text(text: str, max_tokens=MAX_CHUNK_TOKENS) -> List[str]:
         chunks.append("\n".join(current))
     return chunks
 
+def batch_chunks(chunks: List[str], batch_size: int = 2) -> List[List[str]]:
+    return [chunks[i:i+batch_size] for i in range(0, len(chunks), batch_size)]
+
 def build_prompt(chunk: str) -> str:
     return f"""### Unstructured Clinical Records\n\nAnalyze the clinical notes below and summarize with:\n- Diagnostic Patterns\n- Medication Issues\n- Missed Opportunities\n- Inconsistencies\n- Follow-up Recommendations\n\n---\n\n{chunk}\n\n---\nRespond concisely in bullet points with clinical reasoning."""
 
@@ -92,19 +94,16 @@ def init_agent() -> TxAgent:
     agent.init_model()
     return agent
 
-def analyze_serial(agent, chunks: List[str]) -> List[str]:
+def analyze_batches(agent, batches: List[List[str]]) -> List[str]:
     results = []
-    for idx, chunk in enumerate(chunks):
-        prompt = build_prompt(chunk)
-        if estimate_tokens(prompt) > MAX_MODEL_TOKENS:
-            results.append(f"❌ Chunk {idx+1} too long. Skipped.")
-            continue
+    for batch in batches:
+        prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
         response = ""
         try:
             for r in agent.run_gradio_chat(
                 message=prompt,
                 history=[],
-                temperature=0.2,
+                temperature=0.0,
                 max_new_tokens=MAX_NEW_TOKENS,
                 max_token=MAX_MODEL_TOKENS,
                 call_agent=False,
@@ -120,7 +119,8 @@ def analyze_serial(agent, chunks: List[str]) -> List[str]:
                     response += r.content
             results.append(clean_response(response))
         except Exception as e:
-            results.append(f"❌ Error in chunk {idx+1}: {str(e)}")
+            results.append(f"❌ Error in batch: {str(e)}")
+    torch.cuda.empty_cache()
     gc.collect()
     return results
 
@@ -130,7 +130,7 @@ def generate_final_summary(agent, combined: str) -> str:
     for r in agent.run_gradio_chat(
         message=final_prompt,
         history=[],
-        temperature=0.2,
+        temperature=0.0,
         max_new_tokens=MAX_NEW_TOKENS,
         max_token=MAX_MODEL_TOKENS,
         call_agent=False,
@@ -155,13 +155,14 @@ def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Di
     try:
         extracted = extract_text_from_excel(file.name)
         chunks = split_text(extracted)
-        messages.append({"role": "assistant", "content": f"🔍 Split into {len(chunks)} chunks. Analyzing..."})
+        batches = batch_chunks(chunks, batch_size=BATCH_SIZE)
+        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Analyzing..."})
 
-        chunk_results = analyze_serial(agent, chunks)
-        valid = [res for res in chunk_results if not res.startswith("❌")]
+        batch_results = analyze_batches(agent, batches)
+        valid = [res for res in batch_results if not res.startswith("❌")]
 
         if not valid:
-            messages.append({"role": "assistant", "content": "❌ No valid chunk outputs."})
+            messages.append({"role": "assistant", "content": "❌ No valid batch outputs."})
             return messages, None
 
         summary = generate_final_summary(agent, "\n\n".join(valid))