Views
No views yet
transformers library and we advise you to install latest version:pip install transformers>=4.37.01import json
2from typing import Any, Dict, List
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5model_name = "katanemo/Arch-Function-Chat-3B"
6model = AutoModelForCausalLM.from_pretrained(
7 model_name, device_map="auto", torch_dtype="auto", trust_remote_code=True
8)
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10
11# Please use our provided prompt for best performance
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\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{tools}\n</tools>"
15 "\n\nYour task is to decide which functions are needed and collect missing parameters if necessary."
16)
17
18FORMAT_PROMPT = (
19 "\n\nBased on your analysis, provide your response in one of the following JSON formats:"
20 '\n1. If no functions are needed:\n```json\n{"response": "Your response text here"}\n```'
21 '\n2. If functions are needed but some required parameters are missing:\n```json\n{"required_functions": ["func_name1", "func_name2", ...], "clarification": "Text asking for missing parameters"}\n```'
22 '\n3. If functions are needed and all required parameters are available:\n```json\n{"tool_calls": [{"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},... (more tool calls as required)]}\n```'
23)
24
25# Define available tools
26tools = [
27 {
28 "type": "function",
29 "function": {
30 "name": "get_weather",
31 "description": "Get the current weather for a location",
32 "parameters": {
33 "type": "object",
34 "properties": {
35 "location": {
36 "type": "str",
37 "description": "The city and state, e.g. San Francisco, New York",
38 },
39 "unit": {
40 "type": "str",
41 "enum": ["celsius", "fahrenheit"],
42 "description": "The unit of temperature to return",
43 },
44 },
45 "required": ["location"],
46 },
47 },
48 }
49]
50
51
52# Helper function to create the system prompt for our model
53def format_prompt(tools: List[Dict[str, Any]]):
54 tools = "\n".join(
55 [json.dumps(tool["function"], ensure_ascii=False) for tool in tools]
56 )
57 return TASK_PROMPT.format(tools=tools) + FORMAT_PROMPT
58
59
60system_prompt = format_prompt(tools)
61
62messages = [
63 {"role": "system", "content": system_prompt},
64 {"role": "user", "content": "What is the weather in Seattle?"},
65]
66
67model_inputs = tokenizer.apply_chat_template(
68 messages, add_generation_prompt=True, return_tensors="pt"
69).to(model.device)
70
71generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
72
73generated_ids = [
74 output_ids[len(input_ids) :]
75 for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
76]
77
78response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
79print(response)