Update app.py

app.py

@@ -1,24 +1,27 @@
+# Optimized app.py for A100 GPU (safe parallel batching + no stuck + max performance)
+
 import sys
 import os
-import gc
 import json
 import shutil
 import re
 import time
-import pandas as pd
-import gradio as gr
-import torch
-from typing import List, Tuple, Dict, Union
+import gc
+import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Tuple, Dict, Union
 from datetime import datetime
+import pandas as pd
+import gradio as gr
 
 # Constants
 MAX_MODEL_TOKENS = 131072
 MAX_NEW_TOKENS = 4096
 MAX_CHUNK_TOKENS = 8192
 PROMPT_OVERHEAD = 300
-BATCH_SIZE = 4  # 4 chunks per batch
-MAX_WORKERS = 6  # 6 parallel batches
+BATCH_SIZE = 2  # Safer for vLLM
+MAX_PARALLEL_JOBS = 2  # Max threads launched in parallel
+SLEEP_BETWEEN_JOBS = 0.5  # Seconds
 
 # Paths
 persistent_dir = "/data/hf_cache"
@@ -39,6 +42,7 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
+# Utility functions
 def estimate_tokens(text: str) -> int:
     return len(text) // 4 + 1
 
@@ -56,12 +60,12 @@ def extract_text_from_excel(path: str) -> str:
             df = xls.parse(sheet_name).astype(str).fillna("")
         except Exception:
             continue
-        for _, row in df.iterrows():
+        for idx, row in df.iterrows():
             non_empty = [cell.strip() for cell in row if cell.strip()]
             if len(non_empty) >= 2:
-                line = " | ".join(non_empty)
-                if len(line) > 15:
-                    all_text.append(f"[{sheet_name}] {line}")
+                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]:
@@ -80,7 +84,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 = 4) -> List[List[str]]:
+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:
@@ -102,12 +106,13 @@ def init_agent() -> TxAgent:
     agent.init_model()
     return agent
 
-def analyze_batch(agent, batch: List[str]) -> str:
-    prompt = "\n\n".join(build_prompt(chunk) for chunk in batch)
+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=prompt,
+            message=joined_prompt,
             history=[],
             temperature=0.0,
             max_new_tokens=MAX_NEW_TOKENS,
@@ -123,19 +128,21 @@ def analyze_batch(agent, batch: List[str]) -> str:
                         response += m.content
             elif hasattr(r, "content"):
                 response += r.content
+        return clean_response(response)
     except Exception as e:
-        return f"❌ Error in batch: {str(e)}"
-    finally:
-        torch.cuda.empty_cache()
-        gc.collect()
-    return clean_response(response)
+        return f"❌ Error: {str(e)}"
 
 def analyze_batches_parallel(agent, batches: List[List[str]]) -> List[str]:
     results = []
-    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
-        futures = [executor.submit(analyze_batch, agent, batch) for batch in batches]
+    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())
+    torch.cuda.empty_cache()
+    gc.collect()
     return results
 
 def generate_final_summary(agent, combined: str) -> str:
@@ -170,7 +177,7 @@ 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. Analyzing in parallel..."})
+        messages.append({"role": "assistant", "content": f"🔍 Split into {len(batches)} batches. Parallel analyzing..."})
 
         batch_results = analyze_batches_parallel(agent, batches)
         valid = [res for res in batch_results if not res.startswith("❌")]
@@ -194,20 +201,25 @@ def process_report(agent, file, messages: List[Dict[str, str]]) -> Tuple[List[Di
 
 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;}
-    .gr-chatbot, .gr-markdown, .gr-file-upload {border-radius: 16px; background-color: #1b2533;}
-    .gr-chatbot .message {font-size: 16px; padding: 12px; border-radius: 18px;}
-    .gr-chatbot .message.user {background-color: #334155;}
-    .gr-chatbot .message.assistant {background-color: #1e293b;}
+    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><p>Upload a file and analyze medical notes.</p>""")
+        gr.Markdown("""<h2>📄 CPS: Clinical Patient Support System</h2>""")
         with gr.Column():
             chatbot = gr.Chatbot(label="CPS Assistant", height=700, type="messages")
             upload = gr.File(label="Upload Medical File", file_types=[".xlsx"])
             analyze = gr.Button("🧠 Analyze", variant="primary")
             download = gr.File(label="Download Report", visible=False, interactive=False)
+
         state = gr.State(value=[])
 
         def handle_analysis(file, chat):
@@ -225,4 +237,4 @@ if __name__ == "__main__":
         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)
+        sys.exit(1)