File size: 4,892 Bytes
1155704 8c814cd 47f0902 9aeb1dd 47f0902 9a0af2f 6ba76fa 0cec600 9a0af2f 6ba76fa 47f0902 0cec600 1155704 0cec600 47f0902 0cec600 47f0902 0cec600 47f0902 0cec600 47f0902 0cec600 47f0902 0cec600 47f0902 0cec600 1155704 0cec600 47f0902 0cec600 47f0902 0cec600 47f0902 0cec600 6309d92 47f0902 6309d92 9a0af2f 6309d92 9a0af2f 6309d92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import os
import sys
import random
import gradio as gr
# Add `src` directory to Python path
sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
from txagent.txagent import TxAgent
# ==== Environment Setup ====
current_dir = os.path.dirname(os.path.abspath(__file__))
os.environ["MKL_THREADING_LAYER"] = "GNU"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# ==== UI Content ====
DESCRIPTION = '''
<div>
<h1 style="text-align: center;">TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools </h1>
</div>
'''
INTRO = "Precision therapeutics require multimodal adaptive models..."
LICENSE = "DISCLAIMER: THIS WEBSITE DOES NOT PROVIDE MEDICAL ADVICE..."
PLACEHOLDER = '''
<div style="padding: 30px; text-align: center;">
<h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">TxAgent</h1>
<p style="font-size: 18px;">Click clear 🗑️ before asking a new question.</p>
<p style="font-size: 18px;">Click retry 🔄 to see another answer.</p>
</div>
'''
css = """
h1 { text-align: center; }
#duplicate-button {
margin: auto;
color: white;
background: #1565c0;
border-radius: 100vh;
}
.gradio-accordion {
margin-top: 0px !important;
margin-bottom: 0px !important;
}
"""
chat_css = """
.gr-button { font-size: 20px !important; }
.gr-button svg { width: 32px !important; height: 32px !important; }
"""
# ==== Model Settings ====
model_name = "mims-harvard/TxAgent-T1-Llama-3.1-8B"
rag_model_name = "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B"
new_tool_files = {
"new_tool": os.path.join(current_dir, "data", "new_tool.json")
}
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 moderate hepatic impairment?"],
["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
]
# Initialize agent placeholder
agent = None
# ===== Build Gradio UI =====
def create_ui():
with gr.Blocks(css=css) as demo:
gr.Markdown(DESCRIPTION)
gr.Markdown(INTRO)
temperature = gr.Slider(0, 1, step=0.1, value=0.3, label="Temperature")
max_new_tokens = gr.Slider(128, 4096, step=1, value=1024, label="Max New Tokens")
max_tokens = gr.Slider(128, 32000, step=1, value=8192, label="Max Total Tokens")
max_round = gr.Slider(1, 50, step=1, value=30, label="Max Rounds")
multi_agent = gr.Checkbox(label="Enable Multi-agent Reasoning", value=False)
conversation_state = gr.State([])
chatbot = gr.Chatbot(
label="TxAgent",
placeholder=PLACEHOLDER,
height=700,
type="messages",
show_copy_button=True
)
# Retry logic
def handle_retry(history, retry_data, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
agent.update_parameters(seed=random.randint(0, 10000))
new_history = history[:retry_data.index]
prompt = history[retry_data.index]["content"]
result = agent.run_gradio_chat(
new_history + [{"role": "user", "content": prompt}],
temperature, max_new_tokens, max_tokens,
multi_agent, conversation, max_round
)
if hasattr(result, "__iter__") and not isinstance(result, (str, dict, list)):
result = list(result)
return result
chatbot.retry(
handle_retry,
chatbot, chatbot,
temperature, max_new_tokens, max_tokens,
multi_agent, conversation_state, max_round
)
# Main handler
def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
return agent.run_gradio_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round)
gr.ChatInterface(
fn=handle_chat,
chatbot=chatbot,
additional_inputs=[
temperature, max_new_tokens, max_tokens,
multi_agent, conversation_state, max_round
],
examples=question_examples,
css=chat_css,
cache_examples=False,
fill_height=True,
fill_width=True,
stop_btn=True
)
gr.Markdown(LICENSE)
return demo
# === VLLM-safe entry point ===
if __name__ == "__main__":
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()
demo = create_ui()
demo.launch()
|