Views
No views yet

[!NOTE] For evaluation, we use YaRN scaling to deploy the models for Multi-Turn evaluation, and all Arch-Agent models are evaluated with a context length of 64K.
transformers library and we recommend to install latest version:pip install transformers>=4.51.01import json
2from typing import Any, Dict, List
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5model_name = "katanemo/Arch-Agent-1.5B"
6
7model = AutoModelForCausalLM.from_pretrained(
8 model_name, device_map="auto", torch_dtype="auto", trust_remote_code=True
9)
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11
12TASK_PROMPT = (
13 "You are a helpful assistant designed to assist with the user query by making one or more function calls if needed."
14 "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n"
15 "You are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{tool_text}"
16 "\n</tools>\n\nFor each function call, return a json object with function name and arguments within "
17 """<tool_call></tool_call> XML tags:\n<tool_call>\n{{"name": <function-name>, """
18 """"arguments": <args-json-object>}}\n</tool_call>"""
19)
20
21# Define available tools
22tools = [
23 {
24 "type": "function",
25 "function": {
26 "name": "get_weather",
27 "description": "Get the current weather for a location",
28 "parameters": {
29 "type": "object",
30 "properties": {
31 "location": {
32 "type": "str",
33 "description": "The city and state, e.g. San Francisco, New York",
34 },
35 "unit": {
36 "type": "str",
37 "enum": ["celsius", "fahrenheit"],
38 "description": "The unit of temperature to return",
39 },
40 },
41 "required": ["location"],
42 },
43 },
44 }
45]
46
47# Helper function to create the system prompt for our model
48def format_prompt(tools: List[Dict[str, Any]]):
49 tool_text = "\n".join(
50 [json.dumps(tool["function"], ensure_ascii=False) for tool in tools]
51 )
52 return TASK_PROMPT.format(tool_text=tool_text)
53
54system_prompt = format_prompt(tools)
55
56messages = [
57 {"role": "system", "content": system_prompt},
58 {"role": "user", "content": "What is the weather in Seattle?"},
59]
60
61model_inputs = tokenizer.apply_chat_template(
62 messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
63).to(model.device)
64
65generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
66generated_ids = [
67 output_ids[len(input_ids) :]
68 for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
69]
70
71response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
72print(response)