Update app.py

app.py

@@ -8,11 +8,10 @@ from importlib.resources import files
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-tx_app = None  # Will be initialized later in on_start
+tx_app = None  # Global to hold the model
 
 def init_txagent():
-    """Initialize the TxAgent with proper tool file paths"""
-    logger.info("Initializing TxAgent...")
+    logger.info("🔥 Initializing TxAgent...")
 
     tool_files = {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
@@ -21,8 +20,6 @@ def init_txagent():
         "monarch": str(files('tooluniverse.data').joinpath('monarch_tools.json'))
     }
 
-    logger.info(f"Using tool files at: {tool_files}")
-
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -44,7 +41,7 @@ def init_txagent():
     )
 
     agent.init_model()
-    logger.info("Model loading complete")
+    logger.info("✅ Model loading complete")
     return agent
 
 def respond(message, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round):
@@ -81,11 +78,11 @@ def respond(message, chat_history, temperature, max_new_tokens, max_tokens, mult
         logger.error(f"Error in respond function: {str(e)}")
         yield chat_history + [("", f"⚠️ Error: {str(e)}")]
 
-# Define Gradio UI
+# Gradio UI
 with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
     gr.Markdown("# 🧠 TxAgent Biomedical Assistant")
 
-    chatbot = gr.Chatbot(label="Conversation", height=600)
+    chatbot = gr.Chatbot(label="Conversation", height=600, type="messages")
 
     msg = gr.Textbox(
         label="Your medical query",
@@ -119,8 +116,13 @@ with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
         chatbot
     )
 
-    @app.on_start
+    # Hidden button + load trigger to init model safely
+    init_button = gr.Button(visible=False)
+
     def load_model():
         global tx_app
-        logger.info("🔥 Loading TxAgent model in Gradio @on_start")
         tx_app = init_txagent()
+        return gr.update(visible=False)
+
+    app.load(init_button.click(fn=load_model))
+