1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model = AutoModelForCausalLM.from_pretrained("distil-labs/distil-qwen3-0.6b-voice-assistant-banking")
5tokenizer = AutoTokenizer.from_pretrained("distil-labs/distil-qwen3-0.6b-voice-assistant-banking")
6
7TOOLS = [
8 {"type": "function", "function": {"name": "check_balance", "description": "Check the balance of a bank account", "parameters": {"type": "object", "properties": {"account_type": {"type": "string", "enum": ["checking", "savings", "credit"]}}, "required": [], "additionalProperties": False}}},
9 {"type": "function", "function": {"name": "transfer_money", "description": "Transfer money between accounts", "parameters": {"type": "object", "properties": {"amount": {"type": "number"}, "from_account": {"type": "string", "enum": ["checking", "savings"]}, "to_account": {"type": "string", "enum": ["checking", "savings"]}}, "required": [], "additionalProperties": False}}},
10 {"type": "function", "function": {"name": "cancel_card", "description": "Cancel a bank card", "parameters": {"type": "object", "properties": {"card_type": {"type": "string", "enum": ["credit", "debit"]}, "card_last_four": {"type": "string"}, "reason": {"type": "string", "enum": ["lost", "stolen", "damaged", "other"]}}, "required": [], "additionalProperties": False}}},
11 {"type": "function", "function": {"name": "intent_unclear", "description": "Use when the user's intent cannot be determined", "parameters": {"type": "object", "properties": {}, "required": [], "additionalProperties": False}}},
12 {"type": "function", "function": {"name": "greeting", "description": "User is greeting", "parameters": {"type": "object", "properties": {}, "required": [], "additionalProperties": False}}},
13 {"type": "function", "function": {"name": "goodbye", "description": "User is ending the conversation", "parameters": {"type": "object", "properties": {}, "required": [], "additionalProperties": False}}},
14]
15
16messages = [
17 {"role": "system", "content": "You are a tool-calling model working on:\n<task_description>You are a voice assistant for BankCo, a retail bank. The user input is automatically transcribed speech from an ASR system, so it may contain transcription errors, homophones, filler words, or unusual phrasings. Parse the user's request and return the appropriate function call despite any transcription artifacts. If you can identify the intent, call the matching function. Extract any mentioned argument values; omit arguments not mentioned. If you cannot understand what the user wants, call intent_unclear(). Use conversation history to understand context from previous turns.</task_description>\n\nRespond to the conversation history by generating an appropriate tool call that satisfies the user request. Generate only the tool call according to the provided tool schema, do not generate anything else. Always respond with a tool call.\n\n"},
18 {"role": "user", "content": "I need to cancel my credit card ending in 1234"},
19]
20
21text = tokenizer.apply_chat_template(
22 messages, tools=TOOLS, tokenize=False, add_generation_prompt=True,
23 enable_thinking=False,
24)
25inputs = tokenizer(text, return_tensors="pt")
26outputs = model.generate(**inputs, max_new_tokens=256, temperature=0)
27print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))
28# <tool_call>
29# {"name": "cancel_card", "arguments": {"card_type": "credit", "card_last_four": "1234"}}
30# </tool_call>