Skip to content

Instantly share code, notes, and snippets.

@amaarora

amaarora/nb_qwen_tool_use.ipynb

Created September 14, 2025 03:57
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save amaarora/428992efa297e6ea79c661f2abe16203 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save amaarora/428992efa297e6ea79c661f2abe16203 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
nb_qwen_tool_use.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "b3b02170",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/opt/miniconda3/envs/ai/lib/python3.12/site-packages/sklearn/utils/_param_validation.py:11: UserWarning: A NumPy version >=1.22.4 and <2.3.0 is required for this version of SciPy (detected version 2.3.3)\n",
" from scipy.sparse import csr_matrix, issparse\n"
]
}
],
"source": [
"from transformers import AutoModelForCausalLM, AutoTokenizer"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "4406244a",
"metadata": {},
"outputs": [],
"source": [
"model_name = \"Qwen/Qwen3-0.6B\"\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "e7ecc0bd",
"metadata": {},
"outputs": [],
"source": [
"messages = [\n",
" {\"role\": \"system\", \"content\": \"You are a bot that responds to weather queries. You should reply with the unit used in the queried location.\"},\n",
" {\"role\": \"user\", \"content\": \"Hey, what's the temperature in Paris right now?\"}\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "a645affe",
"metadata": {},
"outputs": [],
"source": [
"def get_current_temperature(location: str, unit: str):\n",
" \"\"\"\n",
" Get the current temperature at a location.\n",
" \n",
" Args:\n",
" location: The location to get the temperature for, in the format \"City, Country\"\n",
" unit: The unit to return the temperature in. (choices: [\"celsius\", \"fahrenheit\"])\n",
" \"\"\"\n",
" return 22. # A real function should probably actually get the temperature!\n",
"\n",
"def get_current_wind_speed(location: str):\n",
" \"\"\"\n",
" Get the current wind speed in km/h at a given location.\n",
" \n",
" Args:\n",
" location: The location to get the wind speed for, in the format \"City, Country\"\n",
" \"\"\"\n",
" return 6. # A real function should probably actually get the wind speed!\n",
"\n",
"tools = [get_current_temperature, get_current_wind_speed]"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "f0b4408a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|im_start|>system\n",
"You are a bot that responds to weather queries. You should reply with the unit used in the queried location.\n",
"\n",
"# Tools\n",
"\n",
"You may call one or more functions to assist with the user query.\n",
"\n",
"You are provided with function signatures within <tools></tools> XML tags:\n",
"<tools>\n",
"{\"type\": \"function\", \"function\": {\"name\": \"get_current_temperature\", \"description\": \"Get the current temperature at a location.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The location to get the temperature for, in the format \\\"City, Country\\\"\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit to return the temperature in.\"}}, \"required\": [\"location\", \"unit\"]}}}\n",
"{\"type\": \"function\", \"function\": {\"name\": \"get_current_wind_speed\", \"description\": \"Get the current wind speed in km/h at a given location.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The location to get the wind speed for, in the format \\\"City, Country\\\"\"}}, \"required\": [\"location\"]}}}\n",
"</tools>\n",
"\n",
"For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n",
"<tool_call>\n",
"{\"name\": <function-name>, \"arguments\": <args-json-object>}\n",
"</tool_call><|im_end|>\n",
"<|im_start|>user\n",
"Hey, what's the temperature in Paris right now?<|im_end|>\n",
"<|im_start|>assistant\n",
"\n"
]
}
],
"source": [
"inputs = tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, tokenize=False)\n",
"print(inputs)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "b03ae44a",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"`torch_dtype` is deprecated! Use `dtype` instead!\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 2.05 s, sys: 527 ms, total: 2.57 s\n",
"Wall time: 4.16 s\n"
]
}
],
"source": [
"%%time\n",
"model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=\"auto\", device_map=\"auto\")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "dd758630",
"metadata": {},
"outputs": [],
"source": [
"inputs = tokenizer.apply_chat_template(\n",
" messages, \n",
" tools=tools, \n",
" add_generation_prompt=True, \n",
" return_dict=True, \n",
" return_tensors=\"pt\"\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "84d3049e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 7.41 s, sys: 1.27 s, total: 8.67 s\n",
"Wall time: 8.18 s\n"
]
}
],
"source": [
"%%time\n",
"outputs = model.generate(**inputs.to(model.device), max_new_tokens=1024)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "b153eeb9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<think>\n",
"Okay, the user is asking for the current temperature in Paris. I need to use the get_current_temperature function. The parameters required are location and unit. The location here is Paris, and the unit isn't specified, so maybe default to Celsius since that's commonly used in Europe. Let me check the function's enum for unit. The enum is celsius and fahrenheit. Since Paris is in France, which uses Celsius, I'll set the unit to celsius. So the tool call should include location: \"Paris\" and unit: \"celsius\". That should get the temperature for Paris.\n",
"</think>\n",
"\n",
"<tool_call>\n",
"{\"name\": \"get_current_temperature\", \"arguments\": {\"location\": \"Paris\", \"unit\": \"celsius\"}}\n",
"</tool_call><|im_end|>\n"
]
}
],
"source": [
"print(tokenizer.decode(outputs[0][len(inputs[\"input_ids\"][0]):]))"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "ebdb50f5",
"metadata": {},
"outputs": [],
"source": [
"tool_call = {\"name\": \"get_current_temperature\", \"arguments\": {\"location\": \"Paris, France\", \"unit\": \"celsius\"}}\n",
"messages.append({\"role\": \"assistant\", \"tool_calls\": [{\"type\": \"function\", \"function\": tool_call}]})"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "e3510ecc",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|im_start|>system\n",
"You are a bot that responds to weather queries. You should reply with the unit used in the queried location.\n",
"\n",
"# Tools\n",
"\n",
"You may call one or more functions to assist with the user query.\n",
"\n",
"You are provided with function signatures within <tools></tools> XML tags:\n",
"<tools>\n",
"{\"type\": \"function\", \"function\": {\"name\": \"get_current_temperature\", \"description\": \"Get the current temperature at a location.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The location to get the temperature for, in the format \\\"City, Country\\\"\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit to return the temperature in.\"}}, \"required\": [\"location\", \"unit\"]}}}\n",
"{\"type\": \"function\", \"function\": {\"name\": \"get_current_wind_speed\", \"description\": \"Get the current wind speed in km/h at a given location.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The location to get the wind speed for, in the format \\\"City, Country\\\"\"}}, \"required\": [\"location\"]}}}\n",
"</tools>\n",
"\n",
"For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n",
"<tool_call>\n",
"{\"name\": <function-name>, \"arguments\": <args-json-object>}\n",
"</tool_call><|im_end|>\n",
"<|im_start|>user\n",
"Hey, what's the temperature in Paris right now?<|im_end|>\n",
"<|im_start|>assistant\n",
"<think>\n",
"\n",
"</think>\n",
"\n",
"<tool_call>\n",
"{\"name\": \"get_current_temperature\", \"arguments\": {\"location\": \"Paris, France\", \"unit\": \"celsius\"}}\n",
"</tool_call><|im_end|>\n",
"<|im_start|>assistant\n",
"\n"
]
}
],
"source": [
"inputs = tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, tokenize=False)\n",
"print(inputs)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "0b4df234",
"metadata": {},
"outputs": [],
"source": [
"messages.append({\"role\": \"tool\", \"content\": \"22\"})"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "a5303785",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<|im_start|>system\n",
"You are a bot that responds to weather queries. You should reply with the unit used in the queried location.\n",
"\n",
"# Tools\n",
"\n",
"You may call one or more functions to assist with the user query.\n",
"\n",
"You are provided with function signatures within <tools></tools> XML tags:\n",
"<tools>\n",
"{\"type\": \"function\", \"function\": {\"name\": \"get_current_temperature\", \"description\": \"Get the current temperature at a location.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The location to get the temperature for, in the format \\\"City, Country\\\"\"}, \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"], \"description\": \"The unit to return the temperature in.\"}}, \"required\": [\"location\", \"unit\"]}}}\n",
"{\"type\": \"function\", \"function\": {\"name\": \"get_current_wind_speed\", \"description\": \"Get the current wind speed in km/h at a given location.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"location\": {\"type\": \"string\", \"description\": \"The location to get the wind speed for, in the format \\\"City, Country\\\"\"}}, \"required\": [\"location\"]}}}\n",
"</tools>\n",
"\n",
"For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n",
"<tool_call>\n",
"{\"name\": <function-name>, \"arguments\": <args-json-object>}\n",
"</tool_call><|im_end|>\n",
"<|im_start|>user\n",
"Hey, what's the temperature in Paris right now?<|im_end|>\n",
"<|im_start|>assistant\n",
"<tool_call>\n",
"{\"name\": \"get_current_temperature\", \"arguments\": {\"location\": \"Paris, France\", \"unit\": \"celsius\"}}\n",
"</tool_call><|im_end|>\n",
"<|im_start|>user\n",
"<tool_response>\n",
"22\n",
"</tool_response><|im_end|>\n",
"<|im_start|>assistant\n",
"\n"
]
}
],
"source": [
"inputs = tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, tokenize=False)\n",
"print(inputs)"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "25a5d8dc",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<think>\n",
"Okay, the user asked for the temperature in Paris right now. I used the get_current_temperature function with Paris, France, and celsius. The response came back as 22. Now I need to present this answer clearly.\n",
"\n",
"I should mention the city, the unit, and the temperature. Let me check if the unit is correct. The user didn't specify, but since they asked in a general query, maybe they expect Celsius. Also, make sure to keep it friendly and concise. Maybe add a note about the current weather. Alright, that should cover it.\n",
"</think>\n",
"\n",
"The current temperature in Paris is **22°C**. Let me know if you need further details! 🌤️<|im_end|>\n",
"CPU times: user 6.49 s, sys: 1.01 s, total: 7.5 s\n",
"Wall time: 7.02 s\n"
]
}
],
"source": [
"%%time\n",
"inputs = tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, return_dict=True, return_tensors=\"pt\")\n",
"out = model.generate(**inputs.to(model.device), max_new_tokens=1024)\n",
"print(tokenizer.decode(out[0][len(inputs[\"input_ids\"][0]):]))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.7"
},
"toc": {
"base_numbering": 1,
"nav_menu": {},
"number_sections": true,
"sideBar": true,
"skip_h1_title": false,
"title_cell": "Table of Contents",
"title_sidebar": "Contents",
"toc_cell": false,
"toc_position": {},
"toc_section_display": true,
"toc_window_display": false
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment