Update app.py

app.py

@@ -5,7 +5,6 @@ from txagent import TxAgent
 import spaces
 import gradio as gr
 import os
-import os
 
 # Determine the directory where the current file is located
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -14,7 +13,6 @@ os.environ["MKL_THREADING_LAYER"] = "GNU"
 # Set an environment variable
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-
 DESCRIPTION = '''
 <div>
 <h1 style="text-align: center;">TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools </h1>
@@ -36,9 +34,8 @@ PLACEHOLDER = """
 <div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
    <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">TxAgent</h1>
    <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Tips before using TxAgent:</p>
-   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Please click clear🗑️
- (top-right) to remove previous context before sumbmitting a new question.</p>
-   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Click retry🔄 (below message)  to get multiple versions of the answer.</p>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Please click clear🗑️ (top-right) to remove previous context before submitting a new question.</p>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.55;">Click retry🔄 (below message) to get multiple versions of the answer.</p>
 </div>
 """
 
@@ -71,13 +68,11 @@ chat_css = """
 .gr-button svg { width: 32px !important; height: 32px !important; } /* Enlarges SVG icons */
 """
 
-# model_name = '/n/holylfs06/LABS/mzitnik_lab/Lab/shgao/bioagent/bio/alignment-handbook/data_new/L8-qlora-biov49v9v7v16_32k_chat01_merged'
+# Configuration variables (safe to keep at module level)
 model_name = 'mims-harvard/TxAgent-T1-Llama-3.1-8B'
 rag_model_name = 'mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B'
-
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-
 question_examples = [
     ['Given a 50-year-old patient experiencing severe acute pain and considering the use of the newly approved medication, Journavx, how should the dosage be adjusted considering the presence of moderate hepatic impairment?'],
     ['Given a 50-year-old patient experiencing severe acute pain and considering the use of the newly approved medication, Journavx, how should the dosage be adjusted considering the presence of severe hepatic impairment?'],
@@ -88,20 +83,9 @@ new_tool_files = {
     'new_tool': os.path.join(current_dir, 'data', 'new_tool.json'),
 }
 
-agent = TxAgent(model_name,
-                rag_model_name,
-                tool_files_dict=new_tool_files,
-                force_finish=True,
-                enable_checker=True,
-                step_rag_num=10,
-                seed=100,
-                additional_default_tools=['DirectResponse', 'RequireClarification'])
-agent.init_model()
-
-
 def update_model_parameters(enable_finish, enable_rag, enable_summary,
-                            init_rag_num, step_rag_num, skip_last_k,
-                            summary_mode, summary_skip_last_k, summary_context_length, force_finish, seed):
+                          init_rag_num, step_rag_num, skip_last_k,
+                          summary_mode, summary_skip_last_k, summary_context_length, force_finish, seed):
     # Update model instance parameters dynamically
     updated_params = agent.update_parameters(
         enable_finish=enable_finish,
@@ -116,10 +100,8 @@ def update_model_parameters(enable_finish, enable_rag, enable_summary,
         force_finish=force_finish,
         seed=seed,
     )
-
     return updated_params
 
-
 def update_seed():
     # Update model instance parameters dynamically
     seed = random.randint(0, 10000)
@@ -128,166 +110,166 @@ def update_seed():
     )
     return updated_params
 
-
 def handle_retry(history, retry_data: gr.RetryData, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
     print("Updated seed:", update_seed())
     new_history = history[:retry_data.index]
     previous_prompt = history[retry_data.index]['content']
-
     print("previous_prompt", previous_prompt)
-
     yield from agent.run_gradio_chat(new_history + [{"role": "user", "content": previous_prompt}], temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round)
 
-
 PASSWORD = "mypassword"
 
-# Function to check if the password is correct
-
-
 def check_password(input_password):
     if input_password == PASSWORD:
         return gr.update(visible=True), ""
     else:
         return gr.update(visible=False), "Incorrect password, try again!"
 
-
-conversation_state = gr.State([])
-
-# Gradio block
-chatbot = gr.Chatbot(height=800, placeholder=PLACEHOLDER,
-                     label='TxAgent', type="messages",  show_copy_button=True)
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.Markdown(INTRO)
-    default_temperature = 0.3
-    default_max_new_tokens = 1024
-    default_max_tokens = 81920
-    default_max_round = 30
-    temperature_state = gr.State(value=default_temperature)
-    max_new_tokens_state = gr.State(value=default_max_new_tokens)
-    max_tokens_state = gr.State(value=default_max_tokens)
-    max_round_state = gr.State(value=default_max_round)
-    chatbot.retry(handle_retry, chatbot, chatbot, temperature_state, max_new_tokens_state,
-                  max_tokens_state, gr.Checkbox(value=False, render=False), conversation_state, max_round_state)
-
-    gr.ChatInterface(
-        fn=agent.run_gradio_chat,
-        chatbot=chatbot,
-        fill_height=True, fill_width=True, stop_btn=True,
-        additional_inputs_accordion=gr.Accordion(
-            label="⚙️ Inference Parameters", open=False, render=False),
-        additional_inputs=[
-            temperature_state, max_new_tokens_state, max_tokens_state,
-            gr.Checkbox(
-                label="Activate multi-agent reasoning mode (it requires additional time but offers a more comprehensive analysis).", value=False, render=False),
-            conversation_state,
-            max_round_state,
-            gr.Number(label="Seed", value=100, render=False)
-        ],
-        examples=question_examples,
-        cache_examples=False,
-        css=chat_css,
-    )
-
-    with gr.Accordion("Settings", open=False):
-
-        # Define the sliders
-        temperature_slider = gr.Slider(
-            minimum=0,
-            maximum=1,
-            step=0.1,
-            value=default_temperature,
-            label="Temperature"
-        )
-        max_new_tokens_slider = gr.Slider(
-            minimum=128,
-            maximum=4096,
-            step=1,
-            value=default_max_new_tokens,
-            label="Max new tokens"
-        )
-        max_tokens_slider = gr.Slider(
-            minimum=128,
-            maximum=32000,
-            step=1,
-            value=default_max_tokens,
-            label="Max tokens"
+# Create the Gradio interface
+def create_interface(agent):
+    conversation_state = gr.State([])
+    chatbot = gr.Chatbot(height=800, placeholder=PLACEHOLDER,
+                        label='TxAgent', type="messages", show_copy_button=True)
+
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown(DESCRIPTION)
+        gr.Markdown(INTRO)
+        
+        default_temperature = 0.3
+        default_max_new_tokens = 1024
+        default_max_tokens = 81920
+        default_max_round = 30
+        
+        temperature_state = gr.State(value=default_temperature)
+        max_new_tokens_state = gr.State(value=default_max_new_tokens)
+        max_tokens_state = gr.State(value=default_max_tokens)
+        max_round_state = gr.State(value=default_max_round)
+        
+        chatbot.retry(handle_retry, chatbot, chatbot, temperature_state, max_new_tokens_state,
+                     max_tokens_state, gr.Checkbox(value=False, render=False), conversation_state, max_round_state)
+
+        gr.ChatInterface(
+            fn=agent.run_gradio_chat,
+            chatbot=chatbot,
+            fill_height=True, fill_width=True, stop_btn=True,
+            additional_inputs_accordion=gr.Accordion(
+                label="⚙️ Inference Parameters", open=False, render=False),
+            additional_inputs=[
+                temperature_state, max_new_tokens_state, max_tokens_state,
+                gr.Checkbox(
+                    label="Activate multi-agent reasoning mode (it requires additional time but offers a more comprehensive analysis).", 
+                    value=False, render=False),
+                conversation_state,
+                max_round_state,
+                gr.Number(label="Seed", value=100, render=False)
+            ],
+            examples=question_examples,
+            cache_examples=False,
+            css=chat_css,
         )
-        max_round_slider = gr.Slider(
-            minimum=0,
-            maximum=50,
-            step=1,
-            value=default_max_round,
-            label="Max round")
-
-        # Automatically update states when slider values change
-        temperature_slider.change(
-            lambda x: x, inputs=temperature_slider, outputs=temperature_state)
-        max_new_tokens_slider.change(
-            lambda x: x, inputs=max_new_tokens_slider, outputs=max_new_tokens_state)
-        max_tokens_slider.change(
-            lambda x: x, inputs=max_tokens_slider, outputs=max_tokens_state)
-        max_round_slider.change(
-            lambda x: x, inputs=max_round_slider, outputs=max_round_state)
-
-        password_input = gr.Textbox(
-            label="Enter Password for More Settings", type="password")
-        incorrect_message = gr.Textbox(visible=False, interactive=False)
-        with gr.Accordion("⚙️ Settings", open=False, visible=False) as protected_accordion:
-            with gr.Row():
-                with gr.Column(scale=1):
-                    with gr.Accordion("⚙️ Model Loading", open=False):
-                        model_name_input = gr.Textbox(
-                            label="Enter model path", value=model_name)
-                        load_model_btn = gr.Button(value="Load Model")
-                        load_model_btn.click(
-                            agent.load_models, inputs=model_name_input, outputs=gr.Textbox(label="Status"))
-                with gr.Column(scale=1):
-                    with gr.Accordion("⚙️ Functional Parameters", open=False):
-                        # Create Gradio components for parameter inputs
-                        enable_finish = gr.Checkbox(
-                            label="Enable Finish", value=True)
-                        enable_rag = gr.Checkbox(
-                            label="Enable RAG", value=True)
-                        enable_summary = gr.Checkbox(
-                            label="Enable Summary", value=False)
-                        init_rag_num = gr.Number(
-                            label="Initial RAG Num", value=0)
-                        step_rag_num = gr.Number(
-                            label="Step RAG Num", value=10)
-                        skip_last_k = gr.Number(label="Skip Last K", value=0)
-                        summary_mode = gr.Textbox(
-                            label="Summary Mode", value='step')
-                        summary_skip_last_k = gr.Number(
-                            label="Summary Skip Last K", value=0)
-                        summary_context_length = gr.Number(
-                            label="Summary Context Length", value=None)
-                        force_finish = gr.Checkbox(
-                            label="Force FinalAnswer", value=True)
-                        seed = gr.Number(label="Seed", value=100)
-                        # Button to submit and update parameters
-                        submit_btn = gr.Button("Update Parameters")
-
-                        # Display the updated parameters
-                        updated_parameters_output = gr.JSON()
-
-                        # When button is clicked, update parameters
-                        submit_btn.click(fn=update_model_parameters,
-                                         inputs=[enable_finish, enable_rag, enable_summary, init_rag_num, step_rag_num, skip_last_k,
-                                                 summary_mode, summary_skip_last_k, summary_context_length, force_finish, seed],
-                                         outputs=updated_parameters_output)
-        # Button to submit the password
-        submit_button = gr.Button("Submit")
-
-        # When the button is clicked, check if the password is correct
-        submit_button.click(
-            check_password,
-            inputs=password_input,
-            outputs=[protected_accordion, incorrect_message]
-        )
-    gr.Markdown(LICENSE)
 
+        with gr.Accordion("Settings", open=False):
+            temperature_slider = gr.Slider(
+                minimum=0,
+                maximum=1,
+                step=0.1,
+                value=default_temperature,
+                label="Temperature"
+            )
+            max_new_tokens_slider = gr.Slider(
+                minimum=128,
+                maximum=4096,
+                step=1,
+                value=default_max_new_tokens,
+                label="Max new tokens"
+            )
+            max_tokens_slider = gr.Slider(
+                minimum=128,
+                maximum=32000,
+                step=1,
+                value=default_max_tokens,
+                label="Max tokens"
+            )
+            max_round_slider = gr.Slider(
+                minimum=0,
+                maximum=50,
+                step=1,
+                value=default_max_round,
+                label="Max round")
+
+            temperature_slider.change(
+                lambda x: x, inputs=temperature_slider, outputs=temperature_state)
+            max_new_tokens_slider.change(
+                lambda x: x, inputs=max_new_tokens_slider, outputs=max_new_tokens_state)
+            max_tokens_slider.change(
+                lambda x: x, inputs=max_tokens_slider, outputs=max_tokens_state)
+            max_round_slider.change(
+                lambda x: x, inputs=max_round_slider, outputs=max_round_state)
+
+            password_input = gr.Textbox(
+                label="Enter Password for More Settings", type="password")
+            incorrect_message = gr.Textbox(visible=False, interactive=False)
+            
+            with gr.Accordion("⚙️ Settings", open=False, visible=False) as protected_accordion:
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        with gr.Accordion("⚙️ Model Loading", open=False):
+                            model_name_input = gr.Textbox(
+                                label="Enter model path", value=model_name)
+                            load_model_btn = gr.Button(value="Load Model")
+                            load_model_btn.click(
+                                agent.load_models, inputs=model_name_input, outputs=gr.Textbox(label="Status"))
+                    
+                    with gr.Column(scale=1):
+                        with gr.Accordion("⚙️ Functional Parameters", open=False):
+                            enable_finish = gr.Checkbox(label="Enable Finish", value=True)
+                            enable_rag = gr.Checkbox(label="Enable RAG", value=True)
+                            enable_summary = gr.Checkbox(label="Enable Summary", value=False)
+                            init_rag_num = gr.Number(label="Initial RAG Num", value=0)
+                            step_rag_num = gr.Number(label="Step RAG Num", value=10)
+                            skip_last_k = gr.Number(label="Skip Last K", value=0)
+                            summary_mode = gr.Textbox(label="Summary Mode", value='step')
+                            summary_skip_last_k = gr.Number(label="Summary Skip Last K", value=0)
+                            summary_context_length = gr.Number(label="Summary Context Length", value=None)
+                            force_finish = gr.Checkbox(label="Force FinalAnswer", value=True)
+                            seed = gr.Number(label="Seed", value=100)
+                            
+                            submit_btn = gr.Button("Update Parameters")
+                            updated_parameters_output = gr.JSON()
+                            
+                            submit_btn.click(
+                                fn=update_model_parameters,
+                                inputs=[enable_finish, enable_rag, enable_summary, init_rag_num, 
+                                       step_rag_num, skip_last_k, summary_mode, summary_skip_last_k, 
+                                       summary_context_length, force_finish, seed],
+                                outputs=updated_parameters_output
+                            )
+            
+            submit_button = gr.Button("Submit")
+            submit_button.click(
+                check_password,
+                inputs=password_input,
+                outputs=[protected_accordion, incorrect_message]
+            )
+        
+        gr.Markdown(LICENSE)
+    
+    return demo
 
 if __name__ == "__main__":
+    # Initialize the agent only when running directly
+    agent = TxAgent(
+        model_name,
+        rag_model_name,
+        tool_files_dict=new_tool_files,
+        force_finish=True,
+        enable_checker=True,
+        step_rag_num=10,
+        seed=100,
+        additional_default_tools=['DirectResponse', 'RequireClarification']
+    )
+    agent.init_model()
+    
+    # Create and launch the interface
+    demo = create_interface(agent)
     demo.launch(share=True)