1import json
2import re
3from typing import Optional
4
5from jinja2 import Template
6import torch
7from transformers import AutoModelForCausalLM, AutoTokenizer
8from transformers.utils import get_json_schema
9
10
11system_prompt = Template("""You are an expert in composing functions. You are given a question and a set of possible functions.
12Based on the question, you will need to make one or more function/tool calls to achieve the purpose.
13If none of the functions can be used, point it out and refuse to answer.
14If the given question lacks the parameters required by the function, also point it out.
15
16You have access to the following tools:
17<tools>{{ tools }}</tools>
18
19The output MUST strictly adhere to the following format, and NO other text MUST be included.
20The example format is as follows. Please make sure the parameter type is correct. If no function call is needed, please make the tool calls an empty list '[]'.
21<tool_call>[
22{"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},
23... (more tool calls as required)
24]</tool_call>""")
25
26
27def prepare_messages(
28 query: str,
29 tools: Optional[dict[str, any]] = None,
30 history: Optional[list[dict[str, str]]] = None
31) -> list[dict[str, str]]:
32 """Prepare the system and user messages for the given query and tools.
33
34 Args:
35 query: The query to be answered.
36 tools: The tools available to the user. Defaults to None, in which case if a
37 list without content will be passed to the model.
38 history: Exchange of messages, including the system_prompt from
39 the first query. Defaults to None, the first message in a conversation.
40 """
41 if tools is None:
42 tools = []
43 if history:
44 messages = history.copy()
45 messages.append({"role": "user", "content": query})
46 else:
47 messages = [
48 {"role": "system", "content": system_prompt.render(tools=json.dumps(tools))},
49 {"role": "user", "content": query}
50 ]
51 return messages
52
53
54def parse_response(text: str) -> str | dict[str, any]:
55 """Parses a response from the model, returning either the
56 parsed list with the tool calls parsed, or the
57 model thought or response if couldn't generate one.
58
59 Args:
60 text: Response from the model.
61 """
62 pattern = r"<tool_call>(.*?)</tool_call>"
63 matches = re.findall(pattern, text, re.DOTALL)
64 if matches:
65 return json.loads(matches[0])
66 return text
67
68
69model_name_smollm = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
70model = AutoModelForCausalLM.from_pretrained(model_name_smollm, device_map="auto", torch_dtype="auto", trust_remote_code=True)
71tokenizer = AutoTokenizer.from_pretrained(model_name_smollm)
72
73from datetime import datetime
74import random
75
76def get_current_time() -> str:
77 """Returns the current time in 24-hour format.
78
79 Returns:
80 str: Current time in HH:MM:SS format.
81 """
82 return datetime.now().strftime("%H:%M:%S")
83
84
85def get_random_number_between(min: int, max: int) -> int:
86 """
87 Gets a random number between min and max.
88
89 Args:
90 min: The minimum number.
91 max: The maximum number.
92
93 Returns:
94 A random number between min and max.
95 """
96 return random.randint(min, max)
97
98
99tools = [get_json_schema(get_random_number_between), get_json_schema(get_current_time)]
100
101toolbox = {"get_random_number_between": get_random_number_between, "get_current_time": get_current_time}
102
103query = "Give me a number between 1 and 300"
104
105messages = prepare_messages(query, tools=tools)
106
107inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
108outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
109result = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
110
111tool_calls = parse_response(result)
112# [{'name': 'get_random_number_between', 'arguments': {'min': 1, 'max': 300}}
113
114# Get tool responses
115tool_responses = [toolbox.get(tc["name"])(*tc["arguments"].values()) for tc in tool_calls]
116# [63]
117
118# For the second turn, rebuild the history of messages:
119history = messages.copy()
120# Add the "parsed response"
121history.append({"role": "assistant", "content": result})
122query = "Can you give me the hour?"
123history.append({"role": "user", "content": query})
124
125inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, return_tensors="pt").to(model.device)
126outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
127result = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
128
129tool_calls = parse_response(result)
130tool_responses = [toolbox.get(tc["name"])(*tc["arguments"].values()) for tc in tool_calls]
131# ['07:57:25']
