Update app.py

app.py

@@ -1,21 +1,15 @@
 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
+from datetime import datetime
 import time
 
-# Persistent directory
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
@@ -23,16 +17,12 @@ 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]:
+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
-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"))
@@ -40,146 +30,56 @@ 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 clean_response(text: str) -> str:
+    text = text.encode("utf-8", "ignore").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 parse_excel_to_prompts(file_path: str) -> List[str]:
+    xl = pd.ExcelFile(file_path)
+    df = xl.parse(xl.sheet_names[0], header=0).fillna("")
+    groups = df.groupby("Booking Number")
+    prompts = []
+    for booking, group in groups:
+        records = []
+        for _, row in group.iterrows():
+            records.append(f"- {row['Form Name']}: {row['Form Item']} = {row['Item Response']} ({row['Interview Date']} by {row['Interviewer']})\n{row['Description']}")
+        record_text = "\n".join(records)
+        prompt = f"""
+Patient Booking Number: {booking}
+
+Instructions:
+Analyze the following patient case for missed diagnoses, medication conflicts, incomplete assessments, and any urgent follow-up needed. Summarize under the markdown headings.
+
+Data:
+{record_text}
 
-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()
+### Missed Diagnoses
+- ...
 
-        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)}"})
+### Medication Conflicts
+- ...
 
-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}")
+### Incomplete Assessments
+- ...
 
-def clean_response(text: str) -> str:
-    """Clean TxAgent response to keep only markdown sections with valid findings."""
-    text = sanitize_utf8(text)
-    # Remove tool call artifacts, None, and reasoning
-    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)
-    # Remove extra whitespace and non-markdown content
-    text = re.sub(r"\n{3,}", "\n\n", text)
-    text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)  # Keep markdown-relevant characters
-    
-    # Extract markdown sections with valid findings
-    sections = []
-    current_section = None
-    lines = text.splitlines()
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-        if re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
-            current_section = line
-            sections.append([current_section])
-        elif current_section and re.match(r"-\s*.+", line) and not re.match(r"-\s*No issues identified", line):
-            sections[-1].append(line)
-    
-    # Combine only non-empty sections
-    cleaned = []
-    for section in sections:
-        if len(section) > 1:  # Section has at least one finding
-            cleaned.append("\n".join(section))
-    
-    text = "\n\n".join(cleaned).strip()
-    if not text:
-        text = ""  # Return empty string if no valid findings
-    return text
+### Urgent Follow-up
+- ...
+"""
+        prompts.append(prompt)
+    return prompts
 
 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",
@@ -191,160 +91,59 @@ def init_agent():
         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("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
+        gr.Markdown("<h1 style='text-align: center;'>\ud83e\uddfa Clinical Oversight Assistant (Excel Optimized)</h1>")
         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)
+        file_upload = gr.File(file_types=[".xlsx"], file_count="single")
+        msg_input = gr.Textbox(placeholder="Ask about patient history...", 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):
+        def analyze(message: str, history: List[dict], file) -> tuple:
             history.append({"role": "user", "content": message})
-            history.append({"role": "assistant", "content": "⏳ Extracting text from files..."})
+            history.append({"role": "assistant", "content": "⏳ Processing Excel data..."})
             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 and 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).
-Output ONLY the markdown-formatted findings, with bullet points under each heading. Do NOT include reasoning, tool calls, or intermediate steps. If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text and avoids generic responses.
-
-Example Output:
-### Missed Diagnoses
-- Elevated BP noted without diagnosis. Missed due to inconsistent visits. Risks: stroke. Recommend: BP monitoring, antihypertensives.
-### Medication Conflicts
-- No issues identified.
-### Incomplete Assessments
-- Chest pain not evaluated. Time constraints likely cause. Risks: cardiac issues. Recommend: ECG, stress test.
-### Urgent Follow-up
-- Abnormal creatinine not addressed. Delayed lab review. Risks: renal failure. Recommend: nephrology referral.
-
-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)}"})
+            prompts = parse_excel_to_prompts(file.name)
+            full_output = ""
+
+            for idx, prompt in enumerate(prompts, 1):
+                chunk_output = ""
+                for result in agent.run_gradio_chat(
+                    message=prompt,
+                    history=[],
+                    temperature=0.2,
+                    max_new_tokens=1024,
+                    max_token=4096,
+                    call_agent=False,
+                    conversation=[],
+                ):
+                    if isinstance(result, list):
+                        for r in result:
+                            if hasattr(r, 'content') and r.content:
+                                chunk_output += clean_response(r.content) + "\n"
+                    elif isinstance(result, str):
+                        chunk_output += clean_response(result) + "\n"
+                if chunk_output:
+                    output = f"--- Booking {idx} ---\n{chunk_output.strip()}\n"
+                    history.append({"role": "assistant", "content": output})
+                    full_output += output + "\n"
                     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 and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
-                                        chunk_response += cleaned + "\n\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 and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
-                                chunk_response += cleaned + "\n\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
-                    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"
-
-                # Finalize UI with complete response
-                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."})
-
-                # 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
+            file_hash_value = file_hash(file.name)
+            report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
+            with open(report_path, "w", encoding="utf-8") as f:
+                f.write(full_output)
+            yield history, report_path if os.path.exists(report_path) else 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(