import sys import os import json import shutil import re import gc import time from datetime import datetime from typing import List, Tuple, Dict, Union import pandas as pd import gradio as gr import torch # === 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") file_cache_dir = os.path.join(persistent_dir, "cache") report_dir = os.path.join(persistent_dir, "reports") for d in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]: os.makedirs(d, exist_ok=True) os.environ["HF_HOME"] = model_cache_dir os.environ["TRANSFORMERS_CACHE"] = model_cache_dir current_dir = os.path.dirname(os.path.abspath(__file__)) src_path = os.path.abspath(os.path.join(current_dir, "src")) sys.path.insert(0, src_path) from txagent.txagent import TxAgent # === 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) return text.strip() 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()] 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]: effective_limit = max_tokens - PROMPT_OVERHEAD chunks, current, current_tokens = [], [], 0 for line in text.split("\n"): tokens = estimate_tokens(line) if current_tokens + tokens > effective_limit: if current: chunks.append("\n".join(current)) current, current_tokens = [line], tokens else: current.append(line) current_tokens += tokens if current: chunks.append("\n".join(current)) return chunks 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: 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.""" def init_agent() -> TxAgent: tool_path = os.path.join(tool_cache_dir, "new_tool.json") if not os.path.exists(tool_path): shutil.copy(os.path.abspath("data/new_tool.json"), tool_path) agent = TxAgent( model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B", rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B", tool_files_dict={"new_tool": tool_path}, force_finish=True, enable_checker=True, step_rag_num=4, seed=100 ) agent.init_model() return agent # === Main Processing === def analyze_batches(agent, batches: List[List[str]]) -> List[str]: results = [] 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." final_response = "" for r in agent.run_gradio_chat( message=final_prompt, history=[], temperature=0.0, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[] ): if isinstance(r, str): final_response += r elif isinstance(r, list): for m in r: if hasattr(m, "content"): final_response += m.content elif hasattr(r, "content"): 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"): messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file."}) return messages, None messages.append({"role": "user", "content": f"📂 Processing file: {os.path.basename(file.name)}"}) try: 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..."}) 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 batch outputs."}) return messages, None summary = generate_final_summary(agent, "\n\n".join(valid)) report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md") with open(report_path, 'w', encoding='utf-8') as f: f.write(f"# 🧠 Final Medical Report\n\n{summary}") messages.append({"role": "assistant", "content": f"📊 Final Report:\n\n{summary}"}) messages.append({"role": "assistant", "content": f"✅ Report saved: {os.path.basename(report_path)}"}) return messages, report_path except Exception as e: messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"}) return messages, None def create_ui(agent): with gr.Blocks(css=""" html, body, .gradio-container {background: #0e1621; color: #e0e0e0;} button#component-2 { background: #1e88e5 !important; border: none !important; color: white !important; } button#component-2:hover { background: #1565c0 !important; } """) as demo: gr.Markdown("""

📄 CPS: Clinical Patient Support System

Analyze and summarize unstructured medical files using AI (optimized for A100 GPU).

""") 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", elem_id="analyze-button") download = gr.File(label="Download Report", visible=False, interactive=False) state = gr.State(value=[]) def handle_analysis(file, chat): messages, report_path = process_report(agent, file, chat) 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__": 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)