test / app.py
Ali2206's picture
Update app.py
92abf33 verified
raw
history blame
5.17 kB
import os
import gradio as gr
from txagent import TxAgent
# ========== Configuration ==========
current_dir = os.path.dirname(os.path.abspath(__file__))
os.environ["MKL_THREADING_LAYER"] = "GNU"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Model configuration
MODEL_CONFIG = {
'model_name': 'mims-harvard/TxAgent-T1-Llama-3.1-8B',
'rag_model_name': 'mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B',
'tool_files': {'new_tool': os.path.join(current_dir, 'data', 'new_tool.json')},
'additional_tools': ['DirectResponse', 'RequireClarification'],
'default_params': {
'force_finish': True,
'enable_checker': True,
'step_rag_num': 10,
'seed': 100
}
}
# UI Configuration
UI_CONFIG = {
'description': '''
<div>
<h1 style="text-align: center;">TxAgent: Therapeutic Reasoning AI</h1>
<p style="text-align: center;">Precision therapeutics with multi-step reasoning</p>
</div>
''',
'disclaimer': '''
<div style="color: #666; font-size: 0.9em; margin-top: 20px;">
<strong>Disclaimer:</strong> For informational purposes only, not medical advice.
</div>
'''
}
# Example questions
EXAMPLE_QUESTIONS = [
"How should dosage be adjusted for hepatic impairment with Journavx?",
"Is Xolremdi suitable with Prozac for WHIM syndrome?",
"What are Warfarin-Amiodarone contraindications?"
]
# ========== Application Class ==========
class TxAgentApplication:
def __init__(self):
self.agent = None
self.is_initialized = False
def initialize_agent(self):
if self.is_initialized:
return "Model already initialized"
try:
self.agent = TxAgent(
MODEL_CONFIG['model_name'],
MODEL_CONFIG['rag_model_name'],
tool_files_dict=MODEL_CONFIG['tool_files'],
**MODEL_CONFIG['default_params']
)
self.agent.init_model()
self.is_initialized = True
return "TxAgent initialized successfully"
except Exception as e:
return f"Initialization failed: {str(e)}"
def chat(self, message, chat_history):
if not self.is_initialized:
yield "Error: Please initialize the model first"
return
try:
# Convert to messages format
messages = []
for user, assistant in chat_history:
messages.append({"role": "user", "content": user})
messages.append({"role": "assistant", "content": assistant})
messages.append({"role": "user", "content": message})
# Stream response
full_response = ""
for chunk in self.agent.run_gradio_chat(
messages,
temperature=0.3,
max_new_tokens=1024,
max_tokens=8192,
multi_agent=False,
conversation=[],
max_round=30
):
full_response += chunk
yield [(message, full_response)]
except Exception as e:
yield [(message, f"Error: {str(e)}")]
# ========== Gradio Interface ==========
def create_interface():
app = TxAgentApplication()
with gr.Blocks(title="TxAgent", theme=gr.themes.Soft()) as demo:
gr.Markdown(UI_CONFIG['description'])
# Initialization
with gr.Row():
init_btn = gr.Button("Initialize TxAgent", variant="primary")
init_status = gr.Textbox(label="Status", interactive=False)
# Chat Interface (using modern messages format)
chatbot = gr.Chatbot(
height=600,
label="Conversation",
avatar_images=(
"https://example.com/user.png", # User avatar
"https://example.com/bot.png" # Bot avatar
)
)
with gr.Row():
msg = gr.Textbox(
label="Your Question",
placeholder="Ask about drug interactions or treatments...",
scale=4
)
submit_btn = gr.Button("Submit", variant="primary", scale=1)
# Examples
gr.Examples(
examples=EXAMPLE_QUESTIONS,
inputs=msg,
label="Try these examples:"
)
gr.Markdown(UI_CONFIG['disclaimer'])
# Event Handlers
init_btn.click(
app.initialize_agent,
outputs=init_status
)
msg.submit(
app.chat,
[msg, chatbot],
[chatbot]
)
submit_btn.click(
app.chat,
[msg, chatbot],
[chatbot]
).then(
lambda: "", None, msg
)
return demo
# ========== Main Execution ==========
if __name__ == "__main__":
interface = create_interface()
# Correct launch configuration
interface.launch(
server_name="0.0.0.0",
server_port=7860,
share=True,
enable_queue=True # Enable queue without deprecated parameters
)