import sys import os import pandas as pd import gradio as gr from typing import List, Tuple, Dict, Any, Union, Generator import shutil import re from datetime import datetime import time import asyncio import logging from concurrent.futures import ThreadPoolExecutor, as_completed # Configuration and setup persistent_dir = "/data/hf_cache" os.makedirs(persistent_dir, exist_ok=True) 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 directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]: os.makedirs(directory, 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 # Updated token limits as specified MAX_MODEL_TOKENS = 131072 # TxAgent's max token limit MAX_CHUNK_TOKENS = 32768 # Larger chunks to reduce number of chunks MAX_NEW_TOKENS = 512 # Optimized for fast generation PROMPT_OVERHEAD = 500 # Estimated tokens for prompt template MAX_CONCURRENT = 8 # High concurrency for A100 80GB # Setup logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def clean_response(text: str) -> str: try: text = text.encode('utf-8', 'surrogatepass').decode('utf-8') except UnicodeError: text = text.encode('utf-8', 'replace').decode('utf-8') 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 estimate_tokens(text: str) -> int: return len(text) // 3.5 + 1 # More conservative estimate def extract_text_from_excel(file_path: str) -> str: all_text = [] try: xls = pd.ExcelFile(file_path) for sheet_name in xls.sheet_names: df = xls.parse(sheet_name) df = df.astype(str).fillna("") rows = df.apply(lambda row: " | ".join(row), axis=1) sheet_text = [f"[{sheet_name}] {line}" for line in rows] all_text.extend(sheet_text) except Exception as e: logger.error(f"Error extracting Excel: {str(e)}") raise ValueError(f"Failed to process Excel file: {str(e)}") return "\n".join(all_text) def split_text_into_chunks(text: str) -> List[str]: """Split text into chunks respecting MAX_CHUNK_TOKENS and PROMPT_OVERHEAD""" effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD if effective_max <= 0: raise ValueError("Effective max tokens must be positive") lines = text.split("\n") chunks = [] current_chunk = [] current_tokens = 0 for line in lines: line_tokens = estimate_tokens(line) if current_tokens + line_tokens > effective_max: if current_chunk: chunks.append("\n".join(current_chunk)) current_chunk = [line] current_tokens = line_tokens else: current_chunk.append(line) current_tokens += line_tokens if current_chunk: chunks.append("\n".join(current_chunk)) logger.info(f"Split text into {len(chunks)} chunks") return chunks def build_prompt_from_text(chunk: str) -> str: return f""" ### Unstructured Clinical Records You are reviewing unstructured, mixed-format clinical documentation from various forms, tables, and sheets. **Objective:** Identify patterns, missed diagnoses, inconsistencies, and follow-up gaps. Here is the extracted content chunk: {chunk} Please analyze the above and provide concise responses (max {MAX_NEW_TOKENS} tokens): - Diagnostic Patterns - Medication Issues - Missed Opportunities - Inconsistencies - Follow-up Recommendations """ def init_agent(): default_tool_path = os.path.abspath("data/new_tool.json") target_tool_path = os.path.join(tool_cache_dir, "new_tool.json") if not os.path.exists(target_tool_path): shutil.copy(default_tool_path, target_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": target_tool_path}, force_finish=True, enable_checker=True, step_rag_num=4, seed=100, additional_default_tools=[], max_model_tokens=MAX_MODEL_TOKENS # Pass the updated token limit ) agent.init_model() return agent async def process_chunk(agent: TxAgent, chunk: str, chunk_idx: int) -> Tuple[int, str]: """Process a single chunk with error handling""" try: prompt = build_prompt_from_text(chunk) prompt_tokens = estimate_tokens(prompt) if prompt_tokens > MAX_MODEL_TOKENS: logger.warning(f"Chunk {chunk_idx} prompt too long ({prompt_tokens} tokens)") return chunk_idx, "" response = "" for result in agent.run_gradio_chat( message=prompt, history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[], ): if isinstance(result, str): response += result elif hasattr(result, "content"): response += result.content elif isinstance(result, list): for r in result: if hasattr(r, "content"): response += r.content return chunk_idx, clean_response(response) except Exception as e: logger.error(f"Error processing chunk {chunk_idx}: {str(e)}") return chunk_idx, "" async def process_file(agent: TxAgent, file_path: str) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]: """Process the entire file and yield progress updates""" messages = [] report_path = None try: # Initial messages messages.append({"role": "user", "content": f"Processing file: {os.path.basename(file_path)}"}) messages.append({"role": "assistant", "content": "⏳ Extracting data from Excel..."}) yield messages, None # Extract and chunk text start_time = time.time() text = extract_text_from_excel(file_path) chunks = split_text_into_chunks(text) messages.append({"role": "assistant", "content": f"✅ Extracted {len(chunks)} chunks in {time.time()-start_time:.1f}s"}) yield messages, None # Process chunks in parallel chunk_responses = [None] * len(chunks) with ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor: futures = [] for idx, chunk in enumerate(chunks): future = executor.submit( lambda c, i: asyncio.run(process_chunk(agent, c, i)), chunk, idx ) futures.append(future) messages.append({"role": "assistant", "content": f"🔍 Processing chunk {idx+1}/{len(chunks)}..."}) yield messages, None for future in as_completed(futures): idx, response = future.result() chunk_responses[idx] = response messages.append({"role": "assistant", "content": f"✅ Chunk {idx+1} processed"}) yield messages, None # Combine and summarize combined = "\n\n".join([r for r in chunk_responses if r]) messages.append({"role": "assistant", "content": "📊 Generating final report..."}) yield messages, None final_response = "" for result in agent.run_gradio_chat( message=f"Summarize these clinical findings:\n\n{combined}", history=[], temperature=0.2, max_new_tokens=MAX_NEW_TOKENS*2, # Allow more tokens for summary max_token=MAX_MODEL_TOKENS, call_agent=False, conversation=[], ): if isinstance(result, str): final_response += result elif hasattr(result, "content"): final_response += result.content elif isinstance(result, list): for r in result: if hasattr(r, "content"): final_response += r.content messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(final_response)}" yield messages, None # Save report final_report = f"# Final Clinical Report\n\n{clean_response(final_response)}" timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') report_path = os.path.join(report_dir, f"report_{timestamp}.md") with open(report_path, 'w') as f: f.write(final_report) messages.append({"role": "assistant", "content": f"✅ Report saved: report_{timestamp}.md"}) yield messages, report_path except Exception as e: logger.error(f"Processing failed: {str(e)}") messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"}) yield messages, None def create_ui(agent: TxAgent): """Create the Gradio interface""" with gr.Blocks(title="Clinical Analysis", css=".gradio-container {max-width: 900px}") as demo: gr.Markdown("## 🏥 Clinical Data Analysis (TxAgent)") with gr.Row(): with gr.Column(scale=3): chatbot = gr.Chatbot( label="Analysis Progress", show_copy_button=True, height=600, type="messages" ) with gr.Column(scale=1): file_input = gr.File( label="Upload Excel File", file_types=[".xlsx"], height=100 ) analyze_btn = gr.Button( "🧠 Analyze Data", variant="primary" ) report_output = gr.File( label="Download Report", visible=False ) analyze_btn.click( fn=lambda file: process_file(agent, file.name) if file else ([{"role": "assistant", "content": "❌ Please upload a file"}], None), inputs=[file_input], outputs=[chatbot, report_output] ) return demo if __name__ == "__main__": try: agent = init_agent() demo = create_ui(agent) demo.launch( server_name="0.0.0.0", server_port=7860, show_error=True, allowed_paths=[report_dir], share=False ) except Exception as e: logger.error(f"Application failed: {str(e)}") sys.exit(1)