Update app.py

app.py

@@ -1,15 +1,14 @@
 import sys
 import os
 import pandas as pd
+import json
 import gradio as gr
-from typing import List, Tuple, Dict, Any, Union, Generator
+from typing import List, Tuple, Dict, Any, Union
+import hashlib
 import shutil
 import re
 from datetime import datetime
 import time
-import asyncio
-import logging
-from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # Configuration and setup
 persistent_dir = "/data/hf_cache"
@@ -24,6 +23,7 @@ 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
+os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 src_path = os.path.abspath(os.path.join(current_dir, "src"))
@@ -32,15 +32,10 @@ sys.path.insert(0, src_path)
 from txagent.txagent import TxAgent
 
 # 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 = 4         # Reduced concurrency to avoid vLLM issues
-
-# Setup logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-logger = logging.getLogger(__name__)
+MAX_MODEL_TOKENS = 32768  # Model's maximum sequence length
+MAX_CHUNK_TOKENS = 8192   # Chunk size aligned with max_num_batched_tokens
+MAX_NEW_TOKENS = 2048     # Maximum tokens for generation
+PROMPT_OVERHEAD = 500     # Estimated tokens for prompt template overhead
 
 def clean_response(text: str) -> str:
     try:
@@ -53,9 +48,11 @@ def clean_response(text: str) -> str:
     return text.strip()
 
 def estimate_tokens(text: str) -> int:
-    return len(text) // 3.5 + 1  # Conservative estimate
+    """Estimate the number of tokens based on character length."""
+    return len(text) // 3.5 + 1  # Add 1 to avoid zero estimates
 
 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)
@@ -66,16 +63,18 @@ 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:
-        logger.error(f"Error extracting Excel: {str(e)}")
-        raise ValueError(f"Failed to process Excel file: {str(e)}")
+        raise ValueError(f"Failed to extract text from 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, ensuring each chunk is 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.")
+
     lines = text.split("\n")
     chunks = []
     current_chunk = []
@@ -83,8 +82,8 @@ 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_chunk:
+        if current_tokens + line_tokens > effective_max_tokens:
+            if current_chunk:  # Save the current chunk if it's not empty
                 chunks.append("\n".join(current_chunk))
             current_chunk = [line]
             current_tokens = line_tokens
@@ -94,11 +93,11 @@ 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
 
@@ -110,7 +109,7 @@ Here is the extracted content chunk:
 
 {chunk}
 
-Please analyze the above and provide concise responses (max {MAX_NEW_TOKENS} tokens):
+Please analyze the above and provide:
 - Diagnostic Patterns
 - Medication Issues
 - Missed Opportunities
@@ -119,7 +118,7 @@ Please analyze the above and provide concise responses (max {MAX_NEW_TOKENS} tok
 """
 
 def init_agent():
-    """Initialize TxAgent with conservative settings to avoid vLLM issues"""
+    """Initialize the TxAgent with model and tool configurations."""
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
 
@@ -139,168 +138,205 @@ def init_agent():
     agent.init_model()
     return agent
 
-def process_chunk_sync(agent, chunk: str, chunk_idx: int) -> Tuple[int, str]:
-    """Synchronous wrapper for chunk processing"""
-    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=[],
-            temperature=0.2,
-            max_new_tokens=MAX_NEW_TOKENS,
-            max_token=MAX_MODEL_TOKENS,
-            call_agent=False,
-            conversation=[],
-        ):
-            if isinstance(result, str):
-                response += result
-            elif hasattr(result, "content"):
-                response += result.content
-            elif isinstance(result, list):
-                for r in result:
-                    if hasattr(r, "content"):
-                        response += r.content
-        
-        return chunk_idx, clean_response(response)
-    except Exception as e:
-        logger.error(f"Error processing chunk {chunk_idx}: {str(e)}")
-        return chunk_idx, ""
-
-async def process_file(agent: TxAgent, file_path: str) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
-    messages = []
+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_path is None:
+
+    if file is None or not hasattr(file, "name"):
         messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
-        yield messages, None
-        return
+        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
-        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 sequentially
+        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
+        extracted_text = extract_text_from_excel(file.name)
+        chunks = split_text_into_chunks(extracted_text, max_tokens=MAX_CHUNK_TOKENS)
         chunk_responses = []
-        for idx, chunk in enumerate(chunks):
-            messages.append({"role": "assistant", "content": f"🔍 Processing chunk {idx+1}/{len(chunks)}..."})
-            yield messages, None
+
+        # Process each chunk
+        for i, chunk in enumerate(chunks):
+            messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
             
-            _, response = await process_chunk(agent, chunk, idx)
-            chunk_responses.append(response)
+            prompt = build_prompt_from_text(chunk)
+            prompt_tokens = estimate_tokens(prompt)
+            if prompt_tokens > MAX_MODEL_TOKENS:
+                messages.append({"role": "assistant", "content": f"❌ Chunk {i+1} prompt too long ({prompt_tokens} tokens). Skipping..."})
+                continue
+
+            response = ""
+            try:
+                for result in agent.run_gradio_chat(
+                    message=prompt,
+                    history=[],
+                    temperature=0.2,
+                    max_new_tokens=MAX_NEW_TOKENS,
+                    max_token=MAX_MODEL_TOKENS,
+                    call_agent=False,
+                    conversation=[],
+                ):
+                    if isinstance(result, str):
+                        response += result
+                    elif hasattr(result, "content"):
+                        response += result.content
+                    elif isinstance(result, list):
+                        for r in result:
+                            if hasattr(r, "content"):
+                                response += r.content
+            except Exception as e:
+                messages.append({"role": "assistant", "content": f"❌ Error analyzing chunk {i+1}: {str(e)}"})
+                continue
             
-            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])
+            chunk_responses.append(clean_response(response))
+            messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})
+
+        if not chunk_responses:
+            messages.append({"role": "assistant", "content": "❌ No valid chunk responses to summarize."})
+            return messages, report_path
+
+        # Summarize chunk responses incrementally to avoid token limit
+        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:
+                # Summarize current summary
+                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,
-            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,
+                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
         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"})
+
     except Exception as e:
-        logger.error(f"Processing failed: {str(e)}")
-        messages.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
-        yield messages, None
-def create_ui(agent: TxAgent):
-    """Create the Gradio interface with simplified interaction"""
-    with gr.Blocks(title="Clinical Analysis", css=".gradio-container {max-width: 900px}") as demo:
-        gr.Markdown("## 🏥 Clinical Data Analysis (TxAgent)")
+        messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})
+
+    return messages, report_path
+
+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")
         
         with gr.Row():
             with gr.Column(scale=3):
                 chatbot = gr.Chatbot(
-                    label="Analysis Progress",
+                    label="Clinical Assistant",
                     show_copy_button=True,
                     height=600,
-                    type="messages"
+                    type="messages",
+                    avatar_images=(
+                        None,
+                        "https://i.imgur.com/6wX7Zb4.png"
+                    )
                 )
             with gr.Column(scale=1):
-                file_input = gr.File(
+                file_upload = gr.File(
                     label="Upload Excel File",
                     file_types=[".xlsx"],
                     height=100
                 )
                 analyze_btn = gr.Button(
-                    "🧠 Analyze Data",
+                    "🧠 Analyze Patient History",
                     variant="primary"
                 )
                 report_output = gr.File(
                     label="Download Report",
-                    visible=False
+                    visible=False,
+                    interactive=False
                 )
-        
+
+        # State to maintain chatbot messages
+        chatbot_state = gr.State(value=[])
+
+        def update_ui(file, current_state):
+            messages, report_path = process_final_report(agent, file, current_state)
+            report_update = gr.update(visible=report_path is not None, value=report_path)
+            return messages, report_update, messages
+
         analyze_btn.click(
-            fn=lambda file: process_file(agent, file.name if file else None),
-            inputs=[file_input],
-            outputs=[chatbot, report_output],
-            concurrency_limit=1  # Ensure sequential processing
+            fn=update_ui,
+            inputs=[file_upload, chatbot_state],
+            outputs=[chatbot, report_output, chatbot_state],
+            api_name="analyze"
         )
-    
+
     return demo
 
 if __name__ == "__main__":
     try:
-        # Initialize with conservative settings
         agent = init_agent()
         demo = create_ui(agent)
-        
-        # Launch with stability optimizations
         demo.launch(
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            allowed_paths=[report_dir],
-            share=False,
-            max_threads=4  # Reduced thread count for stability
+            allowed_paths=["/data/hf_cache/reports"],
+            share=False
         )
     except Exception as e:
-        logger.error(f"Application failed: {str(e)}")
+        print(f"Error: {str(e)}")
         sys.exit(1)