135 lines
7.0 KiB
Python
135 lines
7.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate the PRODUCTION fixed chat_template.jinja for SmolLM3-3B.
|
|
Uses a hybrid approach: raw strings for Jinja2, concat for special tokens."""
|
|
from transformers import AutoTokenizer
|
|
tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B")
|
|
|
|
THINK_S = tok.decode([128002])
|
|
THINK_E = tok.decode([128003])
|
|
RESP_S = tok.decode([128013])
|
|
RESP_E = tok.decode([128014])
|
|
TC_S = tok.decode([128015])
|
|
TC_E = tok.decode([128016])
|
|
|
|
# Build the template as a list of chunks
|
|
# Key: any line with special tokens gets its own append with string concat
|
|
T = []
|
|
|
|
# ─── System header (no special tokens) ───
|
|
T.append('{# ───── defaults ───── #}\n')
|
|
T.append('{%- if enable_thinking is not defined -%}\n')
|
|
T.append('{%- set enable_thinking = true -%}\n')
|
|
T.append('{%- endif -%}\n\n')
|
|
T.append('{# ───── reasoning mode ───── #}\n')
|
|
T.append('{%- if enable_thinking -%}\n')
|
|
T.append(' {%- set reasoning_mode = "/think" -%}\n')
|
|
T.append('{%- else -%}\n')
|
|
T.append(' {%- set reasoning_mode = "/no_think" -%}\n')
|
|
T.append('{%- endif -%}\n\n')
|
|
T.append('{# ───── header (system message) ───── #}\n')
|
|
T.append('{{- "<|im_start|>system\\n" -}}\n\n')
|
|
T.append('{%- if messages[0].role == "system" -%}\n')
|
|
T.append(' {%- set system_message = messages[0].content -%}\n')
|
|
T.append(' {%- if "/no_think" in system_message -%}\n')
|
|
T.append(' {%- set reasoning_mode = "/no_think" -%}\n')
|
|
T.append(' {%- elif "/think" in system_message -%}\n')
|
|
T.append(' {%- set reasoning_mode = "/think" -%}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append(' {%- set custom_instructions = system_message.replace("/no_think", "").replace("/think", "").rstrip() -%}\n')
|
|
T.append('{%- endif -%}\n\n')
|
|
T.append('{%- if "/system_override" in system_message -%}\n')
|
|
T.append(' {{- custom_instructions.replace("/system_override", "").rstrip() -}}\n')
|
|
T.append('{%- else -%}\n')
|
|
T.append(' {{- "## Metadata\\n\\n" -}}\n')
|
|
T.append(' {{- "Knowledge Cutoff Date: June 2025\\n" -}}\n')
|
|
T.append(' {%- set today = strftime_now("%d %B %Y") -%}\n')
|
|
T.append(' {{- "Today Date: " ~ today ~ "\\n" -}}\n')
|
|
T.append(' {{- "Reasoning Mode: " + reasoning_mode + "\\n\\n" -}}\n\n')
|
|
T.append(' {{- "## Custom Instructions\\n\\n" -}}\n')
|
|
T.append(' {%- if custom_instructions -%}\n')
|
|
T.append(' {{- custom_instructions + "\\n\\n" -}}\n')
|
|
T.append(' {%- elif reasoning_mode == "/think" -%}\n')
|
|
T.append(' {{- "You are a helpful AI assistant named SmolLM, trained by Hugging Face.\\n\\n" -}}\n')
|
|
T.append(' {%- else -%}\n')
|
|
T.append(' {{- "You are a helpful AI assistant named SmolLM, trained by Hugging Face.\\n\\n" -}}\n')
|
|
T.append(' {%- endif -%}\n\n')
|
|
T.append(' {%- if xml_tools or python_tools or tools -%}\n')
|
|
T.append(' {{- "### Tools\\n\\n" -}}\n')
|
|
T.append(' {%- if xml_tools or tools -%}\n')
|
|
T.append(' {%- if tools -%}\n')
|
|
T.append(' {%- set xml_tools = tools -%}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append(' {%- set ns = namespace(xml_tool_string="You may call one or more functions to assist with the user query.\\nYou are provided with function signatures within <tools></tools> XML tags:\\n\\n<tools>\\n") -%}\n')
|
|
T.append(' {%- for tool in xml_tools[:] -%}\n')
|
|
T.append(' {%- set ns.xml_tool_string = ns.xml_tool_string ~ (tool | tojson) ~ "\\n" -%}\n')
|
|
T.append(' {%- endfor -%}\n')
|
|
|
|
# Tool format instruction - has special tokens
|
|
T.append(' {%- set xml_tool_string = ns.xml_tool_string + "</tools>\\n\\nFor each function call, return a json object with function name and arguments within ' + TC_S + ' XML tags:\\n' + TC_S + '\\n{\\"name\\": <function-name>, \\"arguments\\": <args-json-object>}\\n' + TC_E + '" -%}\n')
|
|
|
|
T.append(' {{- xml_tool_string -}}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append(' {%- if python_tools -%}\n')
|
|
T.append(' {%- set ns = namespace(python_tool_string="You may call one or more functions as python tools.\\n<tools>\\n") -%}\n')
|
|
T.append(' {%- for tool in python_tools[:] -%}\n')
|
|
T.append(' {%- set ns.python_tool_string = ns.python_tool_string ~ (tool | string) ~ "\\n" -%}\n')
|
|
T.append(' {%- endfor -%}\n')
|
|
T.append(' {%- set python_tool_string = ns.python_tool_string + "</tools>\\n\\nThe state persists between code executions." -%}\n')
|
|
T.append(' {{- python_tool_string -}}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append(' {{- "\\n\\n" -}}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append('{%- endif -%}\n')
|
|
T.append('{{- "<|im_end|>\\n" -}}\n\n')
|
|
|
|
# ─── Main loop ───
|
|
T.append('{# ───── main loop ───── #}\n')
|
|
T.append('{%- for message in messages -%}\n')
|
|
T.append(' {%- if message.role == "user" -%}\n')
|
|
T.append(' {{ "<|im_start|>user\\n" + message.content + "<|im_end|>\\n" }}\n')
|
|
T.append(' {%- elif message.role == "assistant" -%}\n')
|
|
T.append(' {% generation %}\n')
|
|
T.append(' {%- if message.tool_calls -%}\n')
|
|
T.append(' {%- set ns = namespace(tc_text="") -%}\n')
|
|
T.append(' {%- for tc in message.tool_calls -%}\n')
|
|
|
|
# FIX: Render tool calls with TC_S/TC_E tokens using ~ (Jinja2 concat)
|
|
T.append(' {%- set ns.tc_text = ns.tc_text ~ "' + TC_S + '\\n{\\"name\\": \\"" ~ tc.function.name ~ "\\", \\"arguments\\": " ~ tc.function.arguments ~ "}\\n' + TC_E + '" -%}\n')
|
|
|
|
T.append(' {%- endfor -%}\n')
|
|
T.append(' {{ "<|im_start|>assistant\\n" ~ (message.content if message.content is string else "") ~ ns.tc_text ~ "<|im_end|>\\n" }}\n')
|
|
T.append(' {%- else -%}\n')
|
|
|
|
# FIX: /think = use think tags, /no_think = plain text (was inverted in original)
|
|
T.append(' {%- if reasoning_mode == "/think" -%}\n')
|
|
T.append(' {{ "<|im_start|>assistant\\n' + THINK_S + '\\n" ~ (message.content if message.content is string else "") ~ "\\n' + THINK_E + '<|im_end|>\\n" }}\n')
|
|
T.append(' {%- else -%}\n')
|
|
T.append(' {{ "<|im_start|>assistant\\n" ~ (message.content if message.content is string else "") ~ "<|im_end|>\\n" }}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append(' {% endgeneration %}\n')
|
|
|
|
# FIX: Tool role with RESP_S/RESP_E tokens
|
|
T.append(' {%- elif message.role == "tool" -%}\n')
|
|
T.append(' {{ "<|im_start|>user\\n' + RESP_S + '\\n" ~ (message.content if message.content is string else "") ~ "\\n' + RESP_E + '<|im_end|>\\n" }}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append('{%- endfor -%}\n\n')
|
|
|
|
# ─── Generation prompt ───
|
|
T.append('{# ───── generation prompt ───── #}\n')
|
|
T.append('{%- if add_generation_prompt -%}\n')
|
|
T.append(' {%- if reasoning_mode == "/think" -%}\n')
|
|
T.append(' {{ "<|im_start|>assistant\\n' + THINK_S + '\\n" }}\n')
|
|
T.append(' {%- else -%}\n')
|
|
T.append(' {{ "<|im_start|>assistant\\n" }}\n')
|
|
T.append(' {%- endif -%}\n')
|
|
T.append('{%- endif -%}\n')
|
|
|
|
template = ''.join(T)
|
|
|
|
with open('/root/chat_template.jinja', 'w', encoding='utf-8') as f:
|
|
f.write(template)
|
|
|
|
print("Production template written to /root/chat_template.jinja")
|
|
print(f"Length: {len(template)} bytes")
|