Update app.py

app.py

@@ -1,404 +1,209 @@
-import sys
-import os
-import pandas as pd
-import pdfplumber
-import json
-import gradio as gr
-from typing import List, Tuple, Optional
+import os, sys, json, re, gc, time, hashlib, logging, shutil, subprocess, multiprocessing as mp
+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 fitz                                     # ⇐ PyMuPDF
 import pyarrow as pa
-import pyarrow.parquet as pq
 import pyarrow.csv as pc
-import numpy as np
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-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")
+import pyarrow.dataset as ds
+import pandas as pd
+import torch, gradio as gr, psutil, numpy as np
+from diskcache import Cache
 
-def file_hash(path: str) -> str:
+# ──────────────────────────────  CONSTANTS  ────────────────────────────── #
+PERSIST = "/data/hf_cache"
+MODEL_CACHE   = os.path.join(PERSIST, "txagent_models")
+TOOL_CACHE    = os.path.join(PERSIST, "tool_cache")
+FILE_CACHE    = os.path.join(PERSIST, "preprocessed")
+REPORT_DIR    = os.path.join(PERSIST, "reports")
+VLLM_CACHEDIR = os.path.join(PERSIST, "vllm_cache")
+
+for d in (MODEL_CACHE, TOOL_CACHE, FILE_CACHE, REPORT_DIR, VLLM_CACHEDIR):
+    os.makedirs(d, exist_ok=True)
+
+os.environ.update(
+    HF_HOME               = MODEL_CACHE,
+    TRANSFORMERS_CACHE    = MODEL_CACHE,
+    VLLM_CACHE_DIR        = VLLM_CACHEDIR,
+    TOKENIZERS_PARALLELISM= "false",
+    CUDA_LAUNCH_BLOCKING  = "1",
+)
+
+# put local `src/` first
+ROOT = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(ROOT, "src"))
+from txagent.txagent import TxAgent                        # noqa: E402
+
+# ───────────────────────────────  LOGGING  ─────────────────────────────── #
+logging.basicConfig(format="%(asctime)s — %(levelname)s — %(message)s",
+                    level=logging.INFO)
+log = logging.getLogger("app")
+
+# ────────────────────────────────  CACHE  ──────────────────────────────── #
+cache = Cache(FILE_CACHE, size_limit=20 * 1024 ** 3)       # 20 GB
+
+# ───────────────────────────────  HELPERS  ─────────────────────────────── #
+def md5(path: str) -> str:
+    h = hashlib.md5()
     with open(path, "rb") as f:
-        return hashlib.md5(f.read()).hexdigest()
-
-def extract_all_pages(file_path: str, progress_callback=None) -> str:
-    try:
-        with pdfplumber.open(file_path) as pdf:
-            total_pages = len(pdf.pages)
-            if total_pages == 0:
-                return ""
-
-        batch_size = 10
-        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 = []
-            with pdfplumber.open(file_path) as pdf:
-                for page in pdf.pages[start:end]:
-                    page_num = start + pdf.pages.index(page)
-                    page_text = page.extract_text() or ""
-                    results.append((page_num, f"=== Page {page_num + 1} ===\n{page_text.strip()}"))
-            return results
-
-        with ThreadPoolExecutor(max_workers=6) as executor:
-            futures = [executor.submit(extract_batch, start, end) for start, end in batches]
-            for future in as_completed(futures):
-                for page_num, text in future.result():
-                    text_chunks[page_num] = text
-                processed_pages += batch_size
-                if progress_callback:
-                    progress_callback(min(processed_pages, total_pages), total_pages)
-
-        return "\n\n".join(filter(None, text_chunks))
-    except Exception as e:
-        logger.error("PDF processing error: %s", e)
-        return f"PDF processing error: {str(e)}"
-
-def excel_to_arrow(file_path: str) -> pa.Table:
-    """Convert Excel file to Arrow table for faster processing"""
-    try:
-        # First try with openpyxl (faster for xlsx)
-        try:
-            df = pd.read_excel(file_path, engine='openpyxl', header=None, dtype=str)
-        except Exception:
-            # Fall back to xlrd if needed
-            df = pd.read_excel(file_path, engine='xlrd', header=None, dtype=str)
-        
-        # Convert to Arrow table
-        table = pa.Table.from_pandas(df.fillna(""))
-        return table
-    except Exception as e:
-        logger.error(f"Error converting Excel to Arrow: {e}")
-        raise
-
-def csv_to_arrow(file_path: str) -> pa.Table:
-    """Convert CSV file to Arrow table for faster processing"""
-    try:
-        read_options = pc.ReadOptions(
-            encoding='utf-8',
-            invalid_row_handler=lambda x: None,
-            column_names=[str(i) for i in range(1000)]  # Generous column count
-        )
-        convert_options = pc.ConvertOptions(
-            strings_can_be_null=True,
-            quoted_strings_can_be_null=True,
-            include_columns=None
-        )
-        table = pc.read_csv(
-            file_path,
-            read_options=read_options,
-            convert_options=convert_options
-        )
-        return table
-    except Exception as e:
-        logger.error(f"Error converting CSV to Arrow: {e}")
-        raise
-
-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:
-            return cache[cache_key]
-
-        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 in ["csv", "xls", "xlsx"]:
-            # Use Arrow for tabular data processing
-            start_time = time.time()
-            
-            if file_type == "csv":
-                table = csv_to_arrow(file_path)
-            else:  # Excel files
-                table = excel_to_arrow(file_path)
-            
-            # Convert to list of lists efficiently
-            content = []
-            for col in table.columns:
-                content.append([str(x) if x is not None else "" for x in col.to_pylist()])
-            
-            # Transpose to get rows
-            rows = list(map(list, zip(*content)))
-            
-            logger.info(f"Processed {len(rows)} rows in {time.time()-start_time:.2f}s")
-            result = json.dumps({
-                "filename": os.path.basename(file_path),
-                "rows": rows,
-                "arrow_processed": True  # Flag for optimized processing
-            })
-        else:
-            result = json.dumps({"error": f"Unsupported file type: {file_type}"})
-
-        cache[cache_key] = 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)
-    # Remove unwanted patterns and tool call artifacts
-    text = re.sub(r"\[.*?\]|\bNone\b|To analyze the patient record excerpt.*?medications\.|Since the previous attempts.*?\.|I need to.*?medications\.|Retrieving tools.*?\.", "", text, flags=re.DOTALL)
-    # Extract only missed diagnoses, ignoring other categories
-    diagnoses = []
-    lines = text.splitlines()
-    in_diagnoses_section = False
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-        if re.match(r"###\s*Missed Diagnoses", line):
-            in_diagnoses_section = True
-            continue
-        if re.match(r"###\s*(Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
-            in_diagnoses_section = False
-            continue
-        if in_diagnoses_section and re.match(r"-\s*.+", line):
-            diagnosis = re.sub(r"^\-\s*", "", line).strip()
-            if diagnosis and not re.match(r"No issues identified", diagnosis, re.IGNORECASE):
-                diagnoses.append(diagnosis)
-    # Join diagnoses into a plain text paragraph
-    text = " ".join(diagnoses)
-    # Clean up extra whitespace and punctuation
-    text = re.sub(r"\s+", " ", text).strip()
-    text = re.sub(r"[^\w\s\.\,\(\)\-]", "", text)
-    return text if text else ""
-
-def summarize_findings(combined_response: str) -> str:
-    # Split response by chunk analyses
-    chunks = combined_response.split("--- Analysis for Chunk")
-    diagnoses = []
-    for chunk in chunks:
-        chunk = chunk.strip()
-        if not chunk or "No oversights identified" in chunk:
-            continue
-        # Extract missed diagnoses from chunk
-        lines = chunk.splitlines()
-        in_diagnoses_section = False
-        for line in lines:
-            line = line.strip()
-            if not line:
-                continue
-            if re.match(r"###\s*Missed Diagnoses", line):
-                in_diagnoses_section = True
-                continue
-            if re.match(r"###\s*(Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
-                in_diagnoses_section = False
-                continue
-            if in_diagnoses_section and re.match(r"-\s*.+", line):
-                diagnosis = re.sub(r"^\-\s*", "", line).strip()
-                if diagnosis and not re.match(r"No issues identified", diagnosis, re.IGNORECASE):
-                    diagnoses.append(diagnosis)
-
-    # Remove duplicates while preserving order
-    seen = set()
-    unique_diagnoses = [d for d in diagnoses if not (d in seen or seen.add(d))]
-    
-    if not unique_diagnoses:
-        return "No missed diagnoses were identified in the provided records."
-
-    # Combine into a single paragraph
-    summary = "Missed diagnoses include " + ", ".join(unique_diagnoses[:-1])
-    if len(unique_diagnoses) > 1:
-        summary += f", and {unique_diagnoses[-1]}"
-    elif len(unique_diagnoses) == 1:
-        summary = "Missed diagnoses include " + unique_diagnoses[0]
-    summary += ", all of which require urgent clinical review to prevent potential adverse outcomes."
-    
-    return summary.strip()
-
-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)
-
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+# ——— PDF ——— #
+def _extract_pg(args):
+    path, pg_no = args
+    with fitz.open(path) as doc:
+        page = doc.load_page(pg_no)
+        text = page.get_text("text")
+    return pg_no, f"=== Page {pg_no+1} ===\n{text.strip()}"
+
+def pdf_to_txt(path: str, progress=None) -> str:
+    doc = fitz.open(path)
+    total = doc.page_count
+    with mp.Pool() as pool:
+        for pg_no, txt in pool.imap_unordered(_extract_pg, [(path, i) for i in range(total)]):
+            if progress: progress(pg_no+1, total)
+            cache.set((path, "pg", pg_no, os.path.getmtime(path)), txt)
+    pages = [cache[(path, "pg", i, os.path.getmtime(path))] for i in range(total)]
+    return "\n\n".join(pages)
+
+# ——— CSV/XLSX ——— #
+def csv_to_arrow(path: str) -> pa.Table:
+    return pc.read_csv(path, read_options=pc.ReadOptions(block_size=1 << 24))  # 16 MiB
+
+def excel_to_arrow(path: str) -> pa.Table:
+    # openpyxl is C‑based; fallback to xlrd only for .xls
+    df = pd.read_excel(path, engine="openpyxl" if path.endswith("x") else "xlrd", dtype=str)
+    return pa.Table.from_pandas(df.fillna(""))
+
+def table_to_rows(tbl: pa.Table) -> List[List[str]]:
+    cols = [col.to_pylist() for col in tbl.columns]
+    return [list(r) for r in zip(*cols)]
+
+def load_tabular(path: str) -> List[List[str]]:
+    key = (path, os.path.getmtime(path))
+    if key in cache:
+        return cache[key]
+    tbl = csv_to_arrow(path) if path.endswith("csv") else excel_to_arrow(path)
+    rows = table_to_rows(tbl)
+    cache[key] = rows
+    return rows
+
+# ——— CLEANERS ——— #
+def strip_tool_noise(txt: str) -> str:
+    txt = re.sub(r"\[.*?TOOL.*?]", "", txt, flags=re.S)
+    txt = re.sub(r"\s+", " ", txt).strip()
+    return txt
+
+def summarize(findings: List[str]) -> str:
+    uniq = list(dict.fromkeys(findings))      # preserve order, dedupe
+    if not uniq:
+        return "No missed diagnoses identified."
+    if len(uniq) == 1:
+        return f"Missed diagnosis: {uniq[0]}."
+    return ("Missed diagnoses include " +
+            ", ".join(uniq[:-1]) +
+            f", and {uniq[-1]}. Please review urgently.")
+
+# ——— MONITOR ——— #
+def sys_usage(tag=""):
+    cpu = psutil.cpu_percent()
+    mem = psutil.virtual_memory()
+    log.info("[%s] CPU %.1f%% — RAM %.0f/%.0f GB",
+             tag, cpu, mem.used/1e9, mem.total/1e9)
+
+# ───────────────────────────────  AGENT  ──────────────────────────────── #
+def init_agent() -> TxAgent:
+    sys_usage("before‑load")
     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,
-        step_rag_num=4,
-        seed=100,
-        additional_default_tools=[],
+        model_name       ="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+        rag_model_name   ="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+        step_rag_num     =4,
+        force_finish     =True,
+        enable_checker   =False,
+        seed             =42
     )
     agent.init_model()
-    log_system_usage("After Load")
-    logger.info("Agent Ready")
+    sys_usage("after‑load")
     return agent
 
-def create_ui(agent):
+AGENT = init_agent()
+
+# ───────────────────────────────  GRADIO  ─────────────────────────────── #
+prompt_tpl = (
+"Analyze the following excerpt (chunk {idx}/{tot}) and list **only** missed diagnoses "
+"with clinical finding + implication in one sentence each.\n\n{chunk}"
+)
+
+def analyze(user_msg, chat_hist, files, prog=gr.Progress()):
+    chat_hist.append({"role":"user", "content":user_msg})
+    yield chat_hist, None, ""
+
+    # ——— ingest files ——— #
+    extracted = ""
+    if files:
+        for f in files:
+            ext = f.name.lower().split(".")[-1]
+            if ext == "pdf":
+                txt = pdf_to_txt(f.name,
+                        progress=lambda cur, tot: prog(cur/tot, desc=f"PDF {cur}/{tot}"))
+                extracted += txt + "\n"
+            elif ext in ("csv", "xls", "xlsx"):
+                rows = load_tabular(f.name)
+                extracted += "\n".join(",".join(r) for r in rows) + "\n"
+        chat_hist.append({"role":"assistant", "content":"✅ Files parsed"})
+        yield chat_hist, None, ""
+
+    # ——— chunk & batch ——— #
+    max_tokens = 6000
+    chunks = [extracted[i:i+max_tokens] for i in range(0, len(extracted), max_tokens)]
+    findings = []
+    for i in range(0, len(chunks), 4):                     # batch of 4
+        batch = chunks[i:i+4]
+        prompts = [prompt_tpl.format(idx=i+j+1, tot=len(chunks), chunk=c[:4000])
+                   for j,c in enumerate(batch)]
+
+        with torch.inference_mode():
+            outs = [list(AGENT.run_gradio_chat(p, [], 0.2, 512, 2048, False, []))[-1]
+                    for p in prompts]
+
+        for out in outs:
+            if out and hasattr(out, "content"):
+                clean = strip_tool_noise(out.content)
+                if clean and "No missed" not in clean:
+                    findings.append(clean)
+
+        prog((i+len(batch))/len(chunks), desc=f"LLM {i+len(batch)}/{len(chunks)}")
+
+    summary = summarize(findings)
+    chat_hist.append({"role":"assistant", "content":summary})
+
+    # save full
+    if files:
+        fn_hash = md5(files[0].name)
+        p = os.path.join(REPORT_DIR, f"{fn_hash}_report.txt")
+        with open(p, "w") as w:
+            w.write("\n".join(findings) + "\n\n" + summary)
+        yield chat_hist, p, summary
+    else:
+        yield chat_hist, None, summary
+
+def ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
-        chatbot = gr.Chatbot(label="Detailed Analysis", height=600, type="messages")
-        final_summary = gr.Markdown(label="Summary of Missed Diagnoses")
-        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 missed diagnoses only. Provide a concise, evidence-based summary as a single paragraph without headings or bullet points. Include specific clinical findings (e.g., 'elevated blood pressure (160/95) on page 10'), their potential implications (e.g., 'may indicate untreated hypertension'), and a recommendation for urgent review. Do not include other oversight categories like medication conflicts. If no missed diagnoses are found, state 'No missed diagnoses identified' in a single sentence.
-Patient Record Excerpt (Chunk {0} of {1}):
-{chunk}
-"""
-
-        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, ""
-
-                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.append({"role": "assistant", "content": "✅ Text extraction complete."})
-                yield history, None, ""
-
-            chunk_size = 6000
-            chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
-            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[:4000]) 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)}")
-
-                    with ThreadPoolExecutor(max_workers=len(batch_chunks)) as executor:
-                        futures = [executor.submit(agent.run_gradio_chat, prompt, [], 0.2, 512, 2048, False, []) for prompt in batch_prompts]
-                        for future in as_completed(futures):
-                            chunk_response = ""
-                            for chunk_output in future.result():
-                                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 + " "
-                                elif isinstance(chunk_output, str) and chunk_output.strip():
-                                    cleaned = clean_response(chunk_output)
-                                    if cleaned:
-                                        chunk_response += cleaned + " "
-                            batch_responses.append(chunk_response.strip())
-                            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 missed diagnoses identified.\n"
-                        history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                        yield history, None, ""
-
-                if combined_response.strip() and not all("No missed diagnoses 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 missed diagnoses 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"Error 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])
+        gr.Markdown("<h1 style='text-align:center'>🩺 Clinical Oversight Assistant</h1>")
+        chat   = gr.Chatbot(height=600, label="Detailed analysis", type="messages")
+        summ   = gr.Markdown(label="Summary of missed diagnoses")
+        files  = gr.File(file_types=[".pdf",".csv",".xls",".xlsx"], file_count="multiple")
+        txtbox = gr.Textbox(placeholder="Ask about potential oversights…", show_label=False)
+        run    = gr.Button("Analyze", variant="primary")
+        dl     = gr.File(label="Download full report")
+
+        run.click(analyze, [txtbox, gr.State([]), files], [chat, dl, summ])
+        txtbox.submit(analyze, [txtbox, gr.State([]), files], [chat, dl, summ])
     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()
+    ui().queue(api_open=False).launch(
+        server_name="0.0.0.0", server_port=7860,
+        allowed_paths=[REPORT_DIR], show_error=True, share=False)