From eed1f3d0c6043bd494424a22443ee198dd56f657 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 12 Sep 2026 12:56:55 +0800 Subject: [PATCH] [Rust Frontend] Support `generation` blocks in HF chat template (#56378) Signed-off-by: Bugen Zhao --- rust/src/chat/src/renderer/hf/generation.rs | 187 ++++++++++++++++++++ rust/src/chat/src/renderer/hf/mod.rs | 17 ++ rust/src/chat/src/renderer/hf/template.rs | 3 + 3 files changed, 207 insertions(+) create mode 100644 rust/src/chat/src/renderer/hf/generation.rs diff --git a/rust/src/chat/src/renderer/hf/generation.rs b/rust/src/chat/src/renderer/hf/generation.rs new file mode 100644 index 0000000000..39e61b5fb6 --- /dev/null +++ b/rust/src/chat/src/renderer/hf/generation.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Support Hugging Face's `{% generation %}` blocks in MiniJinja chat templates. +//! +//! Transformers uses these blocks to track assistant-token masks via scoped call +//! blocks. Inference only needs their rendered text. MiniJinja lacks this tag, so +//! we rewrite `{% generation %}...{% endgeneration %}` to +//! `{% call __hf_generation() %}...{% endcall %}` before compiling the template. +//! The registered `__hf_generation` callback renders and returns `caller()`, +//! preserving the block's local scope just as Transformers' AssistantTracker does. +//! +//! We use MiniJinja's lexer to identify statement tags and replace only their +//! keyword spans. This preserves delimiters and whitespace controls (`-` / `+`), +//! while keeping literals, comments, and raw blocks intact. Original `generation` +//! and `call` pairings are checked before rewriting both to call blocks. + +use minijinja::machinery::{Token, tokenize}; +use minijinja::value::{Kwargs, Value}; +use minijinja::{Error, ErrorKind, State}; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum CallBlockKind { + Call, + Generation, +} + +/// Rewrite Hugging Face's assistant-mask blocks as scoped MiniJinja call blocks. +pub(super) fn rewrite_generation_blocks(source: String) -> Result { + let tokens = tokenize(&source, false, Default::default(), Default::default()) + .collect::, _>>()?; + // Track native calls too: rewriting makes both block kinds use `endcall`, + // so validate the original pairing before that distinction is lost. + let mut blocks = Vec::new(); + let mut replacements = Vec::new(); + + for (index, pair) in tokens.windows(2).enumerate() { + let [(Token::BlockStart, _), (Token::Ident(name), span)] = pair else { + continue; + }; + match *name { + "call" => blocks.push(CallBlockKind::Call), + "endcall" if blocks.pop() == Some(CallBlockKind::Generation) => { + return Err(Error::new( + ErrorKind::SyntaxError, + "expected endgeneration, got endcall", + )); + } + "generation" | "endgeneration" + if matches!(tokens.get(index + 2), Some((Token::BlockEnd, _))) => + { + let replacement = if *name == "generation" { + blocks.push(CallBlockKind::Generation); + "call __hf_generation()" + } else { + if blocks.pop() != Some(CallBlockKind::Generation) { + return Err(Error::new( + ErrorKind::SyntaxError, + "unexpected endgeneration", + )); + } + "endcall" + }; + replacements.push(( + span.start_offset as usize..span.end_offset as usize, + replacement, + )); + } + _ => {} + } + } + if blocks.contains(&CallBlockKind::Generation) { + return Err(Error::new( + ErrorKind::SyntaxError, + "unclosed generation block", + )); + } + + // Replace only keyword spans, preserving delimiters and their whitespace controls. + let mut source = source; + for (span, replacement) in replacements.into_iter().rev() { + source.replace_range(span, replacement); + } + Ok(source) +} + +/// Render the caller with the same local scope as Transformers' AssistantTracker. +pub(super) fn render_generation(state: &State, kwargs: Kwargs) -> Result { + let caller: Value = kwargs.get("caller")?; + kwargs.assert_all_used()?; + caller.call(state, &[]) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use minijinja::Environment; + + use super::*; + + #[test] + fn generation_blocks_reject_invalid_syntax_and_mismatched_end_tags() { + let env = Environment::new(); + for template in [ + "{% generation %}unclosed", + "{% endgeneration %}", + "{% generation extra %}{% endgeneration %}", + "{% generation %}{% endgeneration extra %}", + "{% generation %}{% endcall %}", + "{% call f() %}{% endgeneration %}", + "{% generation %}{% endcall %}{% call f() %}{% endgeneration %}", + "{% generation %}{% if true %}{% endgeneration %}{% endif %}", + ] { + let result = rewrite_generation_blocks(template.to_string()) + .and_then(|source| env.template_from_str(&source).map(|_| ())); + assert!(result.is_err(), "accepted invalid template: {template}"); + } + } + + #[test] + fn rewrite_generation_blocks_preserves_hf_rendering_semantics() { + let cases = [ + ( + "whitespace", + "前 {%- generation -%} 中 {%- endgeneration -%} 后", + ), + ( + "keep_newlines", + "A{%+ generation +%}\nB{%+ endgeneration +%}\nC", + ), + ("literal", "{{ '{% generation %}' }}"), + ("raw", "{% raw %}{%- generation -%}{% endraw %}"), + ("comment", "A{# {% generation %} #}B"), + ( + "scope", + "{% set x = 'outer' %}{% generation %}{% set x = 'inner' %}{{ x }}{% endgeneration %}{{ x }}", + ), + ( + "nested_calls", + "{% macro wrap() %}[{{ caller() }}]{% endmacro %}{% generation %}A{% call wrap() %}{% generation %}B{% endgeneration %}{% endcall %}C{% endgeneration %}", + ), + ]; + let results = cases.map(|(name, source)| { + let template = rewrite_generation_blocks(source.to_string()).unwrap(); + let mut env = Environment::new(); + env.set_trim_blocks(true); + env.set_lstrip_blocks(true); + env.add_function("__hf_generation", render_generation); + let rendered = + env.template_from_str(&template).unwrap().render(Value::UNDEFINED).unwrap(); + (name, rendered) + }); + expect![[r#" + [ + ( + "whitespace", + "前中后", + ), + ( + "keep_newlines", + "A\nB\nC", + ), + ( + "literal", + "{% generation %}", + ), + ( + "raw", + "{%- generation -%}", + ), + ( + "comment", + "AB", + ), + ( + "scope", + "innerouter", + ), + ( + "nested_calls", + "A[B]C", + ), + ] + "#]] + .assert_debug_eq(&results); + } +} diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 7461acf295..5e7ba75e0a 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -25,6 +25,7 @@ use crate::{ mod error; mod format; +mod generation; mod template; mod tojson; @@ -642,6 +643,22 @@ mod tests { .map_err(|_| unreachable!("HF renderer should return text prompt")) } + #[test] + fn generation_blocks_allow_content_format_detection_and_request_overrides() { + let template = "{% for message in messages %}{% generation %}{% for part in message.content %}{{ part.text }}{% endfor %}{% endgeneration %}{% endfor %}"; + let mut request = sample_request(vec![ChatMessage::user("hello")]); + let default = render(Some(template), &request).unwrap(); + request.chat_options.chat_template = Some(template.to_string()); + let overridden = render(Some("unused"), &request).unwrap(); + expect![[r#" + ( + "hello", + "hello", + ) + "#]] + .assert_debug_eq(&(default, overridden)); + } + fn render_mm( template: &str, request: &ChatRequest, diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index b7de0172f3..b890783701 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -20,6 +20,7 @@ use super::error::TemplateError; use super::format::{ ChatTemplateContentFormat, ChatTemplateContentFormatOption, detect_chat_template_content_format, }; +use super::generation::{render_generation, rewrite_generation_blocks}; use super::tojson::hf_tojson_filter; use crate::renderer::hf::{TemplateMessage, TemplateTool}; @@ -36,6 +37,7 @@ fn build_environment(template: String) -> Result> { env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback); env.add_filter("tojson", hf_tojson_filter); + env.add_function("__hf_generation", render_generation); Ok(env) } @@ -107,6 +109,7 @@ pub(super) struct CompiledChatTemplate { impl CompiledChatTemplate { /// Compile the given chat template string into a [`CompiledChatTemplate`]. pub fn new(template: String, content_format: ChatTemplateContentFormatOption) -> Result { + let template = rewrite_generation_blocks(template)?; let content_format = match content_format { ChatTemplateContentFormatOption::Auto => detect_chat_template_content_format(&template), ChatTemplateContentFormatOption::String => ChatTemplateContentFormat::String,