Update app.py

app.py

@@ -2,34 +2,26 @@ import sys
 import os
 import pandas as pd
 import pdfplumber
+import json
 import gradio as gr
-from typing import List, Dict
+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
-from transformers import AutoTokenizer
-from functools import lru_cache
-from difflib import SequenceMatcher
+import time
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Constants
-MAX_TOKENS = 1800
-BATCH_SIZE = 1
-MAX_WORKERS = 2
-CHUNK_SIZE = 5
-MODEL_MAX_TOKENS = 131072
-MAX_TEXT_LENGTH = 500000
-MAX_ROWS_TO_PROCESS = 1000  # Limit for Excel/CSV rows
-
-# Persistent directory setup
+# Persistent directory
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
@@ -37,12 +29,16 @@ 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")
-os.makedirs(report_dir, exist_ok=True)
+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.update({
-    "HF_HOME": model_cache_dir,
-    "TOKENIZERS_PARALLELISM": "false",
-})
+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"))
@@ -50,291 +46,295 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Initialize cache
+# Initialize cache with 10GB limit
 cache = Cache(file_cache_dir, size_limit=10 * 1024**3)
 
-@lru_cache(maxsize=1)
-def get_tokenizer():
-    return AutoTokenizer.from_pretrained("mims-harvard/TxAgent-T1-Llama-3.1-8B")
-
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
 
 def file_hash(path: str) -> str:
-    hash_md5 = hashlib.md5()
     with open(path, "rb") as f:
-        for chunk in iter(lambda: f.read(4096), b""):
-            hash_md5.update(chunk)
-    return hash_md5.hexdigest()
-
-def extract_pdf_page(page, tokenizer, max_tokens=MAX_TOKENS) -> List[str]:
-    try:
-        text = page.extract_text() or ""
-        text = sanitize_utf8(text)
-        if len(text) > MAX_TEXT_LENGTH // 10:
-            text = text[:MAX_TEXT_LENGTH // 10]
-        
-        tokens = tokenizer.encode(text, add_special_tokens=False)
-        if len(tokens) > max_tokens:
-            chunks = []
-            current_chunk = []
-            current_length = 0
-            for token in tokens:
-                if current_length + 1 > max_tokens:
-                    chunks.append(tokenizer.decode(current_chunk))
-                    current_chunk = [token]
-                    current_length = 1
-                else:
-                    current_chunk.append(token)
-                    current_length += 1
-            if current_chunk:
-                chunks.append(tokenizer.decode(current_chunk))
-            return chunks
-        return [text]
-    except Exception as e:
-        logger.warning(f"Error extracting page {page.page_number}: {str(e)}")
-        return []
+        return hashlib.md5(f.read()).hexdigest()
 
-def extract_all_pages(file_path: str) -> List[str]:
+def extract_all_pages(file_path: str, progress_callback=None) -> str:
     try:
-        tokenizer = get_tokenizer()
         with pdfplumber.open(file_path) as pdf:
             total_pages = len(pdf.pages)
             if total_pages == 0:
-                return ["PDF appears to be empty"]
-
-        results = []
-        for i in range(0, min(total_pages, 50)):  # Limit to first 50 pages
-            try:
-                page = pdf.pages[i]
-                chunks = extract_pdf_page(page, tokenizer)
-                for chunk in chunks:
-                    results.append(f"=== Page {i+1} ===\n{chunk}")
-            except Exception as e:
-                logger.warning(f"Error processing page {i+1}: {str(e)}")
-                continue
-        
-        return results if results else ["Could not extract text from PDF"]
+                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(f"PDF processing error: {e}")
-        return [f"PDF processing error: {str(e)}"]
-
-def excel_to_json(file_path: str) -> List[Dict]:
-    engines = ['openpyxl', 'xlrd']
-    for engine in engines:
-        try:
-            with pd.ExcelFile(file_path, engine=engine) as excel_file:
-                sheets = excel_file.sheet_names
-                if not sheets:
-                    return [{"error": "No sheets found"}]
-                
-                results = []
-                for sheet_name in sheets[:3]:  # Limit to first 3 sheets
-                    try:
-                        df = pd.read_excel(
-                            excel_file,
-                            sheet_name=sheet_name,
-                            header=None,
-                            dtype=str,
-                            na_filter=False,
-                            nrows=MAX_ROWS_TO_PROCESS  # Limit rows
-                        )
-                        if not df.empty:
-                            rows = df.head(MAX_ROWS_TO_PROCESS).values.tolist()
-                            results.append({
-                                "filename": os.path.basename(file_path),
-                                "sheet": sheet_name,
-                                "rows": rows,
-                                "type": "excel"
-                            })
-                    except Exception as e:
-                        logger.warning(f"Error processing sheet {sheet_name}: {str(e)}")
-                        continue
-                
-                return results if results else [{"error": "No readable data found"}]
-        except Exception as e:
-            logger.warning(f"Excel engine {engine} failed: {str(e)}")
-            continue
-    
-    return [{"error": "Could not process Excel file with any engine"}]
+        logger.error("PDF processing error: %s", e)
+        return f"PDF processing error: {str(e)}"
 
-def csv_to_json(file_path: str) -> List[Dict]:
+def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str:
     try:
-        df = pd.read_csv(
-            file_path,
-            header=None,
-            dtype=str,
-            encoding_errors='replace',
-            on_bad_lines='skip',
-            nrows=MAX_ROWS_TO_PROCESS  # Limit rows
-        )
-        if df.empty:
-            return [{"error": "CSV file is empty"}]
-            
-        return [{
-            "filename": os.path.basename(file_path),
-            "rows": df.values.tolist(),
-            "type": "csv"
-        }]
-    except Exception as e:
-        logger.error(f"CSV processing error: {e}")
-        return [{"error": f"CSV processing error: {str(e)}"}]
+        file_h = file_hash(file_path)
+        cache_key = f"{file_h}_{file_type}"
+        if cache_key in cache:
+            return cache[cache_key]
 
-def process_file_cached(file_path: str, file_type: str) -> List[Dict]:
-    try:
-        logger.info(f"Processing {file_type} file: {os.path.basename(file_path)}")
-        
         if file_type == "pdf":
-            chunks = extract_all_pages(file_path)
-            return [{
-                "filename": os.path.basename(file_path),
-                "content": chunk,
-                "type": "pdf"
-            } for chunk in chunks]
-            
-        elif file_type in ["xls", "xlsx"]:
-            return excel_to_json(file_path)
-            
+            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":
-            return csv_to_json(file_path)
-            
-        return [{"error": f"Unsupported file type: {file_type}"}]
+            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
+        return result
     except Exception as e:
-        logger.error(f"Error processing file: {e}")
-        return [{"error": f"Error processing file: {str(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:
-    if not text:
-        return ""
-    
-    patterns = [
-        (re.compile(r"\[.*?\]|\bNone\b", re.IGNORECASE), ""),
-        (re.compile(r"\s+"), " "),
-    ]
-    
-    for pattern, repl in patterns:
-        text = pattern.sub(repl, text)
-    
-    return text.strip()
-
-@lru_cache(maxsize=1)
+    text = sanitize_utf8(text)
+    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)
+    text = re.sub(r"\n{3,}", "\n\n", text)
+    text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)
+
+    sections = {}
+    current_section = None
+    lines = text.splitlines()
+    for line in lines:
+        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)
+            if current_section not in sections:
+                sections[current_section] = []
+            continue
+        finding_match = re.match(r"-\s*.+", line)
+        if finding_match and current_section and not re.match(r"-\s*No issues identified", line):
+            sections[current_section].append(line)
+
+    cleaned = []
+    for heading, findings in sections.items():
+        if findings:
+            cleaned.append(f"### {heading}\n" + "\n".join(findings))
+
+    text = "\n\n".join(cleaned).strip()
+    return text if text else ""
+
+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 = {}
+    lines = combined_response.splitlines()
+    current_section = None
+    for line in lines:
+        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)
+            if current_section not in sections:
+                sections[current_section] = []
+            continue
+        finding_match = re.match(r"-\s*(.+)", line)
+        if finding_match and current_section:
+            sections[current_section].append(finding_match.group(1))
+
+    summary_lines = []
+    for heading, findings in sections.items():
+        if findings:
+            summary = f"- **{heading}**: {'; '.join(findings[:2])}. Risks: {heading.lower()} may lead to adverse outcomes. Recommend: urgent review and specialist referral."
+            summary_lines.append(summary)
+
+    if not summary_lines:
+        return "### Summary of Clinical Oversights\nNo critical oversights identified."
+
+    return "### Summary of Clinical Oversights\n" + "\n".join(summary_lines)
+
 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": os.path.join(tool_cache_dir, "new_tool.json")},
+        tool_files_dict={"new_tool": target_tool_path},
         force_finish=True,
         enable_checker=False,
         step_rag_num=4,
         seed=100,
+        additional_default_tools=[],
     )
     agent.init_model()
+    log_system_usage("After Load")
     logger.info("Agent Ready")
     return agent
 
 def create_ui(agent):
-    PROMPT_TEMPLATE = """
-Analyze this patient record excerpt for missed diagnoses (limit response to 500 tokens):
+    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 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}
 """
 
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
-        
-        with gr.Row():
-            with gr.Column(scale=3):
-                chatbot = gr.Chatbot(label="Analysis", height=500, type="messages")
-                msg_input = gr.Textbox(placeholder="Ask about potential oversights...")
-                send_btn = gr.Button("Analyze", variant="primary")
-                file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="single")
-            
-            with gr.Column(scale=1):
-                final_summary = gr.Markdown("## Summary")
-                status = gr.Textbox(label="Status", interactive=False)
-
-        def analyze(message: str, history: List[Dict], file_obj) -> tuple:
+        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:
-                if not file_obj:
-                    return history, "Please upload a file first", "No file uploaded"
-
-                file_path = file_obj.name
-                file_type = os.path.splitext(file_path)[-1].lower().replace(".", "")
-                history.append({"role": "user", "content": message})
-
-                # Process file
-                processed = process_file_cached(file_path, file_type)
-                if "error" in processed[0]:
-                    history.append({"role": "assistant", "content": processed[0]["error"]})
-                    return history, processed[0]["error"], "File processing failed"
-
-                # Prepare chunks
-                chunks = []
-                for item in processed:
-                    if "content" in item:
-                        chunks.append(item["content"])
-                    elif "rows" in item:
-                        rows_text = "\n".join([", ".join(map(str, row)) for row in item["rows"][:100]])
-                        chunks.append(f"=== {item.get('sheet', 'Data')} ===\n{rows_text}")
-
-                if not chunks:
-                    history.append({"role": "assistant", "content": "No processable content found."})
-                    return history, "No processable content found", "Content extraction failed"
-
-                # Analyze each chunk
-                responses = []
-                for i, chunk in enumerate(chunks[:5]):
-                    try:
-                        prompt = PROMPT_TEMPLATE.format(chunk=chunk[:5000])
-                        response = agent.run_quick_summary(prompt, 0.2, 256, 500)
-                        cleaned = clean_response(response)
-                        if cleaned:
-                            responses.append(f"Analysis {i+1}:\n{cleaned}")
-                    except Exception as e:
-                        logger.warning(f"Error analyzing chunk {i+1}: {str(e)}")
-                        continue
-
-                if not responses:
-                    history.append({"role": "assistant", "content": "No valid analysis generated."})
-                    return history, "No valid analysis generated", "Analysis failed"
-
-                summary = "\n\n".join(responses)
-                history.append({"role": "assistant", "content": summary})
-                return history, "Analysis completed", "Success"
+                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 and re.search(r"###\s*\w+", cleaned):
+                                                chunk_response += cleaned + "\n\n"
+                                elif isinstance(chunk_output, str) and chunk_output.strip():
+                                    cleaned = clean_response(m.content)
+                                    if cleaned and re.search(r"###\s*\w+", cleaned):
+                                        chunk_response += cleaned + "\n\n"
+                            batch_responses.append(chunk_response)
+                            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."})
 
-            except Exception as e:
-                logger.error(f"Analysis error: {e}")
-                history.append({"role": "assistant", "content": f"Error: {str(e)}"})
-                return history, f"Error: {str(e)}", "Failed"
-            finally:
-                torch.cuda.empty_cache()
-                gc.collect()
-
-        send_btn.click(
-            analyze,
-            inputs=[msg_input, chatbot, file_upload],
-            outputs=[chatbot, final_summary, status]
-        )
+                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
 
-        msg_input.submit(
-            analyze,
-            inputs=[msg_input, chatbot, file_upload],
-            outputs=[chatbot, final_summary, status]
-        )
+            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.launch(
+        demo.queue(api_open=False).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"Fatal error: {e}")
-        raise
+    finally:
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()