Update app.py

app.py

@@ -1,29 +1,17 @@
-# Optimized app.py for A100 GPU (safe parallel batching + no stuck + max performance)
-
 import sys
 import os
 import json
 import shutil
 import re
-import time
 import gc
-import threading
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import List, Tuple, Dict, Union
+import time
 from datetime import datetime
+from typing import List, Tuple, Dict, Union
 import pandas as pd
 import gradio as gr
+import torch
 
-# Constants
-MAX_MODEL_TOKENS = 131072
-MAX_NEW_TOKENS = 4096
-MAX_CHUNK_TOKENS = 8192
-PROMPT_OVERHEAD = 300
-BATCH_SIZE = 2  # Safer for vLLM
-MAX_PARALLEL_JOBS = 2  # Max threads launched in parallel
-SLEEP_BETWEEN_JOBS = 0.5  # Seconds
-
-# Paths
+# === Configuration ===
 persistent_dir = "/data/hf_cache"
 model_cache_dir = os.path.join(persistent_dir, "txagent_models")
 tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
@@ -42,14 +30,21 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Utility functions
+# === Constants ===
+MAX_MODEL_TOKENS = 131072
+MAX_NEW_TOKENS = 4096
+MAX_CHUNK_TOKENS = 8192
+BATCH_SIZE = 2
+PROMPT_OVERHEAD = 300
+SAFE_SLEEP = 0.5  # seconds between batches
+
+# === Utility Functions ===
 def estimate_tokens(text: str) -> int:
     return len(text) // 4 + 1
 
 def clean_response(text: str) -> str:
     text = re.sub(r"\[.*?\]|\bNone\b", "", text, flags=re.DOTALL)
     text = re.sub(r"\n{3,}", "\n\n", text)
-    text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
     return text.strip()
 
 def extract_text_from_excel(path: str) -> str:
@@ -84,7 +79,7 @@ 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]]:
+def batch_chunks(chunks: List[str], batch_size: int = BATCH_SIZE) -> List[List[str]]:
     return [chunks[i:i+batch_size] for i in range(0, len(chunks), batch_size)]
 
 def build_prompt(chunk: str) -> str:
@@ -106,48 +101,42 @@ def init_agent() -> TxAgent:
     agent.init_model()
     return agent
 
-def process_single_batch(agent, batch: List[str]) -> str:
-    prompts = [build_prompt(chunk) for chunk in batch]
-    joined_prompt = "\n\n".join(prompts)
-    response = ""
-    try:
-        for r in agent.run_gradio_chat(
-            message=joined_prompt,
-            history=[],
-            temperature=0.0,
-            max_new_tokens=MAX_NEW_TOKENS,
-            max_token=MAX_MODEL_TOKENS,
-            call_agent=False,
-            conversation=[]
-        ):
-            if isinstance(r, str):
-                response += r
-            elif isinstance(r, list):
-                for m in r:
-                    if hasattr(m, "content"):
-                        response += m.content
-            elif hasattr(r, "content"):
-                response += r.content
-        return clean_response(response)
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
-
-def analyze_batches_parallel(agent, batches: List[List[str]]) -> List[str]:
+# === Main Processing ===
+def analyze_batches(agent, batches: List[List[str]]) -> List[str]:
     results = []
-    with ThreadPoolExecutor(max_workers=MAX_PARALLEL_JOBS) as executor:
-        futures = []
-        for batch in batches:
-            futures.append(executor.submit(process_single_batch, agent, batch))
-            time.sleep(SLEEP_BETWEEN_JOBS)
-        for future in as_completed(futures):
-            results.append(future.result())
+    for batch in batches:
+        prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
+        try:
+            batch_response = ""
+            for r in agent.run_gradio_chat(
+                message=prompt,
+                history=[],
+                temperature=0.0,
+                max_new_tokens=MAX_NEW_TOKENS,
+                max_token=MAX_MODEL_TOKENS,
+                call_agent=False,
+                conversation=[]
+            ):
+                if isinstance(r, str):
+                    batch_response += r
+                elif isinstance(r, list):
+                    for m in r:
+                        if hasattr(m, "content"):
+                            batch_response += m.content
+                elif hasattr(r, "content"):
+                    batch_response += r.content
+            results.append(clean_response(batch_response))
+            time.sleep(SAFE_SLEEP)
+        except Exception as e:
+            results.append(f"❌ Batch failed: {str(e)}")
+            time.sleep(SAFE_SLEEP * 2)  # longer sleep on error
     torch.cuda.empty_cache()
     gc.collect()
     return results
 
 def generate_final_summary(agent, combined: str) -> str:
-    final_prompt = f"""Provide a structured medical report based on the following summaries:\n\n{combined}\n\nRespond in detailed medical bullet points."""
-    full_report = ""
+    final_prompt = f"Provide a structured medical report based on the following summaries:\n\n{combined}\n\nRespond in detailed medical bullet points."
+    final_response = ""
     for r in agent.run_gradio_chat(
         message=final_prompt,
         history=[],
@@ -158,14 +147,14 @@ def generate_final_summary(agent, combined: str) -> str:
         conversation=[]
     ):
         if isinstance(r, str):
-            full_report += r
+            final_response += r
         elif isinstance(r, list):
             for m in r:
                 if hasattr(m, "content"):
-                    full_report += m.content
+                    final_response += m.content
         elif hasattr(r, "content"):
-            full_report += r.content
-    return clean_response(full_report)
+            final_response += r.content
+    return clean_response(final_response)
 
 def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
     if not file or not hasattr(file, "name"):
@@ -177,9 +166,9 @@ def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Di
         extracted = extract_text_from_excel(file.name)
         chunks = split_text(extracted)
         batches = batch_chunks(chunks, batch_size=BATCH_SIZE)
-        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Parallel analyzing..."})
+        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Analyzing..."})
 
-        batch_results = analyze_batches_parallel(agent, batches)
+        batch_results = analyze_batches(agent, batches)
         valid = [res for res in batch_results if not res.startswith("❌")]
 
         if not valid:
@@ -200,20 +189,11 @@ def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Di
         return messages, None
 
 def create_ui(agent):
-    with gr.Blocks(css="""
-    html, body, .gradio-container {
-        background-color: #0e1621;
-        color: #e0e0e0;
-        font-family: 'Inter', sans-serif;
-    }
-    h2, h3, h4 { color: #89b4fa; font-weight: 600; }
-    button.gr-button-primary {
-        background-color: #007bff !important;
-        color: white !important;
-        font-weight: bold;
-    }
-    """) as demo:
-        gr.Markdown("""<h2>📄 CPS: Clinical Patient Support System</h2>""")
+    with gr.Blocks(css="""html, body, .gradio-container {background: #0e1621; color: #e0e0e0;}""") as demo:
+        gr.Markdown("""
+        <h2>📄 CPS: Clinical Patient Support System</h2>
+        <p>Analyze and summarize unstructured medical files using AI (optimized for A100 GPU).</p>
+        """)
         with gr.Column():
             chatbot = gr.Chatbot(label="CPS Assistant", height=700, type="messages")
             upload = gr.File(label="Upload Medical File", file_types=[".xlsx"])
@@ -227,14 +207,10 @@ def create_ui(agent):
             return messages, gr.update(visible=bool(report_path), value=report_path), messages
 
         analyze.click(fn=handle_analysis, inputs=[upload, state], outputs=[chatbot, download, state])
-
     return demo
 
+# === Main ===
 if __name__ == "__main__":
-    try:
-        agent = init_agent()
-        ui = create_ui(agent)
-        ui.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)
-    except Exception as err:
-        print(f"Startup failed: {err}")
-        sys.exit(1)
+    agent = init_agent()
+    ui = create_ui(agent)
+    ui.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)