Create backend/agent_instance.py
Browse files- backend/agent_instance.py +27 -0
backend/agent_instance.py
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# backend/agent_instance.py
|
2 |
+
import os
|
3 |
+
from txagent.txagent import TxAgent
|
4 |
+
|
5 |
+
def init_agent():
|
6 |
+
model_cache_dir = os.path.expanduser("~/.cache/txagent_models")
|
7 |
+
os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
|
8 |
+
os.environ["HF_HOME"] = model_cache_dir
|
9 |
+
|
10 |
+
model_name = "mims-harvard/TxAgent-T1-Llama-3.1-8B"
|
11 |
+
rag_model_name = "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B"
|
12 |
+
tool_files_dict = {
|
13 |
+
"new_tool": os.path.abspath("data/new_tool.json")
|
14 |
+
}
|
15 |
+
|
16 |
+
agent = TxAgent(
|
17 |
+
model_name=model_name,
|
18 |
+
rag_model_name=rag_model_name,
|
19 |
+
tool_files_dict=tool_files_dict,
|
20 |
+
force_finish=True,
|
21 |
+
enable_checker=True,
|
22 |
+
step_rag_num=10,
|
23 |
+
seed=100,
|
24 |
+
additional_default_tools=[]
|
25 |
+
)
|
26 |
+
agent.init_model()
|
27 |
+
return agent
|