use std::sync::Arc;
use super::tokcfg::{ChatTemplate, fromjson, raise_exception, strftime_now, tojson};
use super::{ContextMixins, HfTokenizerConfigJsonFormatter, JinjaEnvironment};
use either::Either;
use minijinja::{Environment, Value, context};
use serde_json::json;
fn render_default_probe(env: &Environment, messages: serde_json::Value) -> String {
let ctx = context! {
messages => messages,
add_generation_prompt => false,
};
env.get_template("default")
.and_then(|t| t.render(&ctx))
.unwrap_or_default()
}
fn detect_content_array_usage(env: &Environment) -> bool {
let out_array = render_default_probe(
env,
json!([{"role": "user", "content": [{"type": "text", "text": "template_test"}]}]),
);
let out_string =
render_default_probe(env, json!([{"role": "user", "content": "template_test"}]));
out_array.contains("template_test") && !out_string.contains("template_test")
}
fn detect_image_placeholder_template(env: &Environment) -> Option<&'static str> {
let src = env
.get_template("default")
.ok()
.map(|t| t.source().to_string())
.unwrap_or_default();
if src.contains("<|end|>") && src.contains("<|assistant|>") {
return Some("<|image_{n}|>");
}
if src.contains("USER:") && src.contains("ASSISTANT:") {
return Some("<image>");
}
if detect_passthrough_template(env) {
return Some("");
}
None
}
fn detect_passthrough_template(env: &Environment) -> bool {
const PROBE: &str = "\u{1}dynamo_passthrough_probe\u{1}";
let out_string = render_default_probe(env, json!([{"role": "user", "content": PROBE}]));
if out_string.trim() != PROBE {
return false;
}
let out_mixed = render_default_probe(
env,
json!([{"role": "user", "content": [{"type": "text", "text": PROBE}, {"type": "image"}]}]),
);
out_mixed.contains('[') && out_mixed.contains("type")
}
fn remove_known_non_jinja2_tags(template: &str) -> String {
template
.replace("{% generation %}", "")
.replace("{% endgeneration %}", "")
}
fn normalize_dict_method_calls(template: &str) -> String {
let mut out = String::with_capacity(template.len());
let mut i = 0;
while i < template.len() {
if template[i..].starts_with("{#") {
let Some(end) = find_tag_end(template, i + 2, "#}") else {
out.push_str(&template[i..]);
break;
};
out.push_str(&template[i..end]);
i = end;
} else if template[i..].starts_with("{{") {
let Some(end) = find_tag_end(template, i + 2, "}}") else {
out.push_str(&template[i..]);
break;
};
out.push_str("{{");
out.push_str(&normalize_jinja_code_segment(&template[i + 2..end - 2]));
out.push_str("}}");
i = end;
} else if template[i..].starts_with("{%") {
let Some(end) = find_tag_end(template, i + 2, "%}") else {
out.push_str(&template[i..]);
break;
};
let inner = &template[i + 2..end - 2];
if is_jinja_block_name(inner, "raw") {
if let Some(raw_end) = find_raw_block_end(template, end) {
out.push_str(&template[i..raw_end]);
i = raw_end;
} else {
out.push_str(&template[i..]);
break;
}
} else {
out.push_str("{%");
out.push_str(&normalize_jinja_code_segment(inner));
out.push_str("%}");
i = end;
}
} else {
let ch = template[i..].chars().next().expect("valid char boundary");
out.push(ch);
i += ch.len_utf8();
}
}
out
}
fn find_tag_end(template: &str, start: usize, close: &str) -> Option<usize> {
template[start..]
.find(close)
.map(|relative| start + relative + close.len())
}
fn find_raw_block_end(template: &str, start: usize) -> Option<usize> {
let mut i = start;
while let Some(relative_open) = template[i..].find("{%") {
let open = i + relative_open;
let end = find_tag_end(template, open + 2, "%}")?;
if is_jinja_block_name(&template[open + 2..end - 2], "endraw") {
return Some(end);
}
i = end;
}
None
}
fn is_jinja_block_name(inner: &str, name: &str) -> bool {
let trimmed = inner.trim_start();
let trimmed = trimmed
.strip_prefix('-')
.or_else(|| trimmed.strip_prefix('+'))
.unwrap_or(trimmed)
.trim_start();
let Some(rest) = trimmed.strip_prefix(name) else {
return false;
};
rest.chars()
.next()
.is_none_or(|ch| ch.is_whitespace() || ch == '-' || ch == '+')
}
fn normalize_jinja_code_segment(segment: &str) -> String {
let mut out = String::with_capacity(segment.len());
let mut i = 0;
let mut quote: Option<char> = None;
let mut escaped = false;
while i < segment.len() {
let ch = segment[i..].chars().next().expect("valid char boundary");
if let Some(q) = quote {
out.push(ch);
if escaped {
escaped = false;
} else if ch == '\\' {
escaped = true;
} else if ch == q {
quote = None;
}
i += ch.len_utf8();
} else if ch == '\'' || ch == '"' {
quote = Some(ch);
out.push(ch);
i += ch.len_utf8();
} else if segment[i..].starts_with(".items()") {
out.push_str("|items");
i += ".items()".len();
} else {
out.push(ch);
i += ch.len_utf8();
}
}
out
}
fn is_gemma4_reasoning_field_template_source(source: &str) -> bool {
source.contains("<|channel>thought")
&& source.contains("<|tool_call>call:")
&& (source.contains("message.get('reasoning')")
|| source.contains("message.get(\"reasoning\")"))
&& !source.contains("reasoning_content")
}
fn adapt_gemma4_reasoning_template_source(source: &str) -> String {
const OLD_REASONING_BLOCK: &str = "{%- if message.get('reasoning') and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
{{- '<|channel>thought\\n' + message['reasoning'] + '\\n<channel|>'}}
{%- endif -%}";
const NEW_REASONING_BLOCK: &str = "{%- set dyn_gemma4_reasoning_value = message.get('reasoning') or message.get('reasoning_content') -%}
{%- set dyn_gemma4_reasoning_segments = [] -%}
{%- if dyn_gemma4_reasoning_value is string -%}
{%- set dyn_gemma4_reasoning_segments = [dyn_gemma4_reasoning_value] -%}
{%- elif dyn_gemma4_reasoning_value is sequence -%}
{%- set dyn_gemma4_reasoning_segments = dyn_gemma4_reasoning_value -%}
{%- endif -%}
{%- if dyn_gemma4_reasoning_value and not message.get('tool_calls') -%}
{%- for dyn_gemma4_reasoning_segment in dyn_gemma4_reasoning_segments -%}
{%- if dyn_gemma4_reasoning_segment -%}
{{- '<|channel>thought\\n' + dyn_gemma4_reasoning_segment + '\\n<channel|>'}}
{%- endif -%}
{%- endfor -%}
{%- endif -%}";
const OLD_TOOL_LOOP_START: &str = "{%- for tool_call in message['tool_calls'] -%}
{%- set function = tool_call['function'] -%}";
const NEW_TOOL_LOOP_START: &str = "{%- for tool_call in message['tool_calls'] -%}
{%- set dyn_gemma4_reasoning_segment = dyn_gemma4_reasoning_segments[loop.index0] | default('') -%}
{%- if dyn_gemma4_reasoning_segment -%}
{{- '<|channel>thought\\n' + dyn_gemma4_reasoning_segment + '\\n<channel|>'}}
{%- endif -%}
{%- set function = tool_call['function'] -%}";
const OLD_TOOL_LOOP_END: &str = "{{- '}<tool_call|>' -}}
{%- endfor -%}";
const NEW_TOOL_LOOP_END: &str = "{{- '}<tool_call|>' -}}
{%- endfor -%}
{%- set dyn_gemma4_trailing_reasoning = dyn_gemma4_reasoning_segments[message['tool_calls'] | length] | default('') -%}
{%- if dyn_gemma4_trailing_reasoning -%}
{{- '<|channel>thought\\n' + dyn_gemma4_trailing_reasoning + '\\n<channel|>'}}
{%- endif -%}";
if !source.contains(OLD_REASONING_BLOCK)
|| !source.contains(OLD_TOOL_LOOP_START)
|| !source.contains(OLD_TOOL_LOOP_END)
{
return source.to_string();
}
source
.replace(OLD_REASONING_BLOCK, NEW_REASONING_BLOCK)
.replace(OLD_TOOL_LOOP_START, NEW_TOOL_LOOP_START)
.replace(OLD_TOOL_LOOP_END, NEW_TOOL_LOOP_END)
}
fn normalize_chat_template_source(source: &str) -> String {
let source = normalize_dict_method_calls(&remove_known_non_jinja2_tags(source));
if is_gemma4_reasoning_field_template_source(&source) {
adapt_gemma4_reasoning_template_source(&source)
} else {
source
}
}
impl JinjaEnvironment {
fn env(self) -> Environment<'static> {
self.env
}
}
impl Default for JinjaEnvironment {
fn default() -> Self {
let mut env = Environment::new();
env.set_lstrip_blocks(true);
env.set_trim_blocks(true);
JinjaEnvironment { env }
}
}
impl HfTokenizerConfigJsonFormatter {
#[cfg(test)]
pub fn new(config: ChatTemplate, mixins: ContextMixins) -> anyhow::Result<Self> {
Self::with_options(config, mixins, true)
}
pub fn with_options(
config: ChatTemplate,
mixins: ContextMixins,
exclude_tools_when_tool_choice_none: bool,
) -> anyhow::Result<Self> {
let mut env = JinjaEnvironment::default().env();
let chat_template = config.chat_template.as_ref().ok_or(anyhow::anyhow!(
"chat_template field is required in the tokenizer_config.json file"
))?;
env.add_filter("length", |value: Value| -> usize {
use minijinja::value::ValueKind;
match value.kind() {
ValueKind::Undefined | ValueKind::None => 0,
_ => value.len().unwrap_or(0),
}
});
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
env.add_filter("tojson", tojson);
env.add_filter("fromjson", fromjson);
env.add_filter("from_json", fromjson);
env.add_function("raise_exception", raise_exception);
env.add_function("strftime_now", strftime_now);
let mut supports_add_generation_prompt = None;
match &chat_template.0 {
Either::Left(x) => {
if x.contains("add_generation_prompt") {
tracing::debug!(
"Chat template contains `add_generation_prompt` key. This model supports add_generation_prompt."
);
supports_add_generation_prompt = Some(true);
}
let template_cleaned = normalize_chat_template_source(x);
env.add_template_owned("default", template_cleaned.clone())?;
env.add_template_owned("tool_use", template_cleaned)?;
}
Either::Right(map) => {
for t in map {
for (k, v) in t.iter() {
if v.contains("add_generation_prompt") {
match supports_add_generation_prompt {
Some(true) | None => {
tracing::debug!(
"Chat template contains `add_generation_prompt` key. This model supports add_generation_prompt."
);
supports_add_generation_prompt = Some(true);
}
Some(false) => {
tracing::warn!(
"Not all templates contain `add_generation_prompt` key. This model does not support add_generation_prompt."
);
}
}
} else {
supports_add_generation_prompt = Some(false);
}
let template_cleaned = normalize_chat_template_source(v);
env.add_template_owned(k.to_string(), template_cleaned)?;
}
}
if env.templates().count() == 0 {
anyhow::bail!(
"Chat template does not contain a `tool_use` or `default` key. Please ensure it contains at least a `default` key, although `tool_use` should be specified for using tools."
);
}
}
}
let requires_content_arrays = detect_content_array_usage(&env);
let image_placeholder_template = if requires_content_arrays {
None
} else {
detect_image_placeholder_template(&env)
};
let template_handles_reasoning = |name: &str| -> bool {
env.templates()
.find(|(n, _)| *n == name)
.map(|(_, tmpl)| tmpl.source().contains("reasoning_content"))
.unwrap_or(false)
};
let default_template_handles_reasoning = template_handles_reasoning("default");
let tool_use_template_handles_reasoning = template_handles_reasoning("tool_use");
let template_handles_args_string = |name: &str| -> bool {
env.templates()
.find(|(n, _)| *n == name)
.map(|(_, tmpl)| tmpl.source().contains("arguments is string"))
.unwrap_or(false)
};
let default_template_handles_tool_calls_arguments_string =
template_handles_args_string("default");
let tool_use_template_handles_tool_calls_arguments_string =
template_handles_args_string("tool_use");
Ok(HfTokenizerConfigJsonFormatter {
env,
config,
mixins: Arc::new(mixins),
supports_add_generation_prompt: supports_add_generation_prompt.unwrap_or(false),
requires_content_arrays,
exclude_tools_when_tool_choice_none,
default_template_handles_reasoning,
tool_use_template_handles_reasoning,
image_placeholder_template,
default_template_handles_tool_calls_arguments_string,
tool_use_template_handles_tool_calls_arguments_string,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn env_with_default(src: &str) -> Environment<'static> {
let mut env = JinjaEnvironment::default().env();
env.add_template_owned("default", src.to_string()).unwrap();
env
}
#[test]
fn test_remove_known_non_jinja2_tags() {
let template =
"USER: {{ message }} ASSISTANT: {% generation %}Reply here{% endgeneration %}";
let result = remove_known_non_jinja2_tags(template);
assert_eq!(result, "USER: {{ message }} ASSISTANT: Reply here");
}
#[test]
fn test_remove_known_non_jinja2_tags_preserves_standard_tags() {
let template = "{% for item in items %}{{ item }}{% endfor %}";
let result = remove_known_non_jinja2_tags(template);
assert_eq!(result, template);
}
#[test]
fn test_remove_known_non_jinja2_tags_multiple() {
let template = "Start {% generation %}Part 1{% endgeneration %} middle {% generation %}Part 2{% endgeneration %}";
let result = remove_known_non_jinja2_tags(template);
assert_eq!(result, "Start Part 1 middle Part 2");
}
#[test]
fn test_detect_nemotron_parse_passthrough_template() {
let env =
env_with_default("{% for message in messages %}{{ message['content'] }}{% endfor %}");
assert!(
detect_passthrough_template(&env),
"pure pass-through template should be detected"
);
assert!(
!detect_content_array_usage(&env),
"pass-through template renders string content fine, so does not require arrays"
);
assert_eq!(
detect_image_placeholder_template(&env),
Some(""),
"pass-through template should flatten images to an empty placeholder"
);
}
#[test]
fn test_decorated_template_is_not_passthrough() {
let env = env_with_default(
"{% for message in messages %}<|{{ message['role'] }}|>{{ message['content'] }}<|end|>{% endfor %}{% if add_generation_prompt %}<|assistant|>{% endif %}",
);
assert!(
!detect_passthrough_template(&env),
"template that wraps content in role markers is not pass-through"
);
}
#[test]
fn test_string_passthrough_with_native_array_branch_is_not_passthrough() {
let env = env_with_default(
"{% for message in messages %}{% if message['content'] is string %}{{ message['content'] }}{% else %}{% for part in message['content'] %}{% if part['type'] == 'image' %}<image>{% else %}{{ part['text'] }}{% endif %}{% endfor %}{% endif %}{% endfor %}",
);
assert!(
!detect_passthrough_template(&env),
"template with a native content-array branch is not pass-through"
);
assert_eq!(
detect_image_placeholder_template(&env),
None,
"template that natively renders image markers must keep the content array"
);
}
#[test]
fn test_normalize_dict_method_calls_rewrites_items_method() {
let template = "{% for k, v in tool.parameters.properties.items() %}{{ k }}{% endfor %}";
let result = normalize_dict_method_calls(template);
assert_eq!(
result,
"{% for k, v in tool.parameters.properties|items %}{{ k }}{% endfor %}"
);
}
#[test]
fn test_normalize_dict_method_calls_rewrites_expression_items_method() {
let template = "{{ tool.parameters.properties.items() }}";
let result = normalize_dict_method_calls(template);
assert_eq!(result, "{{ tool.parameters.properties|items }}");
}
#[test]
fn test_normalize_dict_method_calls_preserves_literal_text() {
let template = "Do not rewrite literal .items() text.";
let result = normalize_dict_method_calls(template);
assert_eq!(result, template);
}
#[test]
fn test_normalize_dict_method_calls_preserves_comments_raw_and_strings() {
let template = concat!(
"{# comment .items() #}",
"{% raw %}{{ tool.parameters.properties.items() }}{% endraw %}",
"{{ '.items()' }}",
"{{ \".items()\" }}",
);
let result = normalize_dict_method_calls(template);
assert_eq!(result, template);
}
#[test]
fn test_normalize_dict_method_calls_avoids_schema_items_collision() {
let template = normalize_dict_method_calls(
"{% for param_name, param_spec in tool.parameters.properties.items() %}{{ param_name }}={{ param_spec.type }};{% endfor %}",
);
let mut env = Environment::new();
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
env.add_template("t", &template).unwrap();
let tool = json!({
"parameters": {
"properties": {
"items": {"type": "array", "items": {"type": "object"}},
"message": {"type": "string"}
}
}
});
let out = env
.get_template("t")
.unwrap()
.render(context! { tool => tool })
.unwrap();
assert!(out.contains("items=array;"));
assert!(out.contains("message=string;"));
}
#[test]
fn test_minijinja_parses_midchain_dotted_integer_lookup() {
let chat_template: ChatTemplate = serde_json::from_value(serde_json::json!({
"chat_template": r#"{{ m.content.0.type }} {{ "1.5.10" }}"#,
}))
.unwrap();
let formatter =
HfTokenizerConfigJsonFormatter::new(chat_template, ContextMixins::new(&[])).unwrap();
let result = formatter
.env
.get_template("default")
.unwrap()
.render(context! {
m => json!({
"content": [
{
"type": "tool_reference"
}
]
})
})
.unwrap();
assert_eq!(result, "tool_reference 1.5.10");
}
}