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, Generator, Any, Optional
+from typing import List, Dict, Optional, Generator
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
@@ -15,858 +15,401 @@ import logging
 import torch
 import gc
 from diskcache import Cache
+import time
 from transformers import AutoTokenizer
-from pathlib import Path
 
-# ==================== CONFIGURATION ====================
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Directory Setup
-BASE_DIR = Path("/data/hf_cache")
-DIRECTORIES = {
-    "models": BASE_DIR / "txagent_models",
-    "tools": BASE_DIR / "tool_cache",
-    "cache": BASE_DIR / "cache",
-    "reports": BASE_DIR / "reports",
-    "vllm": BASE_DIR / "vllm_cache"
-}
-
-for dir_path in DIRECTORIES.values():
-    dir_path.mkdir(parents=True, exist_ok=True)
-
-# Environment Configuration
-os.environ.update({
-    "HF_HOME": str(DIRECTORIES["models"]),
-    "TRANSFORMERS_CACHE": str(DIRECTORIES["models"]),
-    "VLLM_CACHE_DIR": str(DIRECTORIES["vllm"]),
-    "TOKENIZERS_PARALLELISM": "false",
-    "CUDA_LAUNCH_BLOCKING": "1"
-})
-
-# Add src path for txagent
+# Persistent directory
+persistent_dir = "/data/hf_cache"
+os.makedirs(persistent_dir, exist_ok=True)
+
+model_cache_dir = os.path.join(persistent_dir, "txagent_models")
+tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
+file_cache_dir = os.path.join(persistent_dir, "cache")
+report_dir = os.path.join(persistent_dir, "reports")
+vllm_cache_dir = os.path.join(persistent_dir, "vllm_cache")
+
+for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir, vllm_cache_dir]:
+    os.makedirs(directory, exist_ok=True)
+
+os.environ["HF_HOME"] = model_cache_dir
+os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
+os.environ["VLLM_CACHE_DIR"] = vllm_cache_dir
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+
 current_dir = os.path.dirname(os.path.abspath(__file__))
 src_path = os.path.abspath(os.path.join(current_dir, "src"))
 sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# ==================== CORE COMPONENTS ====================
-class FileProcessor:
-    """Handles all file processing operations"""
-    
-    @staticmethod
-    def extract_pdf_content(file_path: str) -> str:
-        """Extract text from PDF with parallel processing"""
-        try:
+# Initialize cache with 10GB limit
+cache = Cache(file_cache_dir, size_limit=10 * 1024**3)
+
+# Initialize tokenizer for precise chunking
+tokenizer = 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:
+    with open(path, "rb") as f:
+        return hashlib.md5(f.read()).hexdigest()
+
+def extract_all_pages(file_path: str, progress_callback=None) -> str:
+    try:
+        with pdfplumber.open(file_path) as pdf:
+            total_pages = len(pdf.pages)
+            if total_pages == 0:
+                return ""
+
+        batch_size = 10
+        batches = [(i, min(i + batch_size, total_pages)) for i in range(0, total_pages, batch_size)]
+        text_chunks = [""] * total_pages
+        processed_pages = 0
+
+        def extract_batch(start: int, end: int) -> List[tuple]:
+            results = []
             with pdfplumber.open(file_path) as pdf:
-                total_pages = len(pdf.pages)
-                if not total_pages:
-                    return ""
-
-            def process_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)
-                        text = page.extract_text() or ""
-                        results.append((page_num, f"=== Page {page_num + 1} ===\n{text.strip()}"))
-                return results
-
-            batch_size = min(10, total_pages)
-            batches = [(i, min(i + batch_size, total_pages)) for i in range(0, total_pages, batch_size)]
-            text_chunks = [""] * total_pages
-
-            with ThreadPoolExecutor(max_workers=min(6, os.cpu_count() or 4)) as executor:
-                futures = [executor.submit(process_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
-
-            return "\n\n".join(filter(None, text_chunks))
-        except Exception as e:
-            logger.error(f"PDF extraction failed: {e}")
-            return f"PDF processing error: {str(e)}"
-
-    @staticmethod
-    def process_tabular_data(file_path: str, file_type: str) -> List[Dict]:
-        """Process Excel or CSV files"""
+                for page in pdf.pages[start:end]:
+                    page_num = start + pdf.pages.index(page)
+                    page_text = page.extract_text() or ""
+                    results.append((page_num, f"=== Page {page_num + 1} ===\n{page_text.strip()}"))
+            return results
+
+        with ThreadPoolExecutor(max_workers=6) as executor:
+            futures = [executor.submit(extract_batch, start, end) for start, end in batches]
+            for future in as_completed(futures):
+                for page_num, text in future.result():
+                    text_chunks[page_num] = text
+                processed_pages += batch_size
+                if progress_callback:
+                    progress_callback(min(processed_pages, total_pages), total_pages)
+
+        return "\n\n".join(filter(None, text_chunks))
+    except Exception as e:
+        logger.error("PDF processing error: %s", e)
+        return f"PDF processing error: {str(e)}"
+
+def excel_to_json(file_path: str) -> List[Dict]:
+    """Convert Excel file to JSON with optimized processing"""
+    try:
+        # First try with openpyxl (faster for xlsx)
         try:
-            if file_type == "csv":
-                chunks = pd.read_csv(
-                    file_path, 
-                    header=None,
-                    dtype=str,
-                    encoding_errors='replace',
-                    on_bad_lines='skip',
-                    chunksize=10000
-                )
-                df = pd.concat(chunks) if chunks else pd.DataFrame()
-            else:  # Excel
-                try:
-                    df = pd.read_excel(file_path, engine='openpyxl', header=None, dtype=str)
-                except:
-                    df = pd.read_excel(file_path, engine='xlrd', header=None, dtype=str)
-            
+            df = pd.read_excel(file_path, engine='openpyxl', header=None, dtype=str)
+        except Exception:
+            # Fall back to xlrd if needed
+            df = pd.read_excel(file_path, engine='xlrd', header=None, dtype=str)
+        
+        # Convert to list of lists with null handling
+        content = df.where(pd.notnull(df), "").astype(str).values.tolist()
+        
+        return [{
+            "filename": os.path.basename(file_path),
+            "rows": content,
+            "type": "excel"
+        }]
+    except Exception as e:
+        logger.error(f"Error processing Excel file: {e}")
+        return [{"error": f"Error processing Excel file: {str(e)}"}]
+
+def csv_to_json(file_path: str) -> List[Dict]:
+    """Convert CSV file to JSON with optimized processing"""
+    try:
+        # Read CSV in chunks if large
+        chunks = []
+        for chunk in pd.read_csv(
+            file_path,
+            header=None,
+            dtype=str,
+            encoding_errors='replace',
+            on_bad_lines='skip',
+            chunksize=10000
+        ):
+            chunks.append(chunk)
+        
+        df = pd.concat(chunks) if chunks else pd.DataFrame()
+        content = df.where(pd.notnull(df), "").astype(str).values.tolist()
+        
+        return [{
+            "filename": os.path.basename(file_path),
+            "rows": content,
+            "type": "csv"
+        }]
+    except Exception as e:
+        logger.error(f"Error processing CSV file: {e}")
+        return [{"error": f"Error processing CSV file: {str(e)}"}]
+
+def process_file(file_path: str, file_type: str) -> List[Dict]:
+    """Process file based on type and return JSON data"""
+    try:
+        if file_type == "pdf":
+            text = extract_all_pages(file_path)
             return [{
                 "filename": os.path.basename(file_path),
-                "rows": df.where(pd.notnull(df), "").astype(str).values.tolist(),
-                "type": file_type
+                "content": text,
+                "status": "initial",
+                "type": "pdf"
             }]
-        except Exception as e:
-            logger.error(f"{file_type.upper()} processing failed: {e}")
-            return [{"error": f"{file_type.upper()} processing error: {str(e)}"}]
-
-    @classmethod
-    def handle_upload(cls, file_path: str, file_type: str) -> List[Dict]:
-        """Route file processing based on type"""
-        processor_map = {
-            "pdf": cls.extract_pdf_content,
-            "xls": lambda x: cls.process_tabular_data(x, "excel"),
-            "xlsx": lambda x: cls.process_tabular_data(x, "excel"),
-            "csv": lambda x: cls.process_tabular_data(x, "csv")
-        }
-        
-        if file_type not in processor_map:
+        elif file_type in ["xls", "xlsx"]:
+            return excel_to_json(file_path)
+        elif file_type == "csv":
+            return csv_to_json(file_path)
+        else:
             return [{"error": f"Unsupported file type: {file_type}"}]
-        
-        try:
-            result = processor_map[file_type](file_path)
-            if file_type == "pdf":
-                return [{
-                    "filename": os.path.basename(file_path),
-                    "content": result,
-                    "type": "pdf"
-                }]
-            return result
-        except Exception as e:
-            logger.error(f"File processing failed: {e}")
-            return [{"error": f"File processing error: {str(e)}"}]
-
-class TextAnalyzer:
-    """Handles text processing and analysis"""
-    
-    def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained("mims-harvard/TxAgent-T1-Llama-3.1-8B")
-        self.cache = Cache(DIRECTORIES["cache"], size_limit=10*1024**3)
-        
-    def chunk_content(self, text: str, max_tokens: int = 1800) -> List[str]:
-        """Split text into token-limited chunks"""
-        tokens = self.tokenizer.encode(text)
-        return [
-            self.tokenizer.decode(tokens[i:i+max_tokens])
-            for i in range(0, len(tokens), max_tokens)
-        ]
-    
-    def clean_output(self, text: str) -> str:
-        """Clean and format model response"""
-        text = text.encode("utf-8", "ignore").decode("utf-8")
-        text = re.sub(
-            r"\[.*?\]|\bNone\b|To analyze the patient record excerpt.*?medications\."
-            r"|Since the previous attempts.*?\.|I need to.*?medications\."
-            r"|Retrieving tools.*?\.", "", text, flags=re.DOTALL
+    except Exception as e:
+        logger.error("Error processing %s: %s", 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 = 1800) -> List[str]:
+    """Split text into chunks based on token count"""
+    tokens = tokenizer.encode(text)
+    chunks = []
+    for i in range(0, len(tokens), max_tokens):
+        chunk_tokens = tokens[i:i + max_tokens]
+        chunks.append(tokenizer.decode(chunk_tokens))
+    return chunks
+
+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
         )
-        
-        diagnoses = []
-        in_section = False
-        
-        for line in text.splitlines():
+        if result.returncode == 0:
+            used, total, util = result.stdout.strip().split(", ")
+            logger.info("[%s] GPU: %sMB / %sMB | Utilization: %s%%", tag, used, total, util)
+    except Exception as e:
+        logger.error("[%s] GPU/CPU monitor failed: %s", tag, e)
+
+def clean_response(text: str) -> str:
+    text = sanitize_utf8(text)
+    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)
+    diagnoses = []
+    lines = text.splitlines()
+    in_diagnoses_section = False
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        if re.match(r"###\s*Missed Diagnoses", line):
+            in_diagnoses_section = True
+            continue
+        if re.match(r"###\s*(Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
+            in_diagnoses_section = False
+            continue
+        if in_diagnoses_section and re.match(r"-\s*.+", line):
+            diagnosis = re.sub(r"^\-\s*", "", line).strip()
+            if diagnosis and not re.match(r"No issues identified", diagnosis, re.IGNORECASE):
+                diagnoses.append(diagnosis)
+    text = " ".join(diagnoses)
+    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:
+    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
+        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_section = True
+                in_diagnoses_section = True
                 continue
             if re.match(r"###\s*(Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
-                in_section = False
+                in_diagnoses_section = False
                 continue
-            if in_section and re.match(r"-\s*.+", line):
+            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)
-        
-        return " ".join(diagnoses) if diagnoses else ""
-    
-    def generate_summary(self, analysis: str) -> str:
-        """Create concise clinical summary"""
-        findings = []
-        for chunk in analysis.split("--- Analysis for Chunk"):
-            chunk = chunk.strip()
-            if not chunk or "No oversights identified" in chunk:
-                continue
-                
-            in_section = False
-            for line in chunk.splitlines():
-                line = line.strip()
-                if not line:
-                    continue
-                if re.match(r"###\s*Missed Diagnoses", line):
-                    in_section = True
-                    continue
-                if re.match(r"###\s*(Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
-                    in_section = False
-                    continue
-                if in_section and re.match(r"-\s*.+", line):
-                    finding = re.sub(r"^\-\s*", "", line).strip()
-                    if finding and not re.match(r"No issues identified", finding, re.IGNORECASE):
-                        findings.append(finding)
-        
-        unique_findings = list(dict.fromkeys(findings))
-        
-        if not unique_findings:
-            return "No clinical concerns identified in the provided records."
-        
-        if len(unique_findings) > 1:
-            summary = "Potential concerns include: " + ", ".join(unique_findings[:-1])
-            summary += f", and {unique_findings[-1]}"
-        else:
-            summary = "Potential concern identified: " + unique_findings[0]
-        
-        return summary + ". Recommend urgent clinical review."
 
-class ClinicalAgent:
-    """Main application controller"""
-    
-    def __init__(self):
-        self.agent = self._init_agent()
-        self.file_processor = FileProcessor()
-        self.text_analyzer = TextAnalyzer()
-    
-    def _init_agent(self) -> Any:
-        """Initialize the AI agent"""
-        logger.info("Initializing clinical agent...")
-        self._log_system_status("pre-init")
-        
-        tool_path = DIRECTORIES["tools"] / "new_tool.json"
-        if not tool_path.exists():
-            default_tools = Path("data/new_tool.json")
-            if default_tools.exists():
-                shutil.copy(default_tools, 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": str(tool_path)},
-            force_finish=True,
-            enable_checker=False,
-            step_rag_num=4,
-            seed=100,
-            additional_default_tools=[],
-        )
-        agent.init_model()
-        
-        self._log_system_status("post-init")
-        logger.info("Clinical agent ready")
-        return agent
+    seen = set()
+    unique_diagnoses = [d for d in diagnoses if not (d in seen or seen.add(d))]
     
-    def _log_system_status(self, phase: str) -> None:
-        """Log system resource utilization"""
-        try:
-            cpu = psutil.cpu_percent(interval=1)
-            mem = psutil.virtual_memory()
-            logger.info(f"[{phase}] CPU: {cpu:.1f}% | RAM: {mem.used//(1024**2)}MB/{mem.total//(1024**2)}MB")
-            
-            gpu_info = subprocess.run(
-                ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", 
-                 "--format=csv,nounits,noheader"],
-                capture_output=True, text=True
-            )
-            if gpu_info.returncode == 0:
-                used, total, util = gpu_info.stdout.strip().split(", ")
-                logger.info(f"[{phase}] GPU: {used}MB/{total}MB | Util: {util}%")
-        except Exception as e:
-            logger.error(f"Resource monitoring failed: {e}")
-    
-    def process_stream(self, prompt: str, history: List[Dict]) -> Generator[Dict, None, None]:
-        """Stream the agent's responses"""
-        full_response = ""
-        for chunk in self.agent.run_gradio_chat(prompt, [], 0.2, 512, 2048, False, []):
-            if not chunk:
-                continue
-            
-            if isinstance(chunk, list):
-                for msg in chunk:
-                    if hasattr(msg, 'content') and msg.content:
-                        cleaned = self.text_analyzer.clean_output(msg.content)
-                        if cleaned:
-                            full_response += cleaned + " "
-                            yield {"role": "assistant", "content": full_response}
-            elif isinstance(chunk, str) and chunk.strip():
-                cleaned = self.text_analyzer.clean_output(chunk)
-                if cleaned:
-                    full_response += cleaned + " "
-                    yield {"role": "assistant", "content": full_response}
+    if not unique_diagnoses:
+        return "No missed diagnoses were identified in the provided records."
+
+    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."
     
-    def analyze_records(self, message: str, history: List[Dict], files: List) -> Generator[tuple, None, None]:
-        """Main analysis workflow"""
-        outputs = {
-            "chatbot": history.copy(),
-            "download_output": None,
-            "final_summary": "",
-            "progress": {"value": "Initializing...", "visible": True}
-        }
-        yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
-        
-        try:
-            # Add user message
+    return summary.strip()
+
+def init_agent():
+    logger.info("Initializing model...")
+    log_system_usage("Before Load")
+    default_tool_path = os.path.abspath("data/new_tool.json")
+    target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
+    if not os.path.exists(target_tool_path):
+        shutil.copy(default_tool_path, target_tool_path)
+
+    agent = TxAgent(
+        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+        tool_files_dict={"new_tool": target_tool_path},
+        force_finish=True,
+        enable_checker=False,
+        step_rag_num=4,
+        seed=100,
+        additional_default_tools=[],
+    )
+    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.
+Patient Record Excerpt (Chunk {0} of {1}):
+{chunk}
+"""
+
+        def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
             history.append({"role": "user", "content": message})
-            outputs["chatbot"] = history
-            yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
-            
-            # Process files
+            yield history, None, ""
+
             extracted = []
-            file_hash = ""
+            file_hash_value = ""
             
             if files:
+                # Process files in parallel
                 with ThreadPoolExecutor(max_workers=4) as executor:
                     futures = []
                     for f in files:
-                        file_type = Path(f.name).suffix[1:].lower()
+                        file_type = f.name.split(".")[-1].lower()
                         futures.append(executor.submit(
-                            self.file_processor.handle_upload, 
-                            f.name, 
+                            process_file,
+                            f.name,
                             file_type
                         ))
                     
-                    for i, future in enumerate(as_completed(futures), 1):
+                    for future in as_completed(futures):
                         try:
                             extracted.extend(future.result())
-                            outputs["progress"] = self._format_progress(i, len(files), "Processing files")
-                            yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
                         except Exception as e:
-                            logger.error(f"File processing failed: {e}")
-                            extracted.append({"error": str(e)})
-                
-                if files and os.path.exists(files[0].name):
-                    file_hash = hashlib.md5(open(files[0].name, "rb").read()).hexdigest()
-                
-                history.append({"role": "assistant", "content": "✅ Files processed successfully"})
-                outputs.update({
-                    "chatbot": history,
-                    "progress": self._format_progress(len(files), len(files), "Files processed")
-                })
-                yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
-
-            # Analyze content
-            text_content = "\n".join(json.dumps(item) for item in extracted)
-            chunks = self.text_analyzer.chunk_content(text_content)
-            full_analysis = ""
-            
-            for idx, chunk in enumerate(chunks, 1):
-                prompt = f"""
-Analyze this clinical documentation for potential missed diagnoses. Provide:
-1. Specific clinical findings with references (e.g., "Elevated BP (160/95) on page 3")
-2. Their clinical significance
-3. Urgency of review
-Use concise, continuous prose without bullet points. If no concerns, state "No missed diagnoses identified."
-
-Document Excerpt (Part {idx}/{len(chunks)}):
-{chunk[:1750]}
-"""
-                history.append({"role": "assistant", "content": ""})
-                outputs.update({
-                    "chatbot": history,
-                    "progress": self._format_progress(idx, len(chunks), "Analyzing")
-                })
-                yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
-                
-                # Stream analysis
-                chunk_response = ""
-                for update in self.process_stream(prompt, history):
-                    history[-1] = update
-                    chunk_response = update["content"]
-                    outputs.update({
-                        "chatbot": history,
-                        "progress": self._format_progress(idx, len(chunks), "Analyzing")
-                    })
-                    yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
+                            logger.error(f"File processing error: {e}")
+                            extracted.append({"error": f"Error processing file: {str(e)}"})
                 
-                full_analysis += f"--- Analysis Part {idx} ---\n{chunk_response}\n"
-                torch.cuda.empty_cache()
-                gc.collect()
+                file_hash_value = file_hash(files[0].name) if files else ""
+                history.append({"role": "assistant", "content": "✅ File processing complete"})
+                yield history, None, ""
 
-            # Final outputs
-            summary = self.text_analyzer.generate_summary(full_analysis)
-            report_path = DIRECTORIES["reports"] / f"{file_hash}_report.txt" if file_hash else None
+            # Convert extracted data to JSON text
+            text_content = "\n".join(json.dumps(item) for item in extracted)
             
-            if report_path:
-                with open(report_path, "w", encoding="utf-8") as f:
-                    f.write(full_analysis + "\n\nSUMMARY:\n" + summary)
+            # Tokenize and chunk the content properly
+            chunks = tokenize_and_chunk(text_content)
+            combined_response = ""
+            batch_size = 2  # Reduced batch size to prevent token overflow
             
-            outputs.update({
-                "download_output": str(report_path) if report_path and report_path.exists() else None,
-                "final_summary": summary,
-                "progress": {"visible": False}
-            })
-            yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
-
-        except Exception as e:
-            logger.error(f"Analysis failed: {e}")
-            history.append({"role": "assistant", "content": f"❌ Analysis error: {str(e)}"})
-            outputs.update({
-                "chatbot": history,
-                "final_summary": f"Error: {str(e)}",
-                "progress": {"visible": False}
-            })
-            yield (outputs["chatbot"], outputs["download_output"], outputs["final_summary"], outputs["progress"])
-    
-    def _format_progress(self, current: int, total: int, stage: str = "") -> Dict[str, Any]:
-        """Format progress update for UI"""
-        status = f"{stage} - {current}/{total}" if stage else f"{current}/{total}"
-        return {"value": status, "visible": True, "label": f"Progress: {status}"}
-
-    def create_interface(self) -> gr.Blocks:
-        """Build the Gradio interface"""
-        css = """
-        /* ==================== BASE STYLES ==================== */
-        :root {
-            --primary-color: #4f46e5;
-            --primary-dark: #4338ca;
-            --border-radius: 8px;
-            --transition: all 0.3s ease;
-            --shadow: 0 4px 12px rgba(0,0,0,0.1);
-            --font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
-            --background: #ffffff;
-            --text-color: #1e293b;
-            --chat-bg: #f8fafc;
-            --message-bg: #e2e8f0;
-            --panel-bg: rgba(248, 250, 252, 0.9);
-            --panel-dark-bg: rgba(30, 41, 59, 0.9);
-        }
-
-        [data-theme="dark"] {
-            --background: #1e2a44;
-            --text-color: #f1f5f9;
-            --chat-bg: #2d3b55;
-            --message-bg: #475569;
-            --panel-bg: var(--panel-dark-bg);
-        }
-
-        body, .gradio-container {
-            font-family: var(--font-family);
-            background: var(--background);
-            color: var(--text-color);
-            margin: 0;
-            padding: 0;
-            transition: var(--transition);
-        }
-
-        /* ==================== LAYOUT ==================== */
-        .gradio-container {
-            max-width: 1200px;
-            margin: 0 auto;
-            padding: 1.5rem;
-            display: flex;
-            flex-direction: column;
-            gap: 1.5rem;
-        }
-
-        .chat-container {
-            background: var(--chat-bg);
-            border-radius: var(--border-radius);
-            border: 1px solid #e2e8f0;
-            padding: 1.5rem;
-            min-height: 50vh;
-            max-height: 80vh;
-            overflow-y: auto;
-            box-shadow: var(--shadow);
-            margin-bottom: 4rem;
-        }
-
-        .summary-panel {
-            background: var(--panel-bg);
-            border-left: 4px solid var(--primary-color);
-            padding: 1rem;
-            border-radius: var(--border-radius);
-            margin-bottom: 1rem;
-            box-shadow: var(--shadow);
-            backdrop-filter: blur(8px);
-        }
-
-        .upload-area {
-            border: 2px dashed #cbd5e1;
-            border-radius: var(--border-radius);
-            padding: 1.5rem;
-            margin: 0.75rem 0;
-            transition: var(--transition);
-        }
-
-        .upload-area:hover {
-            border-color: var(--primary-color);
-            background: rgba(79, 70, 229, 0.05);
-        }
-
-        /* ==================== COMPONENTS ==================== */
-        .chat__message {
-            margin: 0.75rem 0;
-            padding: 0.75rem 1rem;
-            border-radius: var(--border-radius);
-            max-width: 85%;
-            transition: var(--transition);
-            background: var(--message-bg);
-            border: 1px solid rgba(0,0,0,0.05);
-            animation: messageFade 0.3s ease;
-        }
-
-        .chat__message:hover {
-            transform: translateY(-2px);
-            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        }
-
-        .chat__message.user {
-            background: linear-gradient(135deg, var(--primary-color), var(--primary-dark));
-            color: white;
-            margin-left: auto;
-        }
-
-        .chat__message.assistant {
-            background: var(--message-bg);
-            color: var(--text-color);
-        }
-
-        .input-container {
-            display: flex;
-            align-items: center;
-            gap: 0.75rem;
-            background: var(--chat-bg);
-            padding: 0.75rem 1rem;
-            border-radius: 1.5rem;
-            box-shadow: var(--shadow);
-            position: sticky;
-            bottom: 1rem;
-            z-index: 10;
-        }
-
-        .input__textbox {
-            flex-grow: 1;
-            border: none;
-            background: transparent;
-            color: var(--text-color);
-            outline: none;
-            font-size: 1rem;
-        }
-
-        .input__textbox:focus {
-            border-bottom: 2px solid var(--primary-color);
-        }
-
-        .submit-btn {
-            background: linear-gradient(135deg, var(--primary-color), var(--primary-dark));
-            color: white;
-            border: none;
-            border-radius: 1rem;
-            padding: 0.5rem 1.25rem;
-            font-size: 0.9rem;
-            transition: var(--transition);
-        }
-
-        .submit-btn:hover {
-            transform: scale(1.05);
-        }
-
-        .submit-btn:active {
-            animation: glow 0.3s ease;
-        }
-
-        .tooltip {
-            position: relative;
-        }
-
-        .tooltip:hover::after {
-            content: attr(data-tip);
-            position: absolute;
-            top: -2.5rem;
-            left: 50%;
-            transform: translateX(-50%);
-            background: #1e293b;
-            color: white;
-            padding: 0.4rem 0.8rem;
-            border-radius: 0.4rem;
-            font-size: 0.85rem;
-            max-width: 200px;
-            white-space: normal;
-            text-align: center;
-            z-index: 1000;
-            animation: fadeIn 0.3s ease;
-        }
-
-        .progress-tracker {
-            position: relative;
-            padding: 0.5rem;
-            background: var(--message-bg);
-            border-radius: var(--border-radius);
-            margin-top: 0.75rem;
-            overflow: hidden;
-        }
-
-        .progress-tracker::before {
-            content: '';
-            position: absolute;
-            top: 0;
-            left: 0;
-            height: 100%;
-            width: 0;
-            background: linear-gradient(to right, var(--primary-color), var(--primary-dark));
-            opacity: 0.3;
-            animation: progress 2s ease-in-out infinite;
-        }
-
-        /* ==================== ANIMATIONS ==================== */
-        @keyframes glow {
-            0%, 100% { transform: scale(1); opacity: 1; }
-            50% { transform: scale(1.1); opacity: 0.8; }
-        }
-
-        @keyframes fadeIn {
-            from { opacity: 0; }
-            to { opacity: 1; }
-        }
-
-        @keyframes messageFade {
-            from { opacity: 0; transform: translateY(10px) scale(0.95); }
-            to { opacity: 1; transform: translateY(0) scale(1); }
-        }
-
-        @keyframes progress {
-            0% { width: 0; }
-            50% { width: 60%; }
-            100% { width: 0; }
-        }
-
-        /* ==================== THEMES ==================== */
-        [data-theme="dark"] .chat-container {
-            border-color: #475569;
-        }
-
-        [data-theme="dark"] .upload-area {
-            border-color: #64748b;
-        }
-
-        [data-theme="dark"] .upload-area:hover {
-            background: rgba(79, 70, 229, 0.1);
-        }
-
-        [data-theme="dark"] .summary-panel {
-            border-left-color: #818cf8;
-        }
-
-        /* ==================== MEDIA QUERIES ==================== */
-        @media (max-width: 768px) {
-            .gradio-container {
-                padding: 1rem;
-            }
-
-            .chat-container {
-                min-height: 40vh;
-                max-height: 70vh;
-                margin-bottom: 3.5rem;
-            }
-
-            .summary-panel {
-                padding: 0.75rem;
-            }
-
-            .upload-area {
-                padding: 1rem;
-            }
-
-            .input-container {
-                gap: 0.5rem;
-                padding: 0.5rem;
-            }
-
-            .submit-btn {
-                padding: 0.4rem 1rem;
-            }
-        }
-
-        @media (max-width: 480px) {
-            .chat-container {
-                padding: 1rem;
-                margin-bottom: 3rem;
-            }
-
-            .input-container {
-                flex-direction: column;
-                padding: 0.5rem;
-            }
-
-            .input__textbox {
-                font-size: 0.9rem;
-            }
-
-            .submit-btn {
-                width: 100%;
-                padding: 0.5rem;
-                font-size: 0.85rem;
-            }
-
-            .chat__message {
-                max-width: 90%;
-                padding: 0.5rem 0.75rem;
-            }
-
-            .tooltip:hover::after {
-                top: auto;
-                bottom: -2.5rem;
-                max-width: 80vw;
-            }
-        }
-        """
-        
-        js = """
-        function applyTheme(theme) {
-            document.documentElement.setAttribute('data-theme', theme);
-            localStorage.setItem('theme', theme);
-        }
-
-        document.addEventListener('DOMContentLoaded', () => {
-            const savedTheme = localStorage.getItem('theme') || 'light';
-            applyTheme(savedTheme);
-        });
-        """
-
-        with gr.Blocks(
-            theme=gr.themes.Soft(
-                primary_hue="indigo",
-                secondary_hue="blue",
-                neutral_hue="slate"
-            ),
-            title="Clinical Oversight Assistant",
-            css=css,
-            js=js
-        ) as app:
-            # Header
-            gr.Markdown("""
-            <div style='text-align: center; margin-bottom: 24px;'>
-                <h1 style='color: var(--primary-color); margin-bottom: 8px;'>🩺 Clinical Oversight Assistant</h1>
-                <p style='color: #64748b;'>
-                    AI-powered analysis for identifying potential missed diagnoses in patient records
-                </p>
-            </div>
-            """)
-
-            with gr.Row(equal_height=False):
-                # Main Chat Panel
-                with gr.Column(scale=3):
-                    gr.Markdown(
-                        "<div class='tooltip' data-tip='View conversation history'>**Clinical Analysis Conversation**</div>"
-                    )
-                    chatbot = gr.Chatbot(
-                        label="",
-                        height=650,
-                        show_copy_button=True,
-                        avatar_images=(
-                            "assets/user.png", 
-                            "assets/assistant.png"
-                        ) if Path("assets/user.png").exists() else None,
-                        bubble_full_width=False,
-                        type="messages",
-                        elem_classes=["chat-container"]
-                    )
-
-                # Results Panel
-                with gr.Column(scale=1):
-                    with gr.Group():
-                        gr.Markdown(
-                            "<div class='tooltip' data-tip='Summary of findings'>**Clinical Summary**</div>"
-                        )
-                        final_summary = gr.Markdown(
-                            "<div class='tooltip' data-tip='Analysis results'>Analysis results will appear here...</div>",
-                            elem_classes=["summary-panel"]
+            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(
+                            batch_idx + i + 1,
+                            len(chunks),
+                            chunk=chunk[:1800]  # Conservative chunk size
                         )
+                        for i, chunk in enumerate(batch_chunks)
+                    ]
                     
-                    with gr.Group():
-                        gr.Markdown(
-                            "<div class='tooltip' data-tip='Download report'>**Report Export**</div>"
-                        )
-                        download_output = gr.File(
-                            label="Download Full Analysis",
-                            visible=False,
-                            interactive=False
-                        )
+                    progress((batch_idx) / len(chunks), 
+                           desc=f"Analyzing batch {(batch_idx // batch_size) + 1}/{(len(chunks) + batch_size - 1) // batch_size}")
+                    
+                    # Process batch in parallel
+                    with ThreadPoolExecutor(max_workers=len(batch_prompts)) as executor:
+                        future_to_prompt = {
+                            executor.submit(
+                                agent.run_gradio_chat,
+                                prompt, [], 0.2, 512, 2048, False, []
+                            ): prompt
+                            for prompt in batch_prompts
+                        }
+                        
+                        for future in as_completed(future_to_prompt):
+                            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 + " "
+                            
+                            combined_response += f"--- Analysis for Chunk {batch_idx + 1} ---\n{chunk_response.strip()}\n"
+                            history[-1] = {"role": "assistant", "content": combined_response.strip()}
+                            yield history, None, ""
+                            
+                            # Clean up memory
+                            torch.cuda.empty_cache()
+                            gc.collect()
+
+                # Generate final summary
+                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
 
-            # Input Section
-            with gr.Row():
-                file_upload = gr.File(
-                    file_types=[".pdf", ".csv", ".xls", ".xlsx"],
-                    file_count="multiple",
-                    label="Upload Patient Records",
-                    elem_classes=["upload-area"],
-                    elem_id="file-upload"
-                )
-
-            with gr.Row(elem_classes=["input-container"]):
-                user_input = gr.Textbox(
-                    placeholder="Enter your clinical query or analysis request...",
-                    show_label=False,
-                    container=False,
-                    scale=7,
-                    autofocus=True,
-                    elem_classes=["input__textbox"],
-                    elem_id="user-input"
-                )
-                submit_btn = gr.Button(
-                    "Analyze",
-                    variant="primary",
-                    scale=1,
-                    min_width=120,
-                    elem_classes=["submit-btn"],
-                    elem_id="submit-btn"
-                )
-
-            # Hidden progress tracker
-            progress_tracker = gr.Textbox(
-                label="Analysis Progress",
-                visible=False,
-                interactive=False,
-                elem_classes=["progress-tracker"],
-                elem_id="progress-tracker"
-            )
-
-            # Event handlers
-            submit_btn.click(
-                self.analyze_records,
-                inputs=[user_input, chatbot, file_upload],
-                outputs=[chatbot, download_output, final_summary, progress_tracker],
-                show_progress="hidden"
-            )
-            
-            user_input.submit(
-                self.analyze_records,
-                inputs=[user_input, chatbot, file_upload],
-                outputs=[chatbot, download_output, final_summary, progress_tracker],
-                show_progress="hidden"
-            )
-            
-            app.load(
-                lambda: [[], None, "<div class='tooltip' data-tip='Analysis results'>Analysis results will appear here...</div>", "", None, {"visible": False}],
-                outputs=[chatbot, download_output, final_summary, user_input, file_upload, progress_tracker],
-                queue=False
-            )
+            except Exception as e:
+                logger.error("Analysis error: %s", e)
+                history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
+                yield history, None, f"Error occurred during analysis: {str(e)}"
 
-        return app
+        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
 
-# ==================== APPLICATION ENTRY POINT ====================
 if __name__ == "__main__":
     try:
-        logger.info("Launching Clinical Oversight Assistant...")
-        clinical_app = ClinicalAgent()
-        interface = clinical_app.create_interface()
-        
-        interface.queue(
-            api_open=False,
-            max_size=20
-        ).launch(
+        logger.info("Launching app...")
+        agent = init_agent()
+        demo = create_ui(agent)
+        demo.queue(api_open=False).launch(
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            allowed_paths=[str(DIRECTORIES["reports"])],
+            allowed_paths=[report_dir],
             share=False
         )
-    except Exception as e:
-        logger.error(f"Application failed to start: {e}")
-        raise
     finally:
         if torch.distributed.is_initialized():
             torch.distributed.destroy_process_group()