1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model = AutoModelForCausalLM.from_pretrained("distil-labs/distil-home-assistant-functiongemma")
5tokenizer = AutoTokenizer.from_pretrained("distil-labs/distil-home-assistant-functiongemma")
6
7TOOLS = [
8 {"type": "function", "function": {"name": "toggle_lights", "description": "Turn lights on or off in a specified room", "parameters": {"type": "object", "properties": {"room": {"type": "string", "enum": ["living_room", "bedroom", "kitchen", "bathroom", "office", "hallway"]}, "state": {"type": "string", "enum": ["on", "off"]}}, "required": [], "additionalProperties": False}}},
9 {"type": "function", "function": {"name": "set_thermostat", "description": "Set the temperature for heating or cooling", "parameters": {"type": "object", "properties": {"temperature": {"type": "integer", "minimum": 60, "maximum": 80}, "mode": {"type": "string", "enum": ["heat", "cool", "auto"]}}, "required": [], "additionalProperties": False}}},
10 {"type": "function", "function": {"name": "lock_door", "description": "Lock or unlock a door", "parameters": {"type": "object", "properties": {"door": {"type": "string", "enum": ["front", "back", "garage", "side"]}, "state": {"type": "string", "enum": ["lock", "unlock"]}}, "required": [], "additionalProperties": False}}},
11 {"type": "function", "function": {"name": "get_device_status", "description": "Query the current state of a device or room", "parameters": {"type": "object", "properties": {"device_type": {"type": "string", "enum": ["lights", "thermostat", "door", "all"]}, "room": {"type": "string"}}, "required": [], "additionalProperties": False}}},
12 {"type": "function", "function": {"name": "set_scene", "description": "Activate a predefined scene", "parameters": {"type": "object", "properties": {"scene": {"type": "string", "enum": ["movie_night", "bedtime", "morning", "away", "party"]}}, "required": [], "additionalProperties": False}}},
13 {"type": "function", "function": {"name": "intent_unclear", "description": "Use when the user's intent cannot be determined", "parameters": {"type": "object", "properties": {"reason": {"type": "string", "enum": ["ambiguous", "off_topic", "incomplete", "unsupported_device"]}}, "required": [], "additionalProperties": False}}},
14]
15
16messages = [
17 {"role": "system", "content": "You are a tool-calling model working on:\n<task_description>You are an on-device smart home controller. Given a natural language command from the user, call the appropriate smart home function. If the user does not specify a required value (e.g. which room or what temperature), omit that parameter from the function call. Maintain context across conversation turns to resolve pronouns and sequential commands.</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": "Turn off the living room lights"},
19]
20
21text = tokenizer.apply_chat_template(
22 messages, tools=TOOLS, tokenize=False, add_generation_prompt=True,
23)
24inputs = tokenizer(text, return_tensors="pt")
25outputs = model.generate(**inputs, max_new_tokens=256, temperature=0)
26print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))
27# <tool_call>
28# {"name": "toggle_lights", "arguments": {"room": "living_room", "state": "off"}}
29# </tool_call>