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, Dict, Optional, Generator
+from typing import List
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
@@ -16,22 +16,12 @@ 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__)
 
-# Constants
-MAX_TOKENS = 1800
-BATCH_SIZE = 2
-MAX_WORKERS = 4
-CHUNK_SIZE = 10  # For PDF processing
-
-# Persistent directory setup
+# Persistent directory
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 
@@ -44,13 +34,11 @@ 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,
-    "TRANSFORMERS_CACHE": model_cache_dir,
-    "VLLM_CACHE_DIR": vllm_cache_dir,
-    "TOKENIZERS_PARALLELISM": "false",
-    "CUDA_LAUNCH_BLOCKING": "1"
-})
+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"))
@@ -61,294 +49,174 @@ 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:
-        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 ""
+        return hashlib.md5(f.read()).hexdigest()
 
 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 ""
 
-        # Process in chunks with memory cleanup
-        results = []
-        for chunk_start in range(0, total_pages, CHUNK_SIZE):
-            chunk_end = min(chunk_start + CHUNK_SIZE, total_pages)
-            
+        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:
-                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)
-            
-            # Explicit cleanup
-            del pdf
-            gc.collect()
-        
-        return "\n\n".join(filter(None, results))
+                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}")
+        logger.error("PDF processing error: %s", e)
         return f"PDF processing error: {str(e)}"
 
-def excel_to_json(file_path: str) -> List[Dict]:
-    """Optimized Excel processing with chunking"""
-    try:
-        # Try fastest engines first
-        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"""
+def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str:
     try:
-        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)}"}]
+        file_h = file_hash(file_path)
+        cache_key = f"{file_h}_{file_type}"
+        if cache_key in cache:
+            return cache[cache_key]
 
-@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)
-            return [{
-                "filename": os.path.basename(file_path),
-                "content": text,
-                "status": "initial",
-                "type": "pdf"
-            }]
-        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)
+            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:
-            return [{"error": f"Unsupported file type: {file_type}"}]
+            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 {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)
-    ]
+        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=""):
-    """Optimized system monitoring"""
     try:
-        cpu = psutil.cpu_percent(interval=0.5)
+        cpu = psutil.cpu_percent(interval=1)
         mem = psutil.virtual_memory()
-        logger.info(f"[{tag}] CPU: {cpu:.1f}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB")
-        
-        # GPU monitoring with timeout
-        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")
+        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(f"[{tag}] Monitor failed: {e}")
+        logger.error("[%s] GPU/CPU monitor failed: %s", tag, e)
 
 def clean_response(text: str) -> str:
-    """Enhanced response cleaning with aggressive deduplication"""
-    if not text:
-        return ""
-    
-    # Pre-compiled regex patterns for cleaning
-    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"\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)
-    
-    # Deduplicate near-identical sentences using similarity threshold
-    sentences = text.split(". ")
-    unique_sentences = []
-    seen = set()
-    
-    for s in sentences:
-        if not s:
-            continue
-        # Check similarity with existing sentences
-        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:
-    """Enhanced findings summarization for a single, concise paragraph"""
-    if not combined_response:
-        return "No missed diagnoses were identified in the provided records."
-    
-    # Pre-compiled regex patterns
-    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)
-    
+    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 = []
-    current_section = None
-    
-    for line in combined_response.splitlines():
+    lines = text.splitlines()
+    in_diagnoses_section = False
+    for line in lines:
         line = line.strip()
         if not line:
             continue
-        
-        # Check section headers
-        section_match = section_pattern.match(line)
-        if section_match:
-            current_section = "diagnoses" if section_match.group(1) == "Missed Diagnoses" else None
+        if re.match(r"###\s*Missed Diagnoses", line):
+            in_diagnoses_section = True
             continue
-        
-        # Process diagnosis lines in the correct section
-        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)
-    
-    # Extract findings from non-sectioned text
-    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):
+        if re.match(r"###\s*(Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
+            in_diagnoses_section = False
             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."
-    
+        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))]
     
-    # Create a single paragraph
-    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."
+    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
+    return summary.strip()
 
-@lru_cache(maxsize=1)
 def init_agent():
-    """Cached agent initialization with memory optimization"""
     logger.info("Initializing model...")
     log_system_usage("Before Load")
-    
-    # Tool setup
     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)
 
-    # Initialize with optimized settings
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -360,192 +228,122 @@ def init_agent():
         additional_default_tools=[],
     )
     agent.init_model()
-    
     log_system_usage("After Load")
     logger.info("Agent Ready")
     return agent
 
 def create_ui(agent):
-    """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 concise, evidence-based summary in ONE paragraph without headings, bullet points, or repeating non-clinical data (e.g., name, date of birth, allergies). 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 use external tools, retrieve additional data, or summarize non-clinical information. If no clinical findings are present, state 'No missed diagnoses identified' in ONE sentence. Ignore other oversight categories (e.g., medication conflicts).
+    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}
 """
 
-    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="Detailed Analysis", 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="Summary of Missed Diagnoses")
-                download_output = gr.File(label="Download Full Report")
-                progress_bar = gr.Progress()
-
         def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
-            """Optimized analysis pipeline with memory management"""
             history.append({"role": "user", "content": message})
             yield history, None, ""
 
-            # Process files with caching
-            extracted = []
+            extracted = ""
             file_hash_value = ""
-            
             if files:
-                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[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, ""
+                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 ""
 
-            # Convert to text with memory efficiency
-            text_content = "\n".join(json.dumps(item, ensure_ascii=False) for item in extracted)
-            del extracted
-            gc.collect()
+                history.append({"role": "assistant", "content": "✅ Text extraction complete."})
+                yield history, None, ""
 
-            # Tokenize and chunk
-            chunks = tokenize_and_chunk(text_content)
-            del text_content
-            gc.collect()
-            
+            chunk_size = 6000
+            chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
             combined_response = ""
-            report_path = None
-            seen_responses = set()  # Track unique responses to avoid repetition
-            
+            batch_size = 2
+
             try:
-                # Process in optimized batches
-                for batch_idx in range(0, len(chunks), BATCH_SIZE):
-                    batch_chunks = chunks[batch_idx:batch_idx + BATCH_SIZE]
-                    
-                    # Prepare prompts
-                    batch_prompts = [
-                        PROMPT_TEMPLATE.format(
-                            batch_idx + i + 1,
-                            len(chunks),
-                            chunk=chunk[:1800]  # Conservative size
-                        )
-                        for i, chunk in enumerate(batch_chunks)
-                    ]
-                    
-                    progress(batch_idx / len(chunks), 
-                           desc=f"Analyzing batch {(batch_idx // BATCH_SIZE) + 1}/{(len(chunks) + BATCH_SIZE - 1) // BATCH_SIZE}")
-                    
-                    # Process batch
-                    with ThreadPoolExecutor(max_workers=min(BATCH_SIZE, MAX_WORKERS)) as executor:
-                        futures = {
-                            executor.submit(
-                                agent.run_gradio_chat,
-                                prompt, [], 0.2, 512, 2048, False, []
-                            ): idx
-                            for idx, prompt in enumerate(batch_prompts)
-                        }
-                        
+                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_idx = futures[future]
                             chunk_response = ""
-                            
-                            try:
-                                for chunk_output in future.result():
-                                    if isinstance(chunk_output, (list, str)):
-                                        content = ""
-                                        if isinstance(chunk_output, list):
-                                            content = " ".join(
-                                                clean_response(m.content)
-                                                for m in chunk_output
-                                                if hasattr(m, 'content') and m.content
-                                            )
-                                        elif isinstance(chunk_output, str):
-                                            content = clean_response(chunk_output)
-                                        
-                                        if content and content != "No missed diagnoses identified.":
-                                            # Check for near-duplicate responses
-                                            is_unique = True
-                                            for seen_response in seen_responses:
-                                                if SequenceMatcher(None, content.lower(), seen_response.lower()).ratio() > 0.9:
-                                                    is_unique = False
-                                                    break
-                                            if is_unique:
-                                                chunk_response += content + " "
-                                                seen_responses.add(content)
-                                
-                                if chunk_response:
-                                    combined_response += f"--- Analysis for Chunk {batch_idx + chunk_idx + 1} ---\n{chunk_response.strip()}\n"
-                                    history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                                    yield history, None, ""
-                            finally:
-                                del future
-                                torch.cuda.empty_cache()
-                                gc.collect()
-
-                # Generate final outputs
+                            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)
-                
-                if file_hash_value:
-                    report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                    try:
-                        with open(report_path, "w", encoding="utf-8") as f:
-                            f.write(combined_response + "\n\n" + summary)
-                    except Exception as e:
-                        logger.error(f"Report save failed: {e}")
-                        report_path = None
-                
-                yield history, report_path, summary
+                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(f"Analysis error: {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)}"
-            finally:
-                torch.cuda.empty_cache()
-                gc.collect()
-
-        # Event handlers
-        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]
-        )
-    
+
+        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 optimized app...")
+        logger.info("Launching app...")
         agent = init_agent()
         demo = create_ui(agent)
-        demo.queue(
-            api_open=False,
-            max_size=20
-        ).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()