import sys import os import pandas as pd 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 logging import torch import gc from diskcache import Cache import time import asyncio import pypdfium2 as pdfium import pytesseract from PIL import Image import io # Configure logging and suppress warnings logging.basicConfig(level=logging.INFO) logging.getLogger("pdfminer").setLevel(logging.ERROR) logger = logging.getLogger(__name__) # 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 # Initialize cache with 10GB limit cache = Cache(file_cache_dir, size_limit=10 * 1024**3) 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() async def extract_all_pages_async(file_path: str, progress_callback=None, use_ocr=False) -> str: try: pdf = pdfium.PdfDocument(file_path) total_pages = len(pdf) if total_pages == 0: return "" batch_size = 5 batches = [(i, min(i + batch_size, total_pages)) for i in range(0, total_pages, batch_size)] text_chunks = [""] * total_pages processed_pages = 0 def extract_batch(start: int, end: int) -> List[tuple]: results = [] for i in range(start, end): page = pdf[i] text = page.get_textpage().get_text_range() or "" if not text.strip() and use_ocr: # Fallback to OCR bitmap = page.render(scale=2).to_pil() text = pytesseract.image_to_string(bitmap, lang="eng") results.append((i, f"=== Page {i + 1} ===\n{text.strip()}")) return results loop = asyncio.get_event_loop() with ThreadPoolExecutor(max_workers=4) as executor: futures = [loop.run_in_executor(executor, extract_batch, start, end) for start, end in batches] for future in await asyncio.gather(*futures): for page_num, text in future: text_chunks[page_num] = text logger.debug("Page %d extracted: %s...", page_num + 1, text[:50]) processed_pages += batch_size if progress_callback: progress_callback(min(processed_pages, total_pages), total_pages) pdf.close() extracted_text = "\n\n".join(filter(None, text_chunks)) logger.info("Extracted %d pages, total length: %d chars", total_pages, len(extracted_text)) return extracted_text except Exception as e: logger.error("PDF processing error: %s", e) return f"PDF processing error: {str(e)}" def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str: try: file_h = file_hash(file_path) cache_key = f"{file_h}_{file_type}" if cache_key in cache: logger.info("Using cached extraction for %s", file_path) return cache[cache_key] if file_type == "pdf": # Try without OCR first, fallback to OCR if empty text = asyncio.run(extract_all_pages_async(file_path, progress_callback, use_ocr=False)) if not text.strip() or "PDF processing error" in text: logger.info("Retrying extraction with OCR for %s", file_path) text = asyncio.run(extract_all_pages_async(file_path, progress_callback, use_ocr=True)) 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}"}) cache[cache_key] = result logger.info("Cached extraction for %s, size: %d bytes", file_path, len(result)) return result except Exception as e: logger.error("Error processing %s: %s", os.path.basename(file_path), 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() logger.info("[%s] CPU: %.1f%% | RAM: %dMB / %dMB", tag, cpu, mem.used // (1024**2), mem.total // (1024**2)) 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(", ") logger.info("[%s] GPU: %sMB / %sMB | Utilization: %s%%", tag, used, total, util) except Exception as e: logger.error("[%s] GPU/CPU monitor failed: %s", tag, e) def clean_response(text: str) -> str: text = sanitize_utf8(text) text = text.replace("[", "").replace("]", "").replace("None", "") # Faster string ops text = text.replace("\n\n\n", "\n\n") sections = {} current_section = None for line in text.splitlines(): line = line.strip() if not line: continue section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line) if section_match: current_section = section_match.group(1) sections.setdefault(current_section, []) continue if current_section and line.startswith("- ") and "No issues identified" not in line: sections[current_section].append(line) cleaned = [f"### {heading}\n" + "\n".join(findings) for heading, findings in sections.items() if findings] result = "\n\n".join(cleaned).strip() logger.debug("Cleaned response length: %d chars", len(result)) return result or "" def summarize_findings(combined_response: str) -> str: if not combined_response or all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")): return "### Summary of Clinical Oversights\nNo critical oversights identified in the provided records." sections = {} current_section = None for line in combined_response.splitlines(): line = line.strip() if not line: continue section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line) if section_match: current_section = section_match.group(1) sections.setdefault(current_section, []) continue if current_section and line.startswith("- "): sections[current_section].append(line[2:]) summary_lines = [ f"- **{heading}**: {'; '.join(findings[:1])}. Risks: potential adverse outcomes. Recommend: urgent review." for heading, findings in sections.items() if findings ] result = "### Summary of Clinical Oversights\n" + "\n".join(summary_lines) if summary_lines else "### Summary of Clinical Oversights\nNo critical oversights identified." logger.debug("Summary length: %d chars", len(result)) return result def init_agent(): logger.info("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=False, enable_rag=False, init_rag_num=0, step_rag_num=0, seed=100, additional_default_tools=[], ) agent.init_model() log_system_usage("After Load") logger.info("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="Detailed Analysis", height=600, type="messages") final_summary = gr.Markdown(label="Summary of Clinical Oversights") 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") progress_bar = gr.Progress() prompt_template = """ Analyze the patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown with findings grouped under headings (e.g., 'Missed Diagnoses'). For each finding, include clinical context, risks, and recommendations. Output only markdown bullet points under headings. If no issues, state "No issues identified". Patient Record Excerpt (Chunk {0} of {1}): {chunk} """ async def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()): history.append({"role": "user", "content": message}) yield history, None, "" extracted = "" file_hash_value = "" if files: def update_extraction_progress(current, total): progress(current / total, desc=f"Extracting text... Page {current}/{total}") return history, None, "" futures = [convert_file_to_json(f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files] results = [sanitize_utf8(future) for future in futures] extracted = "\n".join(results) file_hash_value = file_hash(files[0].name) if files else "" history.append({"role": "assistant", "content": "✅ Text extraction complete."}) yield history, None, "" logger.info("Extracted text length: %d chars", len(extracted)) chunk_size = 4000 # Increased slightly chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)] logger.info("Created %d chunks", len(chunks)) combined_response = "" batch_size = 2 try: for batch_idx in range(0, len(chunks), batch_size): batch_chunks = chunks[batch_idx:batch_idx + batch_size] batch_prompts = [prompt_template.format(i + 1, len(chunks), chunk=chunk[:2000]) for i, chunk in enumerate(batch_chunks)] batch_responses = [] progress((batch_idx + 1) / len(chunks), desc=f"Analyzing chunks {batch_idx + 1}-{min(batch_idx + batch_size, len(chunks))}/{len(chunks)}") async def process_chunk(prompt): chunk_response = "" for chunk_output in agent.run_gradio_chat( message=prompt, history=[], temperature=0.2, max_new_tokens=128, max_token=768, 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 and re.search(r"###\s*\w+", cleaned): chunk_response += cleaned + "\n\n" elif isinstance(chunk_output, str) and chunk_output.strip(): cleaned = clean_response(chunk_output) if cleaned and re.search(r"###\s*\w+", cleaned): chunk_response += cleaned + "\n\n" logger.debug("Chunk response length: %d chars", len(chunk_response)) return chunk_response futures = [process_chunk(prompt) for prompt in batch_prompts] batch_responses = await asyncio.gather(*futures) torch.cuda.empty_cache() gc.collect() for chunk_idx, chunk_response in enumerate(batch_responses, batch_idx + 1): if chunk_response: combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n" else: combined_response += f"--- Analysis for Chunk {chunk_idx} ---\nNo oversights identified for this chunk.\n\n" history[-1] = {"role": "assistant", "content": combined_response.strip()} yield history, None, "" if combined_response.strip() and not all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")): history[-1]["content"] = combined_response.strip() else: history.append({"role": "assistant", "content": "No oversights identified in the provided records."}) summary = summarize_findings(combined_response) 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 + "\n\n" + summary) yield history, report_path if report_path and os.path.exists(report_path) else None, summary except Exception as e: logger.error("Analysis error: %s", e) history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"}) yield history, None, f"### Summary of Clinical Oversights\nError occurred during analysis: {str(e)}" send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary]) msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary]) return demo if __name__ == "__main__": try: logger.info("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 ) finally: if torch.distributed.is_initialized(): torch.distributed.destroy_process_group()