Update app.py

app.py

@@ -2,23 +2,18 @@ import sys
 import os
 import pandas as pd
 import pdfplumber
-import json
 import gradio as gr
-from typing import List, Dict, Optional, Generator
+from typing import List, Dict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
 import re
-import psutil
-import subprocess
 import logging
 import torch
 import gc
 from diskcache import Cache
-import time
 from transformers import AutoTokenizer
 from functools import lru_cache
-import numpy as np
 from difflib import SequenceMatcher
 
 # Configure logging
@@ -32,6 +27,7 @@ MAX_WORKERS = 2
 CHUNK_SIZE = 5
 MODEL_MAX_TOKENS = 131072
 MAX_TEXT_LENGTH = 500000
+MAX_ROWS_TO_PROCESS = 1000  # Limit for Excel/CSV rows
 
 # Persistent directory setup
 persistent_dir = "/data/hf_cache"
@@ -41,17 +37,11 @@ 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.makedirs(report_dir, 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"
 })
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -60,7 +50,7 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Initialize cache with 10GB limit
+# Initialize cache
 cache = Cache(file_cache_dir, size_limit=10 * 1024**3)
 
 @lru_cache(maxsize=1)
@@ -99,70 +89,47 @@ def extract_pdf_page(page, tokenizer, max_tokens=MAX_TOKENS) -> List[str]:
                     current_length += 1
             if current_chunk:
                 chunks.append(tokenizer.decode(current_chunk))
-            return [f"=== Page {page.page_number} ===\n{c}" for c in chunks]
-        return [f"=== Page {page.page_number} ===\n{text}"]
+            return chunks
+        return [text]
     except Exception as e:
         logger.warning(f"Error extracting page {page.page_number}: {str(e)}")
         return []
 
-def extract_all_pages(file_path: str, progress_callback=None) -> List[str]:
+def extract_all_pages(file_path: str) -> List[str]:
     try:
         tokenizer = get_tokenizer()
         with pdfplumber.open(file_path) as pdf:
             total_pages = len(pdf.pages)
             if total_pages == 0:
-                logger.error("PDF has 0 pages - may be corrupted or empty")
-                return []
+                return ["PDF appears to be empty"]
 
         results = []
-        total_tokens = 0
-        for chunk_start in range(0, total_pages, CHUNK_SIZE):
-            chunk_end = min(chunk_start + CHUNK_SIZE, total_pages)
-            
-            with pdfplumber.open(file_path) as pdf:
-                with ThreadPoolExecutor(max_workers=min(CHUNK_SIZE, 2)) as executor:
-                    futures = [executor.submit(extract_pdf_page, pdf.pages[i], tokenizer) 
-                             for i in range(chunk_start, chunk_end)]
-                    
-                    for future in as_completed(futures):
-                        page_chunks = future.result()
-                        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("Total tokens exceed model limit. Stopping.")
-                                return results
-                            results.append(chunk)
-                            total_tokens += chunk_tokens
-                    
-                    if progress_callback:
-                        progress_callback(min(chunk_end, total_pages), total_pages)
-            
-            del pdf
-            gc.collect()
-        
-        if not results:
-            logger.error("No content extracted from PDF - may be scanned or encrypted")
-            return ["PDF appears to be empty or unreadable"]
+        for i in range(0, min(total_pages, 50)):  # Limit to first 50 pages
+            try:
+                page = pdf.pages[i]
+                chunks = extract_pdf_page(page, tokenizer)
+                for chunk in chunks:
+                    results.append(f"=== Page {i+1} ===\n{chunk}")
+            except Exception as e:
+                logger.warning(f"Error processing page {i+1}: {str(e)}")
+                continue
         
-        return results
+        return results if results else ["Could not extract text from PDF"]
     except Exception as e:
         logger.error(f"PDF processing error: {e}")
         return [f"PDF processing error: {str(e)}"]
 
 def excel_to_json(file_path: str) -> List[Dict]:
-    """Enhanced Excel processing with multiple engine support"""
-    engines = ['openpyxl', 'xlrd', 'odf']
-    last_error = None
-    
+    engines = ['openpyxl', 'xlrd']
     for engine in engines:
         try:
             with pd.ExcelFile(file_path, engine=engine) as excel_file:
                 sheets = excel_file.sheet_names
                 if not sheets:
-                    return [{"error": "No sheets found in Excel file"}]
+                    return [{"error": "No sheets found"}]
                 
                 results = []
-                for sheet_name in sheets:
+                for sheet_name in sheets[:3]:  # Limit to first 3 sheets
                     try:
                         df = pd.read_excel(
                             excel_file,
@@ -170,99 +137,70 @@ def excel_to_json(file_path: str) -> List[Dict]:
                             header=None,
                             dtype=str,
                             na_filter=False,
-                            engine=engine
+                            nrows=MAX_ROWS_TO_PROCESS  # Limit rows
                         )
                         if not df.empty:
-                            # Convert all cells to string and clean
-                            df = df.applymap(lambda x: str(x).strip() if pd.notna(x) else "")
+                            rows = df.head(MAX_ROWS_TO_PROCESS).values.tolist()
                             results.append({
-                                "filename": f"{os.path.basename(file_path)} - {sheet_name}",
-                                "rows": df.values.tolist(),
-                                "type": "excel",
+                                "filename": os.path.basename(file_path),
                                 "sheet": sheet_name,
-                                "dimensions": f"{len(df)} rows x {len(df.columns)} cols"
+                                "rows": rows,
+                                "type": "excel"
                             })
-                    except Exception as sheet_error:
-                        logger.warning(f"Error processing sheet {sheet_name}: {sheet_error}")
+                    except Exception as e:
+                        logger.warning(f"Error processing sheet {sheet_name}: {str(e)}")
                         continue
                 
-                if results:
-                    logger.info(f"Successfully processed Excel file with {engine} engine")
-                    return results
-        except Exception as engine_error:
-            last_error = engine_error
+                return results if results else [{"error": "No readable data found"}]
+        except Exception as e:
+            logger.warning(f"Excel engine {engine} failed: {str(e)}")
             continue
     
-    return [{"error": f"Failed to process Excel file with all engines. Last error: {str(last_error)}"}]
+    return [{"error": "Could not process Excel file with any engine"}]
 
 def csv_to_json(file_path: str) -> List[Dict]:
     try:
-        chunks = []
-        for chunk in pd.read_csv(
+        df = 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()
+            nrows=MAX_ROWS_TO_PROCESS  # Limit rows
+        )
         if df.empty:
-            return [{"error": "CSV file is empty or could not be read"}]
+            return [{"error": "CSV file is empty"}]
             
         return [{
             "filename": os.path.basename(file_path),
             "rows": df.values.tolist(),
-            "type": "csv",
-            "dimensions": f"{len(df)} rows x {len(df.columns)} cols"
+            "type": "csv"
         }]
     except Exception as e:
         logger.error(f"CSV processing error: {e}")
         return [{"error": f"CSV processing error: {str(e)}"}]
 
-@lru_cache(maxsize=100)
 def process_file_cached(file_path: str, file_type: str) -> List[Dict]:
-    """Enhanced file processing with detailed logging"""
     try:
-        logger.info(f"Processing file: {file_path} (type: {file_type})")
+        logger.info(f"Processing {file_type} file: {os.path.basename(file_path)}")
         
         if file_type == "pdf":
             chunks = extract_all_pages(file_path)
-            if not chunks or (len(chunks) == 1 and "error" in chunks[0]):
-                return [{"error": chunks[0] if chunks else "PDF appears to be empty"}]
             return [{
                 "filename": os.path.basename(file_path),
                 "content": chunk,
-                "status": "initial",
-                "type": "pdf",
-                "page": i+1
-            } for i, chunk in enumerate(chunks)]
+                "type": "pdf"
+            } for chunk in chunks]
             
         elif file_type in ["xls", "xlsx"]:
-            result = excel_to_json(file_path)
-            if "error" in result[0]:
-                logger.error(f"Excel processing failed: {result[0]['error']}")
-            else:
-                logger.info(f"Excel processing successful - found {len(result)} sheets")
-            return result
+            return excel_to_json(file_path)
             
         elif file_type == "csv":
-            result = csv_to_json(file_path)
-            if "error" in result[0]:
-                logger.error(f"CSV processing failed: {result[0]['error']}")
-            else:
-                logger.info(f"CSV processing successful - found {len(result[0]['rows'])} rows")
-            return result
-            
-        else:
-            logger.warning(f"Unsupported file type: {file_type}")
-            return [{"error": f"Unsupported file type: {file_type}"}]
+            return csv_to_json(file_path)
             
+        return [{"error": f"Unsupported file type: {file_type}"}]
     except Exception as e:
-        logger.error(f"Error processing {file_path}: {str(e)}", exc_info=True)
+        logger.error(f"Error processing file: {e}")
         return [{"error": f"Error processing file: {str(e)}"}]
 
 def clean_response(text: str) -> str:
@@ -272,49 +210,25 @@ def clean_response(text: str) -> str:
     patterns = [
         (re.compile(r"\[.*?\]|\bNone\b", re.IGNORECASE), ""),
         (re.compile(r"\s+"), " "),
-        (re.compile(r"[^\w\s\.\,\(\)\-]"), ""),
     ]
     
     for pattern, repl in patterns:
         text = pattern.sub(repl, text)
     
-    sentences = text.split(". ")
-    unique_sentences = []
-    seen = set()
-    
-    for s in sentences:
-        if not s:
-            continue
-        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."
+    return text.strip()
 
 @lru_cache(maxsize=1)
 def init_agent():
     logger.info("Initializing model...")
     
-    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},
+        tool_files_dict={"new_tool": os.path.join(tool_cache_dir, "new_tool.json")},
         force_finish=True,
         enable_checker=False,
         step_rag_num=4,
         seed=100,
-        additional_default_tools=[],
     )
     agent.init_model()
     logger.info("Agent Ready")
@@ -322,8 +236,7 @@ def init_agent():
 
 def create_ui(agent):
     PROMPT_TEMPLATE = """
-Analyze the patient record excerpt for missed diagnoses. Provide detailed, evidence-based analysis.
-Patient Record Excerpt (Chunk {0} of {1}):
+Analyze this patient record excerpt for missed diagnoses (limit response to 500 tokens):
 {chunk}
 """
 
@@ -332,170 +245,89 @@ Patient Record Excerpt (Chunk {0} of {1}):
         
         with gr.Row():
             with gr.Column(scale=3):
-                chatbot = gr.Chatbot(label="Analysis Summary", height=600, value=[])
+                chatbot = gr.Chatbot(label="Analysis", height=500)
                 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")
+                file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="single")
             
             with gr.Column(scale=1):
-                final_summary = gr.Markdown(label="Missed Diagnoses Summary")
-                download_output = gr.File(label="Download Detailed Report")
-                progress_bar = gr.Progress()
+                final_summary = gr.Markdown("## Summary")
+                status = gr.Textbox(label="Status", interactive=False)
 
-        def analyze(message: str, history: List[List[str]], files: List, progress=gr.Progress()):
-            """Enhanced analysis with detailed file processing feedback"""
+        def analyze(message: str, history: List[List[str]], files: List):
             try:
-                if history is None:
-                    history = []
-                
-                history.append([message, None])
-                yield history, None, ""
-
                 if not files:
-                    history[-1][1] = "❌ Please upload a file to analyze"
-                    yield history, None, "No files uploaded"
-                    return
-
-                extracted = []
-                file_hash_value = ""
+                    return history, "Please upload a file first", "No file uploaded"
                 
-                for f in files:
-                    file_type = f.name.split(".")[-1].lower()
-                    logger.info(f"Processing file: {f.name} (type: {file_type})")
-                    
-                    cache_key = f"{file_hash(f.name)}_{file_type}"
-                    if cache_key in cache:
-                        cached_data = cache[cache_key]
-                        if isinstance(cached_data, list) and len(cached_data) > 0:
-                            extracted.extend(cached_data)
-                            history[-1][1] = f"✅ Using cached data for {os.path.basename(f.name)}"
-                            yield history, None, ""
-                            continue
-                    
-                    try:
-                        result = process_file_cached(f.name, file_type)
-                        if "error" in result[0]:
-                            history[-1][1] = f"❌ Error processing {os.path.basename(f.name)}: {result[0]['error']}"
-                            yield history, None, result[0]['error']
-                            return
-                        
-                        cache[cache_key] = result
-                        extracted.extend(result)
-                        history[-1][1] = f"✅ Processed {os.path.basename(f.name)}"
-                        yield history, None, ""
-                    except Exception as e:
-                        logger.error(f"File processing error: {e}", exc_info=True)
-                        history[-1][1] = f"❌ Critical error processing {os.path.basename(f.name)}"
-                        yield history, None, str(e)
-                        return
+                file = files[0]
+                file_type = file.name.split(".")[-1].lower()
+                history.append([message, None])
                 
-                file_hash_value = file_hash(files[0].name) if files else ""
+                # Process file
+                processed = process_file_cached(file.name, file_type)
+                if "error" in processed[0]:
+                    return history, processed[0]["error"], "File processing failed"
                 
-                # Debug extracted content
-                logger.info(f"Extracted content summary:")
-                for item in extracted:
-                    if "content" in item:
-                        logger.info(f"- {item['filename']}: {len(item['content'])} chars")
-                    elif "rows" in item:
-                        logger.info(f"- {item['filename']}: {len(item['rows'])} rows")
-                
-                if not extracted:
-                    history[-1][1] = "❌ No valid content extracted from files"
-                    yield history, None, "No valid content extracted"
-                    return
-
+                # Prepare chunks
                 chunks = []
-                for item in extracted:
+                for item in processed:
                     if "content" in item:
                         chunks.append(item["content"])
                     elif "rows" in item:
-                        # Convert Excel/CSV rows to text
-                        rows_text = "\n".join([", ".join(map(str, row)) for row in item["rows"]])
-                        chunks.append(f"=== {item['filename']} ===\n{rows_text}")
+                        rows_text = "\n".join([", ".join(map(str, row)) for row in item["rows"][:100]])  # Limit rows
+                        chunks.append(f"=== {item.get('sheet', 'Data')} ===\n{rows_text}")
                 
                 if not chunks:
-                    history[-1][1] = "❌ No processable content found in files"
-                    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
+                    return history, "No processable content found", "Content extraction failed"
+                
+                # Process chunks
+                responses = []
+                for i, chunk in enumerate(chunks[:5]):  # Limit to 5 chunks
+                    try:
+                        prompt = PROMPT_TEMPLATE.format(chunk=chunk[:5000])  # Limit chunk size
+                        response = agent.run_quick_summary(prompt, 0.2, 256, 500)  # Limit tokens
+                        cleaned = clean_response(response)
+                        if cleaned:
+                            responses.append(f"Analysis {i+1}:\n{cleaned}")
+                    except Exception as e:
+                        logger.warning(f"Error processing chunk {i+1}: {str(e)}")
+                        continue
+                
+                if not responses:
+                    return history, "No valid analysis generated", "Analysis failed"
+                
+                summary = "\n\n".join(responses)
+                history[-1][1] = summary
+                return history, "Analysis completed", "Success"
                 
-                try:
-                    for batch_idx in range(0, len(chunks), BATCH_SIZE):
-                        batch_chunks = chunks[batch_idx:batch_idx + BATCH_SIZE]
-                        
-                        progress(batch_idx / len(chunks), 
-                               desc=f"Processing batch {(batch_idx // BATCH_SIZE) + 1}/{(len(chunks) + BATCH_SIZE - 1) // BATCH_SIZE}")
-                        
-                        with ThreadPoolExecutor(max_workers=min(BATCH_SIZE, MAX_WORKERS)) as executor:
-                            futures = {
-                                executor.submit(
-                                    agent.run_quick_summary,
-                                    chunk, 0.2, 256, 1024
-                                ): idx
-                                for idx, chunk in enumerate(batch_chunks)
-                            }
-                            
-                            for future in as_completed(futures):
-                                chunk_idx = futures[future]
-                                try:
-                                    response = clean_response(future.result())
-                                    if response:
-                                        combined_response += f"\n--- Analysis for Chunk {batch_idx + chunk_idx + 1} ---\n{response}\n"
-                                        history[-1][1] = combined_response.strip()
-                                        yield history, None, ""
-                                except Exception as e:
-                                    logger.error(f"Chunk processing error: {e}")
-                                    history[-1][1] = f"Error processing chunk: {str(e)}"
-                                    yield history, None, ""
-                                finally:
-                                    del future
-                                    torch.cuda.empty_cache()
-                                    gc.collect()
-
-                    summary = "Analysis complete. " + ("Download full report below." if report_path and os.path.exists(report_path) else "")
-                    history.append(["Analysis completed", None])
-                    history[-1][1] = summary
-                    yield history, report_path, summary
-
-                except Exception as e:
-                    logger.error(f"Analysis error: {e}")
-                    history.append(["Analysis failed", None])
-                    history[-1][1] = f"❌ Error occurred: {str(e)}"
-                    yield history, None, f"Error occurred: {str(e)}"
-                finally:
-                    torch.cuda.empty_cache()
-                    gc.collect()
-
             except Exception as e:
-                logger.error(f"Unexpected error in analysis: {e}")
-                history.append(["System error", None])
-                history[-1][1] = f"❌ System error occurred: {str(e)}"
-                yield history, None, f"System error: {str(e)}"
+                logger.error(f"Analysis error: {e}")
+                return history, f"Error: {str(e)}", "Failed"
+            finally:
+                torch.cuda.empty_cache()
+                gc.collect()
 
         send_btn.click(
-            analyze, 
-            inputs=[msg_input, gr.State([]), file_upload], 
-            outputs=[chatbot, download_output, final_summary]
+            analyze,
+            inputs=[msg_input, chatbot, file_upload],
+            outputs=[chatbot, final_summary, status]
         )
         msg_input.submit(
-            analyze, 
-            inputs=[msg_input, gr.State([]), file_upload], 
-            outputs=[chatbot, download_output, final_summary]
+            analyze,
+            inputs=[msg_input, chatbot, file_upload],
+            outputs=[chatbot, final_summary, status]
         )
     
     return demo
 
 if __name__ == "__main__":
     try:
-        logger.info("Launching app...")
         agent = init_agent()
         demo = create_ui(agent)
-        demo.queue().launch(
+        demo.launch(
             server_name="0.0.0.0",
             server_port=7860,
-            show_error=True
+            share=False
         )
     except Exception as e:
         logger.error(f"Fatal error: {e}")