Update app.py

app.py

@@ -2,11 +2,12 @@ import sys
 import os
 import pandas as pd
 import gradio as gr
-from typing import List, Tuple, Dict, Any, Union, Generator
+from typing import List, Tuple, Dict, Any, Union
 import shutil
 import re
 from datetime import datetime
 import time
+from transformers import AutoTokenizer
 import asyncio
 import logging
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -23,7 +24,7 @@ report_dir = os.path.join(persistent_dir, "reports")
 for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
     os.makedirs(directory, exist_ok=True)
 
-os.environ["HF_HOME"] = model_cache_dir  # Using HF_HOME instead of TRANSFORMERS_CACHE
+os.environ["HF_HOME"] = model_cache_dir  # Using HF_HOME as specified
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 src_path = os.path.abspath(os.path.join(current_dir, "src"))
@@ -31,15 +32,22 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Updated token limits as specified
+# Constants
 MAX_MODEL_TOKENS = 131072  # TxAgent's max token limit
 MAX_CHUNK_TOKENS = 32768   # Larger chunks to reduce number of chunks
 MAX_NEW_TOKENS = 512       # Optimized for fast generation
 PROMPT_OVERHEAD = 500      # Estimated tokens for prompt template
 MAX_CONCURRENT = 8         # High concurrency for A100 80GB
 
+# Initialize tokenizer for precise token counting
+try:
+    tokenizer = AutoTokenizer.from_pretrained("mims-harvard/TxAgent-T1-Llama-3.1-8B")
+except Exception as e:
+    print(f"Warning: Could not load tokenizer, falling back to heuristic: {str(e)}")
+    tokenizer = None
+
 # Setup logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 def clean_response(text: str) -> str:
@@ -53,9 +61,13 @@ def clean_response(text: str) -> str:
     return text.strip()
 
 def estimate_tokens(text: str) -> int:
-    return len(text) // 3.5 + 1  # More conservative estimate
+    """Estimate tokens using tokenizer if available, else fall back to heuristic."""
+    if tokenizer:
+        return len(tokenizer.encode(text, add_special_tokens=False))
+    return len(text) // 3.5 + 1  # Consistent with your heuristic
 
 def extract_text_from_excel(file_path: str) -> str:
+    """Extract text from all sheets in an Excel file."""
     all_text = []
     try:
         xls = pd.ExcelFile(file_path)
@@ -70,12 +82,12 @@ def extract_text_from_excel(file_path: str) -> str:
         raise ValueError(f"Failed to process Excel file: {str(e)}")
     return "\n".join(all_text)
 
-def split_text_into_chunks(text: str) -> List[str]:
-    """Split text into chunks respecting MAX_CHUNK_TOKENS and PROMPT_OVERHEAD"""
-    effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD
-    if effective_max <= 0:
-        raise ValueError("Effective max tokens must be positive")
-    
+def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
+    """Split text into chunks respecting MAX_CHUNK_TOKENS and PROMPT_OVERHEAD."""
+    effective_max_tokens = max_tokens - PROMPT_OVERHEAD
+    if effective_max_tokens <= 0:
+        raise ValueError(f"Effective max tokens ({effective_max_tokens}) must be positive.")
+
     lines = text.split("\n")
     chunks = []
     current_chunk = []
@@ -83,7 +95,7 @@ def split_text_into_chunks(text: str) -> List[str]:
 
     for line in lines:
         line_tokens = estimate_tokens(line)
-        if current_tokens + line_tokens > effective_max:
+        if current_tokens + line_tokens > effective_max_tokens:
             if current_chunk:
                 chunks.append("\n".join(current_chunk))
             current_chunk = [line]
@@ -94,11 +106,12 @@ def split_text_into_chunks(text: str) -> List[str]:
 
     if current_chunk:
         chunks.append("\n".join(current_chunk))
-    
+
     logger.info(f"Split text into {len(chunks)} chunks")
     return chunks
 
 def build_prompt_from_text(chunk: str) -> str:
+    """Build a prompt for analyzing a chunk of clinical data."""
     return f"""
 ### Unstructured Clinical Records
 
@@ -119,6 +132,7 @@ Please analyze the above and provide concise responses (max {MAX_NEW_TOKENS} tok
 """
 
 def init_agent():
+    """Initialize the TxAgent with optimized vLLM settings for A100 80GB."""
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
 
@@ -138,17 +152,19 @@ def init_agent():
     agent.init_model()
     return agent
 
-async def process_chunk(agent: TxAgent, chunk: str, chunk_idx: int) -> Tuple[int, str]:
-    """Process a single chunk with error handling"""
+async def process_chunk(agent, chunk: str, chunk_index: int, total_chunks: int) -> Tuple[int, str, str]:
+    """Process a single chunk and return index, response, and status message."""
+    logger.info(f"Processing chunk {chunk_index+1}/{total_chunks}")
+    prompt = build_prompt_from_text(chunk)
+    prompt_tokens = estimate_tokens(prompt)
+    
+    if prompt_tokens > MAX_MODEL_TOKENS:
+        error_msg = f"❌ Chunk {chunk_index+1} prompt too long ({prompt_tokens} tokens). Skipping..."
+        logger.warning(error_msg)
+        return chunk_index, "", error_msg
+
+    response = ""
     try:
-        prompt = build_prompt_from_text(chunk)
-        prompt_tokens = estimate_tokens(prompt)
-        
-        if prompt_tokens > MAX_MODEL_TOKENS:
-            logger.warning(f"Chunk {chunk_idx} prompt too long ({prompt_tokens} tokens)")
-            return chunk_idx, ""
-        
-        response = ""
         for result in agent.run_gradio_chat(
             message=prompt,
             history=[],
@@ -166,95 +182,143 @@ async def process_chunk(agent: TxAgent, chunk: str, chunk_idx: int) -> Tuple[int
                 for r in result:
                     if hasattr(r, "content"):
                         response += r.content
-        
-        return chunk_idx, clean_response(response)
-    
+        status = f"✅ Chunk {chunk_index+1} analysis complete"
+        logger.info(status)
     except Exception as e:
-        logger.error(f"Error processing chunk {chunk_idx}: {str(e)}")
-        return chunk_idx, ""
+        status = f"❌ Error analyzing chunk {chunk_index+1}: {str(e)}"
+        logger.error(status)
+        response = ""
+    
+    return chunk_index, clean_response(response), status
 
-async def process_file(agent: TxAgent, file_path: str) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
-    """Process the entire file and yield progress updates"""
-    messages = []
+async def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
+    """Process the Excel file and generate a final report."""
+    messages = chatbot_state if chatbot_state else []
     report_path = None
-    
+
+    if file is None or not hasattr(file, "name"):
+        messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
+        return messages, report_path
+
     try:
-        # Initial messages
-        messages.append({"role": "user", "content": f"Processing file: {os.path.basename(file_path)}"})
-        messages.append({"role": "assistant", "content": "⏳ Extracting data from Excel..."})
-        yield messages, None
-        
-        # Extract and chunk text
+        messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
+        messages.append({"role": "assistant", "content": "⏳ Extracting and analyzing data..."})
+
+        # Extract text and split into chunks
         start_time = time.time()
-        text = extract_text_from_excel(file_path)
-        chunks = split_text_into_chunks(text)
-        messages.append({"role": "assistant", "content": f"✅ Extracted {len(chunks)} chunks in {time.time()-start_time:.1f}s"})
-        yield messages, None
-        
-        # Process chunks in parallel
+        extracted_text = extract_text_from_excel(file.name)
+        chunks = split_text_into_chunks(extracted_text, max_tokens=MAX_CHUNK_TOKENS)
+        logger.info(f"Extracted text and split into {len(chunks)} chunks in {time.time() - start_time:.2f} seconds")
+
         chunk_responses = [None] * len(chunks)
-        with ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor:
-            futures = []
-            for idx, chunk in enumerate(chunks):
-                future = executor.submit(
-                    lambda c, i: asyncio.run(process_chunk(agent, c, i)),
-                    chunk, idx
-                )
-                futures.append(future)
-                messages.append({"role": "assistant", "content": f"🔍 Processing chunk {idx+1}/{len(chunks)}..."})
-                yield messages, None
-            
-            for future in as_completed(futures):
-                idx, response = future.result()
-                chunk_responses[idx] = response
-                messages.append({"role": "assistant", "content": f"✅ Chunk {idx+1} processed"})
-                yield messages, None
-        
-        # Combine and summarize
-        combined = "\n\n".join([r for r in chunk_responses if r])
+        batch_size = MAX_CONCURRENT
+
+        # Process chunks in batches
+        for batch_start in range(0, len(chunks), batch_size):
+            batch_chunks = chunks[batch_start:batch_start + batch_size]
+            batch_indices = list(range(batch_start, min(batch_start + batch_size, len(chunks))))
+            logger.info(f"Processing batch {batch_start//batch_size + 1}/{(len(chunks) + batch_size - 1)//batch_size}")
+
+            with ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor:
+                futures = [
+                    executor.submit(lambda c, i: asyncio.run(process_chunk(agent, c, i, len(chunks))), chunk, i)
+                    for i, chunk in zip(batch_indices, batch_chunks)
+                ]
+                for future in as_completed(futures):
+                    chunk_index, response, status = future.result()
+                    chunk_responses[chunk_index] = response
+                    messages.append({"role": "assistant", "content": status})
+
+        # Filter out empty responses
+        chunk_responses = [r for r in chunk_responses if r]
+        if not chunk_responses:
+            messages.append({"role": "assistant", "content": "❌ No valid chunk responses to summarize."})
+            return messages, report_path
+
+        # Summarize chunk responses incrementally
+        summary = ""
+        current_summary_tokens = 0
+        for i, response in enumerate(chunk_responses):
+            response_tokens = estimate_tokens(response)
+            if current_summary_tokens + response_tokens > MAX_MODEL_TOKENS - PROMPT_OVERHEAD - MAX_NEW_TOKENS:
+                summary_prompt = f"Summarize the following analysis:\n\n{summary}\n\nProvide a concise summary."
+                summary_response = ""
+                try:
+                    for result in agent.run_gradio_chat(
+                        message=summary_prompt,
+                        history=[],
+                        temperature=0.2,
+                        max_new_tokens=MAX_NEW_TOKENS,
+                        max_token=MAX_MODEL_TOKENS,
+                        call_agent=False,
+                        conversation=[],
+                    ):
+                        if isinstance(result, str):
+                            summary_response += result
+                        elif hasattr(result, "content"):
+                            summary_response += result.content
+                        elif isinstance(result, list):
+                            for r in result:
+                                if hasattr(r, "content"):
+                                    summary_response += r.content
+                    summary = clean_response(summary_response)
+                    current_summary_tokens = estimate_tokens(summary)
+                except Exception as e:
+                    messages.append({"role": "assistant", "content": f"❌ Error summarizing intermediate results: {str(e)}"})
+                    return messages, report_path
+
+            summary += f"\n\n### Chunk {i+1} Analysis\n{response}"
+            current_summary_tokens += response_tokens
+
+        # Final summarization
+        final_prompt = f"Summarize the key findings from the following analyses:\n\n{summary}"
         messages.append({"role": "assistant", "content": "📊 Generating final report..."})
-        yield messages, None
-        
-        final_response = ""
-        for result in agent.run_gradio_chat(
-            message=f"Summarize these clinical findings:\n\n{combined}",
-            history=[],
-            temperature=0.2,
-            max_new_tokens=MAX_NEW_TOKENS*2,  # Allow more tokens for summary
-            max_token=MAX_MODEL_TOKENS,
-            call_agent=False,
-            conversation=[],
-        ):
-            if isinstance(result, str):
-                final_response += result
-            elif hasattr(result, "content"):
-                final_response += result.content
-            elif isinstance(result, list):
-                for r in result:
-                    if hasattr(r, "content"):
-                        final_response += r.content
-            
-            messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(final_response)}"
-            yield messages, None
-        
-        # Save report
-        final_report = f"# Final Clinical Report\n\n{clean_response(final_response)}"
+
+        final_report_text = ""
+        try:
+            for result in agent.run_gradio_chat(
+                message=final_prompt,
+                history=[],
+                temperature=0.2,
+                max_new_tokens=MAX_NEW_TOKENS * 2,  # Allow more tokens for summary, as in your code
+                max_token=MAX_MODEL_TOKENS,
+                call_agent=False,
+                conversation=[],
+            ):
+                if isinstance(result, str):
+                    final_report_text += result
+                elif hasattr(result, "content"):
+                    final_report_text += result.content
+                elif isinstance(result, list):
+                    for r in result:
+                        if hasattr(r, "content"):
+                            final_report_text += r.content
+        except Exception as e:
+            messages.append({"role": "assistant", "content": f"❌ Error generating final report: {str(e)}"})
+            return messages, report_path
+
+        final_report = f"# Final Clinical Report\n\n{clean_response(final_report_text)}"
+        messages[-1]["content"] = f"📊 Final Report:\n\n{clean_response(final_report_text)}"
+
+        # Save the report
         timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
         report_path = os.path.join(report_dir, f"report_{timestamp}.md")
         
         with open(report_path, 'w') as f:
             f.write(final_report)
-        
-        messages.append({"role": "assistant", "content": f"✅ Report saved: report_{timestamp}.md"})
-        yield messages, report_path
-    
+
+        messages.append({"role": "assistant", "content": f"✅ Report generated and saved: report_{timestamp}.md"})
+        logger.info(f"Total processing time: {time.time() - start_time:.2f} seconds")
+
+        return messages, report_path
+
     except Exception as e:
+        messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})
         logger.error(f"Processing failed: {str(e)}")
-        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
-        yield messages, None
+        return messages, report_path
 
-def create_ui(agent: TxAgent):
-    """Create the Gradio interface"""
+def create_ui(agent):
+    """Create the Gradio interface."""
     with gr.Blocks(title="Clinical Analysis", css=".gradio-container {max-width: 900px}") as demo:
         gr.Markdown("## 🏥 Clinical Data Analysis (TxAgent)")
         
@@ -278,15 +342,29 @@ def create_ui(agent: TxAgent):
                 )
                 report_output = gr.File(
                     label="Download Report",
-                    visible=False
+                    visible=False,
+                    interactive=False
                 )
-        
+
+        # State to maintain chatbot messages
+        chatbot_state = gr.State(value=[])
+
+        async def update_ui(file, current_state):
+            if file is None or not hasattr(file, "name"):
+                messages = current_state if current_state else []
+                messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
+                return messages, None
+            messages, report_path = await process_final_report(agent, file, current_state)
+            report_update = gr.update(visible=report_path is not None, value=report_path)
+            return messages, report_update
+
         analyze_btn.click(
-            fn=lambda file: process_file(agent, file.name) if file else ([{"role": "assistant", "content": "❌ Please upload a file"}], None),
-            inputs=[file_input],
-            outputs=[chatbot, report_output]
+            fn=update_ui,
+            inputs=[file_input, chatbot_state],
+            outputs=[chatbot, report_output],
+            api_name="analyze"
         )
-    
+
     return demo
 
 if __name__ == "__main__":
@@ -298,7 +376,9 @@ if __name__ == "__main__":
             server_port=7860,
             show_error=True,
             allowed_paths=[report_dir],
-            share=False
+            share=False,
+            inline=False,
+            max_threads=40
         )
     except Exception as e:
         logger.error(f"Application failed: {str(e)}")