Update app.py

app.py

@@ -130,7 +130,8 @@ def init_agent():
         enable_checker=True,
         step_rag_num=8,
         seed=100,
-        additional_default_tools=[]
+        additional_default_tools=[],
+        device_map="auto"
     )
     agent.init_model()
     return agent
@@ -140,7 +141,7 @@ def create_ui(agent: TxAgent):
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
         gr.Markdown("<h3 style='text-align: center;'>Identify potential oversights in patient care</h3>")
 
-        chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
+        chatbot = gr.Chatbot(label="Analysis", height=600)
         file_upload = gr.File(
             label="Upload Medical Records",
             file_types=[".pdf", ".csv", ".xls", ".xlsx"],
@@ -149,23 +150,26 @@ def create_ui(agent: TxAgent):
         msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
         send_btn = gr.Button("Analyze", variant="primary")
         conversation_state = gr.State([])
-        download_output = gr.File(label="Download Full Report (after tools finish)")
+        download_output = gr.File(label="Download Full Report")
 
         def analyze_potential_oversights(message: str, history: list, conversation: list, files: list):
             start_time = time.time()
             try:
-                history.append({"role": "user", "content": message})
+                # Initialize conversation
+                history.append((message, "Analyzing records for potential oversights..."))
                 yield history, None
 
+                # Process files
                 extracted_data = ""
                 file_hash_value = ""
-                if files:
+                if files and isinstance(files, list):
                     with ThreadPoolExecutor(max_workers=4) as executor:
-                        futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files if hasattr(f, 'name')]
-                        results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
-                        extracted_data = "\n".join(results)
-                        file_hash_value = file_hash(files[0].name)
+                        futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) 
+                                 for f in files if hasattr(f, 'name')]
+                        extracted_data = "\n".join([sanitize_utf8(f.result()) for f in as_completed(futures)])
+                        file_hash_value = file_hash(files[0].name) if files else ""
 
+                # Medical oversight analysis prompt
                 analysis_prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
@@ -174,10 +178,16 @@ def create_ui(agent: TxAgent):
 
 Medical Records:\n{extracted_data[:15000]}
 
-### Potential Oversights:\n"""
+Provide ONLY the potential oversights in this format:
 
-                response = ""
-                for chunk in agent.run_gradio_chat(
+### Potential Oversights:
+1. [Missed diagnosis] - [Evidence from records]
+2. [Medication issue] - [Supporting data]
+3. [Assessment gap] - [Relevant findings]"""
+
+                # Generate and stream response
+                full_response = ""
+                generator = agent.run_gradio_chat(
                     message=analysis_prompt,
                     history=[],
                     temperature=0.2,
@@ -185,23 +195,45 @@ Medical Records:\n{extracted_data[:15000]}
                     max_token=4096,
                     call_agent=False,
                     conversation=conversation
-                ):
-                    partial = ""
-                    if isinstance(chunk, str):
-                        partial = chunk
-                    elif isinstance(chunk, list):
-                        partial = "".join([c.content for c in chunk if hasattr(c, 'content')])
-                    response += partial
-                    history[-1] = {"role": "assistant", "content": response.strip()}
-                    yield history, None
-
-                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                return history, report_path if os.path.exists(report_path) else None
+                )
+
+                for update in generator:
+                    if not update:
+                        continue
+                    
+                    if isinstance(update, str):
+                        full_response += update
+                    elif isinstance(update, list):
+                        full_response += "".join([msg.content for msg in update if hasattr(msg, 'content')])
+                    
+                    # Clean and update the response
+                    cleaned = full_response.replace("[TOOL_CALLS]", "").strip()
+                    if cleaned:
+                        history[-1] = (message, cleaned)
+                        yield history, None
+
+                # Final cleaned response
+                final_output = full_response.replace("[TOOL_CALLS]", "").strip()
+                if not final_output:
+                    final_output = "No clear oversights identified. Recommend comprehensive review."
+
+                # Prepare report path if available
+                report_path = None
+                if file_hash_value:
+                    possible_report = os.path.join(report_dir, f"{file_hash_value}_report.txt")
+                    if os.path.exists(possible_report):
+                        report_path = possible_report
+
+                history[-1] = (message, final_output)
+                print(f"Final analysis:\n{final_output}")
+                yield history, report_path
 
             except Exception as e:
-                history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
-                return history, None
+                print(f"Analysis error: {str(e)}")
+                history[-1] = (message, f"❌ Analysis failed: {str(e)}")
+                yield history, None
 
+        # UI event handlers
         inputs = [msg_input, chatbot, conversation_state, file_upload]
         outputs = [chatbot, download_output]
         send_btn.click(analyze_potential_oversights, inputs=inputs, outputs=outputs)
@@ -218,11 +250,28 @@ Medical Records:\n{extracted_data[:15000]}
 if __name__ == "__main__":
     print("Initializing medical analysis agent...")
     agent = init_agent()
+    
+    print("Performing warm-up call...")
+    try:
+        warm_up = agent.run_gradio_chat(
+            message="Warm up",
+            history=[],
+            temperature=0.1,
+            max_new_tokens=10,
+            max_token=100,
+            call_agent=False,
+            conversation=[]
+        )
+        for _ in warm_up:
+            pass
+    except Exception as e:
+        print(f"Warm-up error: {str(e)}")
+
     print("Launching interface...")
     demo = create_ui(agent)
-    demo.queue().launch(
+    demo.queue(concurrency_count=2).launch(
         server_name="0.0.0.0",
         server_port=7860,
         show_error=True,
         share=True
-    )
+    )