scripts/schema.py

"""Single source of truth for the tool schema, developer prompt, and prompt
rendering — imported by both step 1 (dataset synthesis) and step 4 (export gate).

Keeping these in one module guarantees the developer+tools prefix the model is
trained on is byte-identical to the one the export gate (and the Unity runtime)
renders. Never redefine TOOLS or DEV anywhere else.
"""
import os
from transformers import AutoTokenizer

import paths

# ============================ tool schema ============================
# Rooms apply ONLY to the light tools. TV / computer / vacuum / speaker are
# single devices and take no room. play_music has a small genre enum (default
# jazz). web_search is the general-knowledge escape hatch.
TOOLS = [
    {"type": "function", "function": {
        "name": "turn_on_light",
        "description": "Turn on the light in a room, optionally at a brightness percent.",
        "parameters": {"type": "object", "properties": {
            "room": {"type": "string", "description": "Which room's light."},
            "brightness": {"type": "integer", "description": "Brightness percent 0-100 (default 100)."},
        }, "required": ["room"]}}},
    {"type": "function", "function": {
        "name": "turn_off_light",
        "description": "Turn off the light in a room.",
        "parameters": {"type": "object", "properties": {
            "room": {"type": "string", "description": "Which room's light."},
        }, "required": ["room"]}}},
    {"type": "function", "function": {
        "name": "set_light_color",
        "description": "Change a room's light color.",
        "parameters": {"type": "object", "properties": {
            "room": {"type": "string", "description": "Which room's light."},
            "color": {"type": "string", "description": "Color name.",
                      "enum": ["red", "orange", "yellow", "green", "blue",
                               "purple", "pink", "white", "warm", "cool"]},
        }, "required": ["room", "color"]}}},
    {"type": "function", "function": {
        "name": "turn_on_tv", "description": "Turn on the TV.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "turn_off_tv", "description": "Turn off the TV.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "turn_on_computer", "description": "Turn on the computer.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "turn_off_computer", "description": "Turn off the computer.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "start_vacuum", "description": "Start the robot vacuum.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "play_music",
        "description": "Start playing music, optionally of a genre (default jazz).",
        "parameters": {"type": "object", "properties": {
            "genre": {"type": "string", "description": "Music genre.",
                      "enum": ["rock", "jazz", "lofi"]},
        }, "required": []}}},
    {"type": "function", "function": {
        "name": "stop_music", "description": "Stop the music.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "set_volume",
        "description": "Set the speaker volume to an absolute percent.",
        "parameters": {"type": "object", "properties": {
            "volume": {"type": "integer", "description": "Volume percent 0-100."},
        }, "required": ["volume"]}}},
    {"type": "function", "function": {
        "name": "get_volume", "description": "Get the current speaker volume.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "get_weather",
        "description": "Get the current weather / temperature for a city.",
        "parameters": {"type": "object", "properties": {
            "city": {"type": "string", "description": "City name (corrected if misspelled)."},
        }, "required": ["city"]}}},
    {"type": "function", "function": {
        "name": "get_time",
        "description": "Get the current time; no city = the device's current city.",
        "parameters": {"type": "object", "properties": {
            "city": {"type": "string", "description": "Optional city for world time."},
        }, "required": []}}},
    {"type": "function", "function": {
        "name": "get_location",
        "description": "Get the device's current location (city / area).",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "web_search",
        "description": "Search the web for general knowledge / facts and answer from the results.",
        "parameters": {"type": "object", "properties": {
            "query": {"type": "string", "description": "The search query."},
        }, "required": ["query"]}}},
]

DEV = ("You convert the user's command (English or Korean) into one or more function calls "
       "with correctly filled arguments, in the order the user asked. Rooms apply only to lights. "
       "If the request does not match any function, reply briefly in text instead of calling.")

# Held-out evaluation queries — kept OUT of the training set (step 1 filters on
# this) and used by the step 4 ORT parity gate, so the gate always tests unseen
# phrasings. Keep the two in sync.
BENCH = {
    "거실 불 30퍼센트로 켜줘", "turn on the living room light at 30 percent", "마루 불 켜줘",
    "거실 불 꺼줘", "turn off the kitchen light", "거실 조명 파란색으로 바꿔줘",
    "티비 켜줘", "turn on the tv", "컴퓨터 켜줘", "turn off the computer", "청소기 돌려줘",
    "음악 틀어줘", "play music", "락 음악 틀어줘", "play some jazz", "로파이 틀어줘",
    "음악 꺼줘", "볼륨 30으로 해줘", "set the volume to 60", "볼륨 몇이야", "what's the volume",
    "서울 날씨 어때?", "what's the weather in london right now?", "뉴욕은 지금 몇 시야?",
    "지금 몇시야?", "여기가 어디야?", "프랑스 수도가 어디야", "what is the capital of france?",
    "부엌, 거실, 안방 불 켜줘", "마루에 빨간 불 켜고 재즈 틀어줘",
}

# ============================ prompt rendering ============================
# The tokenizer + chat template come from the base model; merging LoRA leaves
# both untouched, so training and export render identically.
_tok = AutoTokenizer.from_pretrained(paths.BASE_MODEL)
_chat_template = open(os.path.join(paths.BASE_MODEL, "chat_template.jinja")).read()


def prompt_only(user):
    """The rendered prompt up to (and including) the model's generation turn."""
    return _tok.apply_chat_template(
        [{"role": "developer", "content": DEV}, {"role": "user", "content": user}],
        tools=TOOLS, chat_template=_chat_template,
        add_generation_prompt=True, tokenize=False)


def render_calls(user, calls):
    """Render one training row for a tool-calling turn.

    `calls` is a list of (name, args_dict). Returns (prompt, completion) where
    completion is just the assistant's function-call span.
    """
    prompt = prompt_only(user)
    tool_calls = [{"type": "function", "function": {"name": n, "arguments": a}} for n, a in calls]
    full = _tok.apply_chat_template(
        [{"role": "developer", "content": DEV}, {"role": "user", "content": user},
         {"role": "assistant", "tool_calls": tool_calls}],
        tools=TOOLS, chat_template=_chat_template,
        add_generation_prompt=False, tokenize=False)
    return prompt, full[len(prompt):]


def render_text(user, reply):
    """Render one training row for a plain-text refusal turn."""
    prompt = prompt_only(user)
    full = _tok.apply_chat_template(
        [{"role": "developer", "content": DEV}, {"role": "user", "content": user},
         {"role": "assistant", "content": reply}],
        tools=TOOLS, chat_template=_chat_template,
        add_generation_prompt=False, tokenize=False)
    return prompt, full[len(prompt):]