Update app.py

app.py

@@ -27,11 +27,11 @@ logger = logging.getLogger(__name__)
 
 # Constants
 MAX_TOKENS = 1800
-BATCH_SIZE = 1  # Reduced to minimize memory pressure
+BATCH_SIZE = 1
 MAX_WORKERS = 2
-CHUNK_SIZE = 5  # Smaller chunks for PDF processing
+CHUNK_SIZE = 5
 MODEL_MAX_TOKENS = 131072
-MAX_TEXT_LENGTH = 500000  # Limit raw text length before tokenization
+MAX_TEXT_LENGTH = 500000
 
 # Persistent directory setup
 persistent_dir = "/data/hf_cache"
@@ -63,17 +63,14 @@ 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""):
@@ -81,12 +78,10 @@ def file_hash(path: str) -> str:
     return hash_md5.hexdigest()
 
 def extract_pdf_page(page, tokenizer, max_tokens=MAX_TOKENS) -> List[str]:
-    """Extract and chunk a single page with token limit"""
     try:
         text = page.extract_text() or ""
         text = sanitize_utf8(text)
-        if len(text) > MAX_TEXT_LENGTH // 10:  # Per-page text limit
-            logger.warning(f"Page {page.page_number} text too long ({len(text)}). Truncating.")
+        if len(text) > MAX_TEXT_LENGTH // 10:
             text = text[:MAX_TEXT_LENGTH // 10]
         
         tokens = tokenizer.encode(text, add_special_tokens=False)
@@ -111,7 +106,6 @@ def extract_pdf_page(page, tokenizer, max_tokens=MAX_TOKENS) -> List[str]:
         return []
 
 def extract_all_pages(file_path: str, progress_callback=None) -> List[str]:
-    """Extract PDF pages with early token-based chunking"""
     try:
         tokenizer = get_tokenizer()
         with pdfplumber.open(file_path) as pdf:
@@ -134,7 +128,7 @@ def extract_all_pages(file_path: str, progress_callback=None) -> List[str]:
                         for chunk in page_chunks:
                             chunk_tokens = len(tokenizer.encode(chunk, add_special_tokens=False))
                             if total_tokens + chunk_tokens > MODEL_MAX_TOKENS:
-                                logger.warning(f"Total tokens ({total_tokens + chunk_tokens}) exceed model limit ({MODEL_MAX_TOKENS}). Stopping.")
+                                logger.warning(f"Total tokens exceed model limit. Stopping.")
                                 return results
                             results.append(chunk)
                             total_tokens += chunk_tokens
@@ -151,31 +145,56 @@ def extract_all_pages(file_path: str, progress_callback=None) -> List[str]:
         return [f"PDF processing error: {str(e)}"]
 
 def excel_to_json(file_path: str) -> List[Dict]:
-    """Optimized Excel processing with chunking"""
     try:
-        for engine in ['openpyxl', 'xlrd']:
+        # Try with openpyxl first
+        try:
+            with pd.ExcelFile(file_path, engine='openpyxl') as excel_file:
+                sheets = excel_file.sheet_names
+                results = []
+                for sheet_name in sheets:
+                    df = pd.read_excel(
+                        excel_file,
+                        sheet_name=sheet_name,
+                        header=None,
+                        dtype=str,
+                        na_filter=False
+                    )
+                    if not df.empty:
+                        results.append({
+                            "filename": f"{os.path.basename(file_path)} - {sheet_name}",
+                            "rows": df.values.tolist(),
+                            "type": "excel"
+                        })
+                return results if results else [{"error": "No data found in any sheet"}]
+        except Exception as openpyxl_error:
+            # Fallback to 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")
+                with pd.ExcelFile(file_path, engine='xlrd') as excel_file:
+                    sheets = excel_file.sheet_names
+                    results = []
+                    for sheet_name in sheets:
+                        df = pd.read_excel(
+                            excel_file,
+                            sheet_name=sheet_name,
+                            header=None,
+                            dtype=str,
+                            na_filter=False
+                        )
+                        if not df.empty:
+                            results.append({
+                                "filename": f"{os.path.basename(file_path)} - {sheet_name}",
+                                "rows": df.values.tolist(),
+                                "type": "excel"
+                            })
+                    return results if results else [{"error": "No data found in any sheet"}]
+            except Exception as xlrd_error:
+                logger.error(f"Excel processing failed: {xlrd_error}")
+                return [{"error": f"Excel processing failed: {str(xlrd_error)}"}]
     except Exception as e:
-        logger.error(f"Excel processing error: {e}")
-        return [{"error": f"Excel processing error: {str(e)}"}]
+        logger.error(f"Excel file opening error: {e}")
+        return [{"error": f"Excel file opening error: {str(e)}"}]
 
 def csv_to_json(file_path: str) -> List[Dict]:
-    """Optimized CSV processing with chunking"""
     try:
         chunks = []
         for chunk in pd.read_csv(
@@ -201,7 +220,6 @@ def csv_to_json(file_path: str) -> List[Dict]:
 
 @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":
             chunks = extract_all_pages(file_path)
@@ -218,80 +236,17 @@ def process_file_cached(file_path: str, file_type: str) -> List[Dict]:
         else:
             return [{"error": f"Unsupported file type: {file_type}"}]
     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 with early validation"""
-    if len(text) > MAX_TEXT_LENGTH:
-        logger.warning(f"Text length ({len(text)}) exceeds limit ({MAX_TEXT_LENGTH}). Truncating.")
-        text = text[:MAX_TEXT_LENGTH]
-
-    tokenizer = get_tokenizer()
-    tokens = tokenizer.encode(text, add_special_tokens=False)
-    if len(tokens) > MODEL_MAX_TOKENS:
-        logger.error(f"Token count ({len(tokens)}) exceeds model limit ({MODEL_MAX_TOKENS}).")
-        return [text[:MAX_TEXT_LENGTH // 10]]  # Fallback to small chunk
-
-    chunks = []
-    current_chunk = []
-    current_length = 0
-
-    for token in tokens:
-        if current_length + 1 > max_tokens:
-            chunks.append(tokenizer.decode(current_chunk))
-            current_chunk = [token]
-            current_length = 1
-        else:
-            current_chunk.append(token)
-            current_length += 1
-    
-    if current_chunk:
-        chunks.append(tokenizer.decode(current_chunk))
-    
-    return chunks
-
-def log_system_usage(tag=""):
-    """Optimized system monitoring"""
-    try:
-        cpu = psutil.cpu_percent(interval=0.5)
-        mem = psutil.virtual_memory()
-        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(f"[{tag}] Monitor failed: {e}")
+        logger.error(f"Error processing file: {e}")
+        return [{"error": f"Error processing file: {str(e)}"}]
 
 def clean_response(text: str) -> str:
-    """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:
@@ -314,74 +269,11 @@ def clean_response(text: str) -> str:
             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."
-    
-    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 = []
-    current_section = None
-    
-    for line in combined_response.splitlines():
-        line = line.strip()
-        if not line:
-            continue
-        
-        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)
-    
-    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))]
-    
-    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
-
 @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")
@@ -399,15 +291,12 @@ 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 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).
+Analyze the patient record excerpt for missed diagnoses. Provide detailed, evidence-based analysis.
 Patient Record Excerpt (Chunk {0} of {1}):
 {chunk}
 """
@@ -417,8 +306,8 @@ Patient Record Excerpt (Chunk {0} of {1}):
         
         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)
+                chatbot = gr.Chatbot(label="Analysis Summary", height=600)
+                msg_input = gr.Textbox(placeholder="Ask about potential oversights...")
                 send_btn = gr.Button("Analyze", variant="primary")
                 file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
             
@@ -428,7 +317,6 @@ Patient Record Excerpt (Chunk {0} of {1}):
                 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, ""
 
@@ -442,31 +330,38 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     
                     if cache_key in cache:
                         extracted.extend(cache[cache_key])
+                        history.append({"role": "assistant", "content": f"Using cached data for {os.path.basename(f.name)}"})
+                        yield history, None, ""
                     else:
                         result = process_file_cached(f.name, file_type)
-                        cache[cache_key] = result
-                        extracted.extend(result)
+                        if result and not (len(result) == 1 and "error" in result[0]):
+                            cache[cache_key] = result
+                            extracted.extend(result)
+                            history.append({"role": "assistant", "content": f"Processed {os.path.basename(f.name)}"})
+                            yield history, None, ""
+                        else:
+                            error_msg = result[0]["error"] if result else "Unknown error"
+                            history.append({"role": "assistant", "content": f"Failed to process {os.path.basename(f.name)}: {error_msg}"})
+                            yield history, None, error_msg
+                            return
                 
                 file_hash_value = file_hash(files[0].name) if files else ""
-                history.append({"role": "assistant", "content": "✅ File processing complete"})
-                yield history, None, ""
 
             if not extracted:
-                history.append({"role": "assistant", "content": "❌ No valid content extracted. Please upload a supported file."})
-                yield history, None, "No valid content extracted."
+                history.append({"role": "assistant", "content": "❌ No valid content extracted"})
+                yield history, None, "No valid content extracted"
+                return
+
+            chunks = [item["content"] for item in extracted if "content" in item]
+            if not chunks:
+                history.append({"role": "assistant", "content": "❌ No processable content found"})
+                yield history, None, "No processable content found"
                 return
 
             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:
-                chunks = [item["content"] for item in extracted if "content" in item]
-                if not chunks:
-                    history.append({"role": "assistant", "content": "❌ No processable content found in the file."})
-                    yield history, None, "No processable content found."
-                    return
-                
                 for batch_idx in range(0, len(chunks), BATCH_SIZE):
                     batch_chunks = chunks[batch_idx:batch_idx + BATCH_SIZE]
                     
@@ -483,7 +378,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                            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 = {
+                        futures = {
                             executor.submit(
                                 agent.run_quick_summary,
                                 chunk, 0.2, 256, 1024
@@ -491,74 +386,31 @@ Patient Record Excerpt (Chunk {0} of {1}):
                             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, ""
-                            except Exception as e:
-                                logger.error(f"Quick summary error for chunk {batch_idx + chunk_idx + 1}: {e}")
-                                history[-1] = {"role": "assistant", "content": f"Error processing chunk {batch_idx + chunk_idx + 1}: {str(e)}"}
-                                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]
+                        for future in as_completed(futures):
+                            chunk_idx = futures[future]
                             try:
-                                for chunk_output in future.result():
-                                    if isinstance(chunk_output, list):
-                                        for msg in chunk_output:
-                                            if isinstance(msg, gr.ChatMessage) and msg.content:
-                                                cleaned_content = clean_response(msg.content)
-                                                if cleaned_content and cleaned_content != "No missed diagnoses identified.":
-                                                    combined_response += cleaned_content + "\n"
-                                                    history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                                                    yield history, report_path, ""
+                                response = clean_response(future.result())
+                                if response:
+                                    combined_response += f"--- Analysis for Chunk {batch_idx + chunk_idx + 1} ---\n{response}\n"
+                                    history[-1] = {"role": "assistant", "content": combined_response.strip()}
+                                    yield history, None, ""
                             except Exception as e:
-                                logger.error(f"Detailed analysis error for chunk {batch_idx + chunk_idx + 1}: {e}")
-                                history[-1] = {"role": "assistant", "content": f"Error in detailed analysis for chunk {batch_idx + chunk_idx + 1}: {str(e)}"}
+                                logger.error(f"Chunk processing error: {e}")
+                                history[-1] = {"role": "assistant", "content": f"Error processing chunk: {str(e)}"}
                                 yield history, None, ""
                             finally:
                                 del future
                                 torch.cuda.empty_cache()
                                 gc.collect()
 
-                summary = summarize_findings(combined_response)
-                
-                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
+                summary = "Analysis complete. " + ("Download full report below." if report_path and os.path.exists(report_path) else "")
+                history.append({"role": "assistant", "content": "Analysis completed successfully"})
+                yield history, report_path, summary
 
             except Exception as 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)}"
+                yield history, None, f"Error occurred: {str(e)}"
             finally:
                 torch.cuda.empty_cache()
                 gc.collect()
@@ -578,22 +430,14 @@ Patient Record Excerpt (Chunk {0} of {1}):
 
 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().launch(
             server_name="0.0.0.0",
             server_port=7860,
-            show_error=True,
-            allowed_paths=[report_dir],
-            share=False
+            show_error=True
         )
     except Exception as e:
         logger.error(f"Fatal error: {e}")
-        raise
-    finally:
-        if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
+        raise