Update app.py

app.py

@@ -4,7 +4,7 @@ import pandas as pd
 import pdfplumber
 import json
 import gradio as gr
-from typing import List
+from typing import List, Dict, Optional, Generator
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
@@ -16,12 +16,22 @@ import torch
 import gc
 from diskcache import Cache
 import time
+from transformers import AutoTokenizer
+from functools import lru_cache
+import numpy as np
+from difflib import SequenceMatcher
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Persistent directory
+# Constants
+MAX_TOKENS = 1800
+BATCH_SIZE = 2
+MAX_WORKERS = 4
+CHUNK_SIZE = 10  # For PDF processing
+
+# Persistent directory setup
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
@@ -34,11 +44,13 @@ 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"
+os.environ.update({
+    "HF_HOME": model_cache_dir,
+    "TRANSFORMERS_CACHE": model_cache_dir,
+    "VLLM_CACHE_DIR": vllm_cache_dir,
+    "TOKENIZERS_PARALLELISM": "false",
+    "CUDA_LAUNCH_BLOCKING": "1"
+})
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 src_path = os.path.abspath(os.path.join(current_dir, "src"))
@@ -49,169 +61,275 @@ from txagent.txagent import TxAgent
 # Initialize cache with 10GB limit
 cache = Cache(file_cache_dir, size_limit=10 * 1024**3)
 
+# Initialize tokenizer for precise chunking (with caching)
+@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:
+    """Optimized UTF-8 sanitization"""
     return text.encode("utf-8", "ignore").decode("utf-8")
 
 def file_hash(path: str) -> str:
+    """Optimized file hashing with buffer reading"""
+    hash_md5 = hashlib.md5()
     with open(path, "rb") as f:
-        return hashlib.md5(f.read()).hexdigest()
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+def extract_pdf_page(page) -> str:
+    """Optimized single page extraction"""
+    try:
+        text = page.extract_text() or ""
+        return f"=== Page {page.page_number} ===\n{text.strip()}"
+    except Exception as e:
+        logger.warning(f"Error extracting page {page.page_number}: {str(e)}")
+        return ""
 
 def extract_all_pages(file_path: str, progress_callback=None) -> str:
+    """Optimized PDF extraction with memory management"""
     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 = []
+        results = []
+        for chunk_start in range(0, total_pages, CHUNK_SIZE):
+            chunk_end = min(chunk_start + CHUNK_SIZE, total_pages)
+            
             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))
+                with ThreadPoolExecutor(max_workers=min(CHUNK_SIZE, 4)) as executor:
+                    futures = [executor.submit(extract_pdf_page, pdf.pages[i]) 
+                             for i in range(chunk_start, chunk_end)]
+                    
+                    for future in as_completed(futures):
+                        results.append(future.result())
+                        
+                    if progress_callback:
+                        progress_callback(min(chunk_end, total_pages), total_pages)
+            
+            del pdf
+            gc.collect()
+        
+        return "\n\n".join(filter(None, results))
     except Exception as e:
-        logger.error("PDF processing error: %s", e)
+        logger.error(f"PDF processing error: {e}")
         return f"PDF processing error: {str(e)}"
 
-def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str:
+def excel_to_json(file_path: str) -> List[Dict]:
+    """Optimized Excel processing with chunking"""
+    try:
+        for engine in ['openpyxl', 'xlrd']:
+            try:
+                df = pd.read_excel(
+                    file_path, 
+                    engine=engine, 
+                    header=None, 
+                    dtype=str,
+                    na_filter=False
+                )
+                return [{
+                    "filename": os.path.basename(file_path),
+                    "rows": df.values.tolist(),
+                    "type": "excel"
+                }]
+            except Exception:
+                continue
+        raise Exception("No suitable Excel engine found")
+    except Exception as e:
+        logger.error(f"Excel processing error: {e}")
+        return [{"error": f"Excel processing error: {str(e)}"}]
+
+def csv_to_json(file_path: str) -> List[Dict]:
+    """Optimized CSV processing with chunking"""
     try:
-        file_h = file_hash(file_path)
-        cache_key = f"{file_h}_{file_type}"
-        if cache_key in cache:
-            return cache[cache_key]
+        chunks = []
+        for chunk in pd.read_csv(
+            file_path,
+            header=None,
+            dtype=str,
+            encoding_errors='replace',
+            on_bad_lines='skip',
+            chunksize=10000,
+            na_filter=False
+        ):
+            chunks.append(chunk)
+        
+        df = pd.concat(chunks) if chunks else pd.DataFrame()
+        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)}"}]
 
+@lru_cache(maxsize=100)
+def process_file_cached(file_path: str, file_type: str) -> List[Dict]:
+    """Cached file processing with memory optimization"""
+    try:
         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})
+            text = extract_all_pages(file_path)
+            return [{
+                "filename": os.path.basename(file_path),
+                "content": text,
+                "status": "initial",
+                "type": "pdf"
+            }]
         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})
+            return excel_to_json(file_path)
+        elif file_type == "csv":
+            return csv_to_json(file_path)
         else:
-            result = json.dumps({"error": f"Unsupported file type: {file_type}"})
-
-        cache[cache_key] = result
-        return result
+            return [{"error": f"Unsupported file type: {file_type}"}]
     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)}"})
+        logger.error(f"Error processing {os.path.basename(file_path)}: {e}")
+        return [{"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"}]
+
+def tokenize_and_chunk(text: str, max_tokens: int = MAX_TOKENS) -> List[str]:
+    """Optimized tokenization and chunking"""
+    tokenizer = get_tokenizer()
+    tokens = tokenizer.encode(text, add_special_tokens=False)
+    return [
+        tokenizer.decode(tokens[i:i + max_tokens])
+        for i in range(0, len(tokens), max_tokens)
+    ]
 
 def log_system_usage(tag=""):
+    """Optimized system monitoring"""
     try:
-        cpu = psutil.cpu_percent(interval=1)
+        cpu = psutil.cpu_percent(interval=0.5)
         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)
+        logger.info(f"[{tag}] CPU: {cpu:.1f}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB")
+        
+        try:
+            result = subprocess.run(
+                ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,nounits,noheader"],
+                capture_output=True, 
+                text=True,
+                timeout=2
+            )
+            if result.returncode == 0:
+                used, total, util = result.stdout.strip().split(", ")
+                logger.info(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
+        except subprocess.TimeoutExpired:
+            logger.warning(f"[{tag}] GPU monitoring timed out")
     except Exception as e:
-        logger.error("[%s] GPU/CPU monitor failed: %s", tag, e)
+        logger.error(f"[{tag}] Monitor failed: {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
+    """Enhanced response cleaning with aggressive deduplication"""
+    if not text:
+        return ""
+    
+    patterns = [
+        (re.compile(r"\[.*?\]|\bNone\b", re.IGNORECASE), ""),
+        (re.compile(r"(The patient record excerpt provides|Patient record excerpt contains).*?(John Doe|general information).*?\.", re.IGNORECASE), ""),
+        (re.compile(r"To (analyze|proceed).*?medications\.", re.IGNORECASE), ""),
+        (re.compile(r"Since the previous attempts.*?\.", re.IGNORECASE), ""),
+        (re.compile(r"I need to.*?results\.", re.IGNORECASE), ""),
+        (re.compile(r"(Therefore, )?(Retrieving|I will start by retrieving) tools.*?\.", re.IGNORECASE), ""),
+        (re.compile(r"This requires reviewing.*?\.", re.IGNORECASE), ""),
+        (re.compile(r"Given the context, it is important to review.*?\.", re.IGNORECASE), ""),
+        (re.compile(r"Final Analysis\s*", re.IGNORECASE), ""),
+        (re.compile(r"Therefore, no missed diagnoses can be identified.*?\.", re.IGNORECASE), ""),
+        (re.compile(r"\s+"), " "),
+        (re.compile(r"[^\w\s\.\,\(\)\-]"), ""),
+        (re.compile(r"(No missed diagnoses identified\.)\s*\1+", re.IGNORECASE), r"\1"),
+    ]
+    
+    for pattern, repl in patterns:
+        text = pattern.sub(repl, text)
+    
+    sentences = text.split(". ")
+    unique_sentences = []
+    seen = set()
+    
+    for s in sentences:
+        if not s:
             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 ""
+        is_unique = True
+        for seen_s in seen:
+            if SequenceMatcher(None, s.lower(), seen_s.lower()).ratio() > 0.9:
+                is_unique = False
+                break
+        if is_unique:
+            unique_sentences.append(s)
+            seen.add(s)
+    
+    text = ". ".join(unique_sentences).strip()
+    
+    return text if text else "No missed diagnoses identified."
 
 def summarize_findings(combined_response: str) -> str:
-    # Split response by chunk analyses
-    chunks = combined_response.split("--- Analysis for Chunk")
+    """Enhanced findings summarization for a single, concise paragraph"""
+    if not combined_response:
+        return "No missed diagnoses were identified in the provided records."
+    
+    diagnosis_pattern = re.compile(r"-\s*(.+)$")
+    section_pattern = re.compile(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)")
+    no_issues_pattern = re.compile(r"No issues identified|No missed diagnoses identified", re.IGNORECASE)
+    
     diagnoses = []
-    for chunk in chunks:
-        chunk = chunk.strip()
-        if not chunk or "No oversights identified" in chunk:
+    current_section = None
+    
+    for line in combined_response.splitlines():
+        line = line.strip()
+        if not line:
             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):
+        
+        section_match = section_pattern.match(line)
+        if section_match:
+            current_section = "diagnoses" if section_match.group(1) == "Missed Diagnoses" else None
+            continue
+        
+        if current_section == "diagnoses":
+            diagnosis_match = diagnosis_pattern.match(line)
+            if diagnosis_match and not no_issues_pattern.search(line):
+                diagnosis = diagnosis_match.group(1).strip()
+                if diagnosis:
                     diagnoses.append(diagnosis)
-
-    # Remove duplicates while preserving order
+    
+    medication_pattern = re.compile(r"medications includ(?:e|ing|ed) ([^\.]+)", re.IGNORECASE)
+    evaluation_pattern = re.compile(r"psychiatric evaluation.*?mention of ([^\.]+)", re.IGNORECASE)
+    
+    for line in combined_response.splitlines():
+        line = line.strip()
+        if not line or no_issues_pattern.search(line):
+            continue
+        
+        med_match = medication_pattern.search(line)
+        if med_match:
+            meds = med_match.group(1).strip()
+            diagnoses.append(f"use of medications ({meds}), suggesting an undiagnosed psychiatric condition requiring urgent review")
+        
+        eval_match = evaluation_pattern.search(line)
+        if eval_match:
+            details = eval_match.group(1).strip()
+            diagnoses.append(f"psychiatric evaluation noting {details}, indicating a potential missed psychiatric diagnosis requiring urgent review")
+    
+    if not diagnoses:
+        return "No missed diagnoses were identified in the provided records."
+    
     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."
+    summary = "The patient record indicates missed diagnoses including "
+    summary += ", ".join(unique_diagnoses[:-1])
+    summary += f", and {unique_diagnoses[-1]}" if len(unique_diagnoses) > 1 else unique_diagnoses[0]
+    summary += ". These findings suggest potential oversights in the patient's medical evaluation and require urgent clinical review to prevent adverse outcomes."
     
-    return summary.strip()
+    return summary
 
+@lru_cache(maxsize=1)
 def init_agent():
+    """Cached agent initialization with memory optimization"""
     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):
@@ -228,122 +346,188 @@ def init_agent():
         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("<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.
+    """Optimized UI creation with pre-compiled templates"""
+    PROMPT_TEMPLATE = """
+Analyze the patient record excerpt for missed diagnoses, focusing ONLY on clinical findings such as symptoms, medications, or evaluation results provided in the excerpt. Provide a detailed, evidence-based analysis using all available tools (e.g., Tool_RAG, CallAgent) to identify potential oversights. Include specific findings (e.g., 'elevated blood pressure (160/95)'), their implications (e.g., 'may indicate untreated hypertension'), and recommend urgent review. Treat medications or psychiatric evaluations as potential missed diagnoses. Do NOT repeat non-clinical information (e.g., name, date of birth, allergies). If no clinical findings are present, state 'No missed diagnoses identified' in ONE sentence. Ignore other oversight categories (e.g., medication conflicts).
 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 Summary", height=600, type="messages")
+                msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
+                send_btn = gr.Button("Analyze", variant="primary")
+                file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
+            
+            with gr.Column(scale=1):
+                final_summary = gr.Markdown(label="Missed Diagnoses Summary")
+                download_output = gr.File(label="Download Detailed Report")
+                progress_bar = gr.Progress()
+
         def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
+            """Optimized analysis pipeline with quick summary and background report"""
             history.append({"role": "user", "content": message})
             yield history, None, ""
 
-            extracted = ""
+            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."})
+                for f in files:
+                    file_type = f.name.split(".")[-1].lower()
+                    cache_key = f"{file_hash(f.name)}_{file_type}"
+                    
+                    if cache_key in cache:
+                        extracted.extend(cache[cache_key])
+                    else:
+                        result = process_file_cached(f.name, file_type)
+                        cache[cache_key] = result
+                        extracted.extend(result)
+                
+                file_hash_value = file_hash(files[0].name) if files else ""
+                history.append({"role": "assistant", "content": "✅ File processing 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
+            text_content = "\n".join(json.dumps(item, ensure_ascii=False) for item in extracted)
+            del extracted
+            gc.collect()
 
+            chunks = tokenize_and_chunk(text_content)
+            del text_content
+            gc.collect()
+            
+            combined_response = ""
+            report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
+            seen_responses = set()
+            
             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."})
+                for batch_idx in range(0, len(chunks), BATCH_SIZE):
+                    batch_chunks = chunks[batch_idx:batch_idx + BATCH_SIZE]
+                    
+                    batch_prompts = [
+                        PROMPT_TEMPLATE.format(
+                            batch_idx + i + 1,
+                            len(chunks),
+                            chunk=chunk[:1800]
+                        )
+                        for i, chunk in enumerate(batch_chunks)
+                    ]
+                    
+                    progress(batch_idx / len(chunks), 
+                           desc=f"Processing batch {(batch_idx // BATCH_SIZE) + 1}/{(len(chunks) + BATCH_SIZE - 1) // BATCH_SIZE}")
+                    
+                    with ThreadPoolExecutor(max_workers=min(BATCH_SIZE, MAX_WORKERS)) as executor:
+                        quick_futures = {
+                            executor.submit(
+                                agent.run_quick_summary,
+                                chunk, 0.2, 256, 1024
+                            ): idx
+                            for idx, chunk in enumerate(batch_chunks)
+                        }
+                        
+                        for future in as_completed(quick_futures):
+                            chunk_idx = quick_futures[future]
+                            try:
+                                quick_response = clean_response(future.result())
+                                if quick_response and quick_response != "No missed diagnoses identified.":
+                                    is_unique = True
+                                    for seen_response in seen_responses:
+                                        if SequenceMatcher(None, quick_response.lower(), seen_response.lower()).ratio() > 0.9:
+                                            is_unique = False
+                                            break
+                                    if is_unique:
+                                        combined_response += f"--- Quick Analysis for Chunk {batch_idx + chunk_idx + 1} ---\n{quick_response}\n"
+                                        seen_responses.add(quick_response)
+                                        history[-1] = {"role": "assistant", "content": combined_response.strip()}
+                                        yield history, None, ""
+                            finally:
+                                del future
+                                torch.cuda.empty_cache()
+                                gc.collect()
+
+                    # Start background detailed analysis
+                    with ThreadPoolExecutor(max_workers=min(BATCH_SIZE, MAX_WORKERS)) as executor:
+                        detailed_futures = {
+                            executor.submit(
+                                agent.run_gradio_chat,
+                                prompt, [], 0.2, 512, 2048, False, None, 3, None, 0, None, report_path
+                            ): idx
+                            for idx, prompt in enumerate(batch_prompts)
+                        }
+                        
+                        for future in as_completed(detailed_futures):
+                            chunk_idx = detailed_futures[future]
+                            try:
+                                for chunk_output in future.result():
+                                    if isinstance(chunk_output, list):
+                                        for msg in chunk_output:
+                                            if isinstance(msg, ChatMessage) and msg.content:
+                                                combined_response += clean_response(msg.content) + "\n"
+                                                history[-1] = {"role": "assistant", "content": combined_response.strip()}
+                                                yield history, report_path, ""
+                            finally:
+                                del future
+                                torch.cuda.empty_cache()
+                                gc.collect()
 
                 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
+                
+                if report_path and os.path.exists(report_path):
+                    history.append({"role": "assistant", "content": "Detailed report ready for download."})
+                    yield history, report_path, summary
+                else:
+                    history.append({"role": "assistant", "content": "Detailed report still processing."})
+                    yield history, None, summary
 
             except Exception as e:
-                logger.error("Analysis error: %s", e)
+                logger.error(f"Analysis error: {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])
+            finally:
+                torch.cuda.empty_cache()
+                gc.collect()
+
+        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...")
+        logger.info("Launching optimized app...")
         agent = init_agent()
         demo = create_ui(agent)
-        demo.queue(api_open=False).launch(
+        demo.queue(
+            api_open=False,
+            max_size=20
+        ).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()