Update app.py

app.py

@@ -2,12 +2,11 @@ import sys
 import os
 import pandas as pd
 import gradio as gr
-from typing import List, Tuple, Dict, Any, Union
+from typing import List, Tuple, Dict, Any, Union, Generator
 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
@@ -33,22 +32,15 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Constants
+# Updated token limits as specified
 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:
@@ -62,13 +54,9 @@ def clean_response(text: str) -> str:
     return text.strip()
 
 def estimate_tokens(text: str) -> int:
-    """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
+    return len(text) // 3.5 + 1  # More conservative estimate
 
 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)
@@ -79,15 +67,16 @@ def extract_text_from_excel(file_path: str) -> str:
             sheet_text = [f"[{sheet_name}] {line}" for line in rows]
             all_text.extend(sheet_text)
     except Exception as e:
-        raise ValueError(f"Failed to extract text from Excel file: {str(e)}")
+        logger.error(f"Error extracting Excel: {str(e)}")
+        raise ValueError(f"Failed to process Excel file: {str(e)}")
     return "\n".join(all_text)
 
-def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
-    """Split text into chunks within token limits, accounting for 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.")
-
+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")
+    
     lines = text.split("\n")
     chunks = []
     current_chunk = []
@@ -95,7 +84,7 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> Lis
 
     for line in lines:
         line_tokens = estimate_tokens(line)
-        if current_tokens + line_tokens > effective_max_tokens:
+        if current_tokens + line_tokens > effective_max:
             if current_chunk:
                 chunks.append("\n".join(current_chunk))
             current_chunk = [line]
@@ -106,11 +95,11 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> Lis
 
     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
 
@@ -122,7 +111,7 @@ Here is the extracted content chunk:
 
 {chunk}
 
-Please analyze the above and provide:
+Please analyze the above and provide concise responses (max {MAX_NEW_TOKENS} tokens):
 - Diagnostic Patterns
 - Medication Issues
 - Missed Opportunities
@@ -131,7 +120,6 @@ Please analyze the above and provide:
 """
 
 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")
 
@@ -146,24 +134,23 @@ def init_agent():
         enable_checker=True,
         step_rag_num=4,
         seed=100,
-        additional_default_tools=[]
+        additional_default_tools=[],
+        max_model_tokens=MAX_MODEL_TOKENS  # Pass the updated token limit
     )
     agent.init_model()
     return agent
 
-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 = ""
+async def process_chunk(agent: TxAgent, chunk: str, chunk_idx: int) -> Tuple[int, str]:
+    """Process a single chunk with error handling"""
     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=[],
@@ -181,205 +168,139 @@ async def process_chunk(agent, chunk: str, chunk_index: int, total_chunks: int)
                 for r in result:
                     if hasattr(r, "content"):
                         response += r.content
-        status = f"✅ Chunk {chunk_index+1} analysis complete"
-        logger.info(status)
-    except Exception as e:
-        status = f"❌ Error analyzing chunk {chunk_index+1}: {str(e)}"
-        logger.error(status)
-        response = ""
+        
+        return chunk_idx, clean_response(response)
     
-    return chunk_index, clean_response(response), status
+    except Exception as e:
+        logger.error(f"Error processing chunk {chunk_idx}: {str(e)}")
+        return chunk_idx, ""
 
-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 []
+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 = []
     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:
-        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
+        # 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
         start_time = time.time()
-        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")
-
+        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
         chunk_responses = [None] * len(chunks)
-        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}"
+        with ThreadPoolExecutor(max_workers=MAX_CONCURRENT) as executor:
+            futures = []
+            for idx, chunk in enumerate(chunks):
+                futures.append(executor.submit(
+                    lambda c, i: asyncio.run(process_chunk(agent, c, i)),
+                    chunk, idx
+                )
+                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])
         messages.append({"role": "assistant", "content": "📊 Generating final report..."})
-
-        final_report_text = ""
-        try:
-            for result in agent.run_gradio_chat(
-                message=final_prompt,
-                history=[],
-                temperature=0.2,
-                max_new_tokens=MAX_NEW_TOKENS,
-                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"# \U0001f9e0 Final Patient Report\n\n{clean_response(final_report_text)}"
-        messages[-1]["content"] = f"📊 Final Report:\n\n{clean_response(final_report_text)}"
-
-        # Save the 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)}"
         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 generated and saved: report_{timestamp}.md"})
-        logger.info(f"Total processing time: {time.time() - start_time:.2f} seconds")
-
-        return messages, report_path
-
+        
+        messages.append({"role": "assistant", "content": f"✅ Report saved: report_{timestamp}.md"})
+        yield 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)}")
-        return messages, report_path
+        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
+        yield messages, None
 
-async def create_ui(agent):
-    """Create the Gradio UI for the patient history analysis tool."""
-    with gr.Blocks(title="Patient History Chat", css=".gradio-container {max-width: 900px !important}") as demo:
-        gr.Markdown("## 🏥 Patient History Analysis Tool")
+def create_ui(agent: TxAgent):
+    """Create the Gradio interface"""
+    with gr.Blocks(title="Clinical Analysis", css=".gradio-container {max-width: 900px}") as demo:
+        gr.Markdown("## 🏥 Clinical Data Analysis (TxAgent)")
         
         with gr.Row():
             with gr.Column(scale=3):
                 chatbot = gr.Chatbot(
-                    label="Clinical Assistant",
+                    label="Analysis Progress",
                     show_copy_button=True,
                     height=600,
-                    type="messages",
-                    avatar_images=(
-                        None,
-                        "https://i.imgur.com/6wX7Zb4.png"
-                    )
+                    type="messages"
                 )
             with gr.Column(scale=1):
-                file_upload = gr.File(
+                file_input = gr.File(
                     label="Upload Excel File",
                     file_types=[".xlsx"],
                     height=100
                 )
                 analyze_btn = gr.Button(
-                    "🧠 Analyze Patient History",
+                    "🧠 Analyze Data",
                     variant="primary"
                 )
                 report_output = gr.File(
                     label="Download Report",
-                    visible=False,
-                    interactive=False
+                    visible=False
                 )
-
-        # State to maintain chatbot messages
-        chatbot_state = gr.State(value=[])
-
-        async def update_ui(file, current_state):
-            messages = current_state if current_state else []
-            messages, report_path = await process_final_report(agent, file, messages)
-            report_update = gr.update(visible=report_path is not None, value=report_path)
-            return messages, report_update, messages
-
+        
         analyze_btn.click(
-            fn=update_ui,
-            inputs=[file_upload, chatbot_state],
-            outputs=[chatbot, report_output, chatbot_state],
-            api_name="analyze"
+            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]
         )
-
+    
     return demo
 
 if __name__ == "__main__":
     try:
         agent = init_agent()
-        demo = asyncio.run(create_ui(agent))
+        demo = create_ui(agent)
         demo.launch(
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            allowed_paths=["/data/hf_cache/reports"],
-            share=False,
-            inline=False,
-            max_threads=40
+            allowed_paths=[report_dir],
+            share=False
         )
     except Exception as e:
-        print(f"Error: {str(e)}")
+        logger.error(f"Application failed: {str(e)}")
         sys.exit(1)