import sys import os import pandas as pd import pdfplumber import json import gradio as gr from typing import List from concurrent.futures import ThreadPoolExecutor, as_completed import hashlib import shutil import re import psutil import subprocess import multiprocessing from functools import partial import time # Persistent directory 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") vllm_cache_dir = os.path.join(persistent_dir, "vllm_cache") for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir, vllm_cache_dir]: os.makedirs(directory, exist_ok=True) os.environ["HF_HOME"] = model_cache_dir os.environ["TRANSFORMERS_CACHE"] = model_cache_dir os.environ["VLLM_CACHE_DIR"] = vllm_cache_dir os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["CUDA_LAUNCH_BLOCKING"] = "1" 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 def sanitize_utf8(text: str) -> str: return text.encode("utf-8", "ignore").decode("utf-8") def file_hash(path: str) -> str: with open(path, "rb") as f: return hashlib.md5(f.read()).hexdigest() def extract_page_range(file_path: str, start_page: int, end_page: int) -> str: """Extract text from a range of PDF pages.""" try: text_chunks = [] with pdfplumber.open(file_path) as pdf: for page in pdf.pages[start_page:end_page]: page_text = page.extract_text() or "" text_chunks.append(f"=== Page {start_page + pdf.pages.index(page) + 1} ===\n{page_text.strip()}") return "\n\n".join(text_chunks) except Exception: return "" def extract_all_pages(file_path: str, progress_callback=None) -> str: """Extract text from all pages of a PDF using parallel processing.""" try: with pdfplumber.open(file_path) as pdf: total_pages = len(pdf.pages) if total_pages == 0: return "" # Use 6 processes (adjust based on CPU cores) num_processes = min(6, multiprocessing.cpu_count()) pages_per_process = max(1, total_pages // num_processes) # Create page ranges for parallel processing ranges = [(i * pages_per_process, min((i + 1) * pages_per_process, total_pages)) for i in range(num_processes)] if ranges[-1][1] != total_pages: ranges[-1] = (ranges[-1][0], total_pages) # Process page ranges in parallel with multiprocessing.Pool(processes=num_processes) as pool: extract_func = partial(extract_page_range, file_path) results = [] for idx, result in enumerate(pool.starmap(extract_func, ranges)): results.append(result) if progress_callback: processed_pages = min((idx + 1) * pages_per_process, total_pages) progress_callback(processed_pages, total_pages) return "\n\n".join(filter(None, results)) except Exception as e: return f"PDF processing error: {str(e)}" def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str: try: h = file_hash(file_path) cache_path = os.path.join(file_cache_dir, f"{h}.json") if os.path.exists(cache_path): with open(cache_path, "r", encoding="utf-8") as f: return f.read() if file_type == "pdf": text = extract_all_pages(file_path, progress_callback) result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"}) elif file_type == "csv": df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str, skip_blank_lines=False, on_bad_lines="skip") content = df.fillna("").astype(str).values.tolist() result = json.dumps({"filename": os.path.basename(file_path), "rows": content}) elif file_type in ["xls", "xlsx"]: try: df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str) except Exception: df = pd.read_excel(file_path, engine="xlrd", header=None, dtype=str) content = df.fillna("").astype(str).values.tolist() result = json.dumps({"filename": os.path.basename(file_path), "rows": content}) else: result = json.dumps({"error": f"Unsupported file type: {file_type}"}) with open(cache_path, "w", encoding="utf-8") as f: f.write(result) return result except Exception as e: return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"}) def log_system_usage(tag=""): try: cpu = psutil.cpu_percent(interval=1) mem = psutil.virtual_memory() print(f"[{tag}] CPU: {cpu}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB") result = subprocess.run( ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,nounits,noheader"], capture_output=True, text=True ) if result.returncode == 0: used, total, util = result.stdout.strip().split(", ") print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%") except Exception as e: print(f"[{tag}] GPU/CPU monitor failed: {e}") def clean_response(text: str) -> str: text = sanitize_utf8(text) text = re.sub(r"\[TOOL_CALLS\].*", "", text, flags=re.DOTALL) text = re.sub(r"\n{3,}", "\n\n", text).strip() return text def init_agent(): print("🔁 Initializing model...") log_system_usage("Before Load") 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=[], ) agent.init_model() log_system_usage("After Load") print("✅ Agent Ready") return agent def create_ui(agent): with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("

🩺 Clinical Oversight Assistant

") chatbot = gr.Chatbot(label="Analysis", height=600, type="messages") file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple") msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False) send_btn = gr.Button("Analyze", variant="primary") download_output = gr.File(label="Download Full Report") def analyze(message: str, history: List[dict], files: List): history.append({"role": "user", "content": message}) history.append({"role": "assistant", "content": "⏳ Extracting text from files..."}) yield history, None extracted = "" file_hash_value = "" if files: # Progress callback for extraction total_pages = 0 processed_pages = 0 def update_extraction_progress(current, total): nonlocal processed_pages, total_pages processed_pages = current total_pages = total animation = ["🌀", "🔄", "⚙️", "🔃"][(int(time.time() * 2) % 4)] history[-1] = {"role": "assistant", "content": f"Extracting text... {animation} Page {processed_pages}/{total_pages}"} return history, None with ThreadPoolExecutor(max_workers=6) as executor: futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files] results = [sanitize_utf8(f.result()) for f in as_completed(futures)] extracted = "\n".join(results) file_hash_value = file_hash(files[0].name) if files else "" history.pop() # Remove extraction message history.append({"role": "assistant", "content": "✅ Text extraction complete."}) yield history, None # Split extracted text into chunks of ~6,000 characters chunk_size = 6000 chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)] combined_response = "" prompt_template = """ You are a medical analysis assistant. Analyze the following patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown format under these headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, and Urgent Follow-up. For each finding, include: - Clinical context (why the issue was missed or relevant details from the record). - Potential risks if unaddressed (e.g., disease progression, adverse events). - Actionable recommendations (e.g., tests, referrals, medication adjustments). If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text, formatted as markdown with bullet points under each heading, and avoids generic or static responses. Patient Record Excerpt (Chunk {0} of {1}): {chunk} ### Missed Diagnoses - ... ### Medication Conflicts - ... ### Incomplete Assessments - ... ### Urgent Follow-up - ... """ try: # Process each chunk and stream results in real-time for chunk_idx, chunk in enumerate(chunks, 1): # Update UI with chunk progress animation = ["🔍", "📊", "🧠", "🔎"][(int(time.time() * 2) % 4)] history.append({"role": "assistant", "content": f"Analyzing records... {animation} Chunk {chunk_idx}/{len(chunks)}"}) yield history, None prompt = prompt_template.format(chunk_idx, len(chunks), chunk=chunk[:4000]) # Truncate to avoid token limits chunk_response = "" for chunk_output in agent.run_gradio_chat( message=prompt, history=[], temperature=0.2, max_new_tokens=1024, max_token=4096, call_agent=False, conversation=[], ): if chunk_output is None: continue if isinstance(chunk_output, list): for m in chunk_output: if hasattr(m, 'content') and m.content: cleaned = clean_response(m.content) if cleaned: chunk_response += cleaned + "\n" # Update UI with partial response if history[-1]["content"].startswith("Analyzing"): history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"} else: history[-1]["content"] = f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}" yield history, None elif isinstance(chunk_output, str) and chunk_output.strip(): cleaned = clean_response(chunk_output) if cleaned: chunk_response += cleaned + "\n" # Update UI with partial response if history[-1]["content"].startswith("Analyzing"): history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"} else: history[-1]["content"] = f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}" yield history, None # Append completed chunk response to combined response combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n" # Finalize UI with complete response if combined_response: history[-1]["content"] = combined_response.strip() else: history.append({"role": "assistant", "content": "No oversights identified."}) # Generate report file report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None if report_path: with open(report_path, "w", encoding="utf-8") as f: f.write(combined_response) yield history, report_path if report_path and os.path.exists(report_path) else None except Exception as e: print("🚨 ERROR:", e) history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"}) yield history, None send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output]) msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output]) return demo if __name__ == "__main__": print("🚀 Launching app...") agent = init_agent() demo = create_ui(agent) demo.queue(api_open=False).launch( server_name="0.0.0.0", server_port=7860, show_error=True, allowed_paths=[report_dir], share=False )