Update app.py

app.py

@@ -1,11 +1,8 @@
 import sys
 import os
 import pandas as pd
-import json
 import gradio as gr
-from typing import List, Tuple, Dict, Any, Union
-import hashlib
-import shutil
+from typing import List, Tuple
 import re
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -16,16 +13,15 @@ os.makedirs(persistent_dir, exist_ok=True)
 
 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")
 
-for d in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
+for d in [model_cache_dir, tool_cache_dir, report_dir]:
     os.makedirs(d, exist_ok=True)
 
 os.environ["HF_HOME"] = model_cache_dir
 os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src"))
 from txagent.txagent import TxAgent
 
 MAX_MODEL_TOKENS = 32768
@@ -39,9 +35,6 @@ def clean_response(text: str) -> str:
     text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
     return text.strip()
 
-def estimate_tokens(text: str) -> int:
-    return len(text) // 3.5 + 1
-
 def extract_text_from_excel(file_path: str) -> str:
     all_text = []
     xls = pd.ExcelFile(file_path)
@@ -52,47 +45,41 @@ def extract_text_from_excel(file_path: str) -> str:
         all_text.extend(sheet_text)
     return "\n".join(all_text)
 
-def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS, max_chunks: int = 30) -> List[str]:
-    effective_max = max_tokens - PROMPT_OVERHEAD
-    lines, chunks, curr_chunk, curr_tokens = text.split("\n"), [], [], 0
+def split_text_into_chunks(text: str) -> List[str]:
+    effective_max = MAX_CHUNK_TOKENS - PROMPT_OVERHEAD
+    lines, chunks, curr_chunk = text.split("\n"), [], []
+    curr_tokens = sum(len(line.split()) for line in curr_chunk)
+    
     for line in lines:
-        t = estimate_tokens(line)
-        if curr_tokens + t > effective_max:
+        line_tokens = len(line.split())
+        if curr_tokens + line_tokens > effective_max:
             if curr_chunk:
                 chunks.append("\n".join(curr_chunk))
-            if len(chunks) >= max_chunks:
-                break
-            curr_chunk, curr_tokens = [line], t
+            curr_chunk, curr_tokens = [line], line_tokens
         else:
             curr_chunk.append(line)
-            curr_tokens += t
-    if curr_chunk and len(chunks) < max_chunks:
+            curr_tokens += line_tokens
+    if curr_chunk:
         chunks.append("\n".join(curr_chunk))
     return chunks
 
 def build_prompt_from_text(chunk: str) -> str:
-    return f"""
-### Unstructured Clinical Records
-
-Analyze the following clinical notes and provide a detailed, concise summary focusing on:
-- Diagnostic Patterns
-- Medication Issues
-- Missed Opportunities
+    return f"""Analyze these clinical notes and provide:
+- Diagnostic patterns
+- Medication issues 
+- Missed opportunities
 - Inconsistencies
-- Follow-up Recommendations
+- Follow-up recommendations
 
----
+Respond with clear bullet points:
 
-{chunk}
-
----
-Respond in well-structured bullet points with medical reasoning.
-"""
+{chunk}"""
 
 def init_agent():
     tool_path = os.path.join(tool_cache_dir, "new_tool.json")
     if not os.path.exists(tool_path):
-        shutil.copy(os.path.abspath("data/new_tool.json"), tool_path)
+        import shutil
+        shutil.copy("data/new_tool.json", 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",
@@ -105,68 +92,62 @@ def init_agent():
     agent.init_model()
     return agent
 
-def process_final_report(agent, file, chatbot_state: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], Union[str, None, str]]:
-    messages = chatbot_state if chatbot_state else []
-    if file is None or not hasattr(file, "name"):
-        return messages + [("assistant", "❌ Please upload a valid Excel file.")], None, ""
+def process_final_report(agent, file, chatbot_state: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], str]:
+    messages = chatbot_state.copy() if chatbot_state else []
+    
+    if file is None:
+        messages.append(("assistant", "❌ Please upload a valid Excel file."))
+        return messages, None
 
     messages.append(("user", f"Processing Excel file: {os.path.basename(file.name)}"))
-    text = extract_text_from_excel(file.name)
-    chunks = split_text_into_chunks(text)
-    chunk_responses = [None] * len(chunks)
-
-    def analyze_chunk(i, chunk):
-        prompt = build_prompt_from_text(chunk)
-        response = ""
-        for res 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(res, str):
-                response += res
-            elif hasattr(res, "content"):
-                response += res.content
-            elif isinstance(res, list):
-                for r in res:
-                    if hasattr(r, "content"):
-                        response += r.content
-        return i, clean_response(response)
-
-    with ThreadPoolExecutor(max_workers=1) as executor:
-        futures = [executor.submit(analyze_chunk, i, c) for i, c in enumerate(chunks)]
-        for f in as_completed(futures):
-            i, result = f.result()
-            chunk_responses[i] = result
-
-    valid = [r for r in chunk_responses if r and not r.startswith("❌")]
-    if not valid:
-        return messages + [("assistant", "❌ No valid chunk results.")], None, ""
-
-    summary_prompt = f"Summarize this analysis in a final structured report:\n\n" + "\n\n".join(valid)
-    messages.append(("assistant", "📊 Generating final report..."))
-
-    final_report = ""
-    for res 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(res, str):
-            final_report += res
-        elif hasattr(res, "content"):
-            final_report += res.content
-
-    cleaned = clean_response(final_report)
-    report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
-    with open(report_path, 'w') as f:
-        f.write(f"# 🧠 Final Patient Report\n\n{cleaned}")
-
-    # Add the report content to the chat messages
-    messages.append(("assistant", f"✅ Report generated and saved: {os.path.basename(report_path)}"))
-    messages.append(("assistant", f"## Final Report\n\n{cleaned}"))
+    yield messages, None
     
-    return messages, report_path, cleaned
+    try:
+        text = extract_text_from_excel(file.name)
+        chunks = split_text_into_chunks(text)
+        
+        messages.append(("assistant", "🔍 Analyzing clinical data..."))
+        yield messages, None
+        
+        full_report = []
+        for i, chunk in enumerate(chunks, 1):
+            prompt = build_prompt_from_text(chunk)
+            response = ""
+            
+            for res 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(res, str):
+                    response += res
+                elif hasattr(res, "content"):
+                    response += res.content
+            
+            cleaned = clean_response(response)
+            full_report.append(cleaned)
+            
+            # Update progress in chat
+            progress_msg = f"✅ Analyzed section {i}/{len(chunks)}"
+            if len(messages) > 2 and "Analyzed section" in messages[-1][1]:
+                messages[-1] = ("assistant", progress_msg)
+            else:
+                messages.append(("assistant", progress_msg))
+            yield messages, None
+        
+        # Generate final report
+        final_report = "## 🧠 Final Clinical Report\n\n" + "\n\n".join(full_report)
+        report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
+        with open(report_path, 'w') as f:
+            f.write(final_report)
+        
+        messages.append(("assistant", f"✅ Report generated and saved: {os.path.basename(report_path)}"))
+        messages.append(("assistant", final_report))
+        yield messages, report_path
+        
+    except Exception as e:
+        messages.append(("assistant", f"❌ Error: {str(e)}"))
+        yield messages, None
 
 def create_ui(agent):
     with gr.Blocks(css="""
@@ -211,24 +192,25 @@ def create_ui(agent):
             margin-top: 10px;
             border: 1px solid #2c3344;
         }
+        .bullet-points {
+            margin-left: 20px;
+        }
     """) as demo:
-        gr.Markdown("""# 🧠 Clinical Reasoning Assistant
+        gr.Markdown("""# Clinical Reasoning Assistant
 Upload clinical Excel records below and click **Analyze** to generate a medical summary.
 """)
-        chatbot = gr.Chatbot(label="Chatbot", elem_classes="chatbot", type="tuples")
+        
+        chatbot = gr.Chatbot(label="Chatbot", elem_classes="chatbot")
         file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
         analyze_btn = gr.Button("Analyze")
         report_output = gr.File(label="Download Report", visible=False)
-        chatbot_state = gr.State(value=[])
-
-        def update_ui(file, current_state):
-            messages, report_path, final_text = process_final_report(agent, file, current_state)
-            return messages, gr.update(visible=report_path is not None, value=report_path), messages
-
+        chatbot_state = gr.State([])
+        
         analyze_btn.click(
-            fn=update_ui,
-            inputs=[file_upload, chatbot_state],
-            outputs=[chatbot, report_output, chatbot_state]
+            fn=process_final_report,
+            inputs=[file_upload, chatbot_state, gr.State(agent)],
+            outputs=[chatbot, report_output],
+            show_progress="hidden"
         )
 
     return demo
@@ -237,7 +219,12 @@ if __name__ == "__main__":
     try:
         agent = init_agent()
         demo = create_ui(agent)
-        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)
+        demo.launch(
+            server_name="0.0.0.0", 
+            server_port=7860, 
+            allowed_paths=["/data/hf_cache/reports"], 
+            share=False
+        )
     except Exception as e:
         print(f"Error: {str(e)}")
         sys.exit(1)