535 lines
19 KiB
Rust
535 lines
19 KiB
Rust
//! Chat template support for tokenizers using Jinja2 templates
|
|
//!
|
|
//! This module provides functionality to apply chat templates to messages,
|
|
//! similar to HuggingFace transformers' apply_chat_template method.
|
|
|
|
use std::{collections::HashMap, fs};
|
|
|
|
use anyhow::{anyhow, Result};
|
|
use minijinja::{
|
|
context,
|
|
machinery::{
|
|
ast::{Expr, Stmt},
|
|
parse, WhitespaceConfig,
|
|
},
|
|
syntax::SyntaxConfig,
|
|
value::Kwargs,
|
|
Environment, Error as MinijinjaError, ErrorKind, Value,
|
|
};
|
|
use serde::Serialize;
|
|
use serde_json::{self, ser::PrettyFormatter, Value as JsonValue};
|
|
|
|
/// Chat template content format
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
|
pub enum ChatTemplateContentFormat {
|
|
/// Content is a simple string
|
|
#[default]
|
|
String,
|
|
/// Content is a list of structured parts (OpenAI format)
|
|
OpenAI,
|
|
}
|
|
|
|
impl std::fmt::Display for ChatTemplateContentFormat {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
Self::String => write!(f, "string"),
|
|
Self::OpenAI => write!(f, "openai"),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Detect the content format expected by a Jinja2 chat template
|
|
///
|
|
/// This implements the same detection logic as SGLang's detect_jinja_template_content_format
|
|
/// which uses AST parsing to look for content iteration patterns.
|
|
///
|
|
/// Returns:
|
|
/// - ChatTemplateContentFormat::OpenAI if template expects structured content (list of parts)
|
|
/// - ChatTemplateContentFormat::String if template expects simple string content
|
|
pub fn detect_chat_template_content_format(template: &str) -> ChatTemplateContentFormat {
|
|
// Use AST-based detection (enabled by default)
|
|
if let Some(format) = detect_format_with_ast(template) {
|
|
return format;
|
|
}
|
|
|
|
// Default to string format if AST parsing fails
|
|
ChatTemplateContentFormat::String
|
|
}
|
|
|
|
/// Flags tracking which OpenAI-style patterns we've seen
|
|
#[derive(Default, Debug, Clone, Copy)]
|
|
struct Flags {
|
|
saw_iteration: bool,
|
|
saw_structure: bool,
|
|
saw_assignment: bool,
|
|
saw_macro: bool,
|
|
}
|
|
|
|
impl Flags {
|
|
fn any(self) -> bool {
|
|
self.saw_iteration || self.saw_structure || self.saw_assignment || self.saw_macro
|
|
}
|
|
}
|
|
|
|
/// Single-pass AST detector with scope tracking
|
|
struct Detector<'a> {
|
|
ast: &'a Stmt<'a>,
|
|
/// Message loop vars currently in scope (e.g., `message`, `m`, `msg`)
|
|
scope: std::collections::VecDeque<String>,
|
|
scope_set: std::collections::HashSet<String>,
|
|
flags: Flags,
|
|
}
|
|
|
|
impl<'a> Detector<'a> {
|
|
fn new(ast: &'a Stmt<'a>) -> Self {
|
|
Self {
|
|
ast,
|
|
scope: std::collections::VecDeque::new(),
|
|
scope_set: std::collections::HashSet::new(),
|
|
flags: Flags::default(),
|
|
}
|
|
}
|
|
|
|
fn run(mut self) -> Flags {
|
|
self.walk_stmt(self.ast);
|
|
self.flags
|
|
}
|
|
|
|
fn push_scope(&mut self, var: String) {
|
|
self.scope.push_back(var.clone());
|
|
self.scope_set.insert(var);
|
|
}
|
|
|
|
fn pop_scope(&mut self) {
|
|
if let Some(v) = self.scope.pop_back() {
|
|
self.scope_set.remove(&v);
|
|
}
|
|
}
|
|
|
|
fn is_var_access(expr: &Expr, varname: &str) -> bool {
|
|
matches!(expr, Expr::Var(v) if v.id == varname)
|
|
}
|
|
|
|
fn is_const_str(expr: &Expr, value: &str) -> bool {
|
|
matches!(expr, Expr::Const(c) if c.value.as_str() == Some(value))
|
|
}
|
|
|
|
fn is_numeric_const(expr: &Expr) -> bool {
|
|
matches!(expr, Expr::Const(c) if c.value.is_number())
|
|
}
|
|
|
|
/// Check if expr is varname.content or varname["content"]
|
|
fn is_var_dot_content(expr: &Expr, varname: &str) -> bool {
|
|
match expr {
|
|
Expr::GetAttr(g) => Self::is_var_access(&g.expr, varname) && g.name == "content",
|
|
Expr::GetItem(g) => {
|
|
Self::is_var_access(&g.expr, varname)
|
|
&& Self::is_const_str(&g.subscript_expr, "content")
|
|
}
|
|
// Unwrap filters/tests that just wrap the same expr
|
|
Expr::Filter(f) => f
|
|
.expr
|
|
.as_ref()
|
|
.is_some_and(|e| Self::is_var_dot_content(e, varname)),
|
|
Expr::Test(t) => Self::is_var_dot_content(&t.expr, varname),
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
/// Check if expr accesses .content on any variable in our scope, or any descendant of it.
|
|
fn is_any_scope_var_content(&self, expr: &Expr) -> bool {
|
|
let mut current_expr = expr;
|
|
loop {
|
|
// Check if current level matches <scopeVar>.content
|
|
if self
|
|
.scope_set
|
|
.iter()
|
|
.any(|v| Self::is_var_dot_content(current_expr, v))
|
|
{
|
|
return true;
|
|
}
|
|
// Walk up the expression tree
|
|
match current_expr {
|
|
Expr::GetAttr(g) => current_expr = &g.expr,
|
|
Expr::GetItem(g) => current_expr = &g.expr,
|
|
_ => return false,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn walk_stmt(&mut self, stmt: &Stmt) {
|
|
// Early exit if we've already detected an OpenAI pattern
|
|
if self.flags.any() {
|
|
return;
|
|
}
|
|
|
|
match stmt {
|
|
Stmt::Template(t) => {
|
|
for ch in &t.children {
|
|
self.walk_stmt(ch);
|
|
}
|
|
}
|
|
// {% for message in messages %}
|
|
Stmt::ForLoop(fl) => {
|
|
// Detect "for X in messages" → push X into scope
|
|
if let Expr::Var(iter) = &fl.iter {
|
|
if iter.id == "messages" {
|
|
if let Expr::Var(target) = &fl.target {
|
|
self.push_scope(target.id.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Also detect "for ... in message.content" or "for ... in content"
|
|
// - Iterating directly over <scopeVar>.content => OpenAI style
|
|
if self.is_any_scope_var_content(&fl.iter) {
|
|
self.flags.saw_iteration = true;
|
|
}
|
|
// - Iterating over a local var named "content"
|
|
if matches!(&fl.iter, Expr::Var(v) if v.id == "content") {
|
|
self.flags.saw_iteration = true;
|
|
}
|
|
|
|
for b in &fl.body {
|
|
self.walk_stmt(b);
|
|
}
|
|
|
|
// Pop scope if we pushed it
|
|
if let Expr::Var(iter) = &fl.iter {
|
|
if iter.id == "messages" && matches!(&fl.target, Expr::Var(_)) {
|
|
self.pop_scope();
|
|
}
|
|
}
|
|
}
|
|
Stmt::IfCond(ic) => {
|
|
self.inspect_expr_for_structure(&ic.expr);
|
|
for b in &ic.true_body {
|
|
self.walk_stmt(b);
|
|
}
|
|
for b in &ic.false_body {
|
|
self.walk_stmt(b);
|
|
}
|
|
}
|
|
Stmt::EmitExpr(e) => {
|
|
self.inspect_expr_for_structure(&e.expr);
|
|
}
|
|
// {% set content = message.content %}
|
|
Stmt::Set(s) => {
|
|
if Self::is_var_access(&s.target, "content")
|
|
&& self.is_any_scope_var_content(&s.expr)
|
|
{
|
|
self.flags.saw_assignment = true;
|
|
}
|
|
}
|
|
Stmt::Macro(m) => {
|
|
// Heuristic: macro that checks type (via `is` test) and also has any loop
|
|
let mut has_type_check = false;
|
|
let mut has_loop = false;
|
|
Self::scan_macro_body(&m.body, &mut has_type_check, &mut has_loop);
|
|
if has_type_check && has_loop {
|
|
self.flags.saw_macro = true;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
fn inspect_expr_for_structure(&mut self, expr: &Expr) {
|
|
if self.flags.saw_structure {
|
|
return;
|
|
}
|
|
|
|
match expr {
|
|
// content[0] or message.content[0]
|
|
Expr::GetItem(gi) => {
|
|
if (matches!(&gi.expr, Expr::Var(v) if v.id == "content")
|
|
|| self.is_any_scope_var_content(&gi.expr))
|
|
&& Self::is_numeric_const(&gi.subscript_expr)
|
|
{
|
|
self.flags.saw_structure = true;
|
|
}
|
|
}
|
|
// content|length or message.content|length
|
|
Expr::Filter(f) => {
|
|
if f.name == "length" {
|
|
if let Some(inner) = &f.expr {
|
|
// Box derefs automatically, so `&**inner` is `&Expr`
|
|
let inner_ref: &Expr = inner;
|
|
let is_content_var = matches!(inner_ref, Expr::Var(v) if v.id == "content");
|
|
if is_content_var || self.is_any_scope_var_content(inner_ref) {
|
|
self.flags.saw_structure = true;
|
|
}
|
|
}
|
|
} else if let Some(inner) = &f.expr {
|
|
let inner_ref: &Expr = inner;
|
|
self.inspect_expr_for_structure(inner_ref);
|
|
}
|
|
}
|
|
// content is sequence/iterable OR message.content is sequence/iterable
|
|
Expr::Test(t) => {
|
|
if t.name == "sequence" || t.name == "iterable" || t.name == "string" {
|
|
if matches!(&t.expr, Expr::Var(v) if v.id == "content")
|
|
|| self.is_any_scope_var_content(&t.expr)
|
|
{
|
|
self.flags.saw_structure = true;
|
|
}
|
|
} else {
|
|
self.inspect_expr_for_structure(&t.expr);
|
|
}
|
|
}
|
|
Expr::GetAttr(g) => {
|
|
// Keep walking; nested expressions can hide structure checks
|
|
self.inspect_expr_for_structure(&g.expr);
|
|
}
|
|
// Handle binary operations like: if (message.content is string) and other_cond
|
|
Expr::BinOp(op) => {
|
|
self.inspect_expr_for_structure(&op.left);
|
|
self.inspect_expr_for_structure(&op.right);
|
|
}
|
|
// Handle unary operations like: if not (message.content is string)
|
|
Expr::UnaryOp(op) => {
|
|
self.inspect_expr_for_structure(&op.expr);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
fn scan_macro_body(body: &[Stmt], has_type_check: &mut bool, has_loop: &mut bool) {
|
|
for s in body {
|
|
if *has_type_check && *has_loop {
|
|
return;
|
|
}
|
|
|
|
match s {
|
|
Stmt::IfCond(ic) => {
|
|
if matches!(&ic.expr, Expr::Test(_)) {
|
|
*has_type_check = true;
|
|
}
|
|
Self::scan_macro_body(&ic.true_body, has_type_check, has_loop);
|
|
Self::scan_macro_body(&ic.false_body, has_type_check, has_loop);
|
|
}
|
|
Stmt::ForLoop(fl) => {
|
|
*has_loop = true;
|
|
Self::scan_macro_body(&fl.body, has_type_check, has_loop);
|
|
}
|
|
Stmt::Template(t) => {
|
|
Self::scan_macro_body(&t.children, has_type_check, has_loop);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// AST-based detection using minijinja's unstable machinery
|
|
/// Single-pass detector with scope tracking
|
|
fn detect_format_with_ast(template: &str) -> Option<ChatTemplateContentFormat> {
|
|
let ast = match parse(
|
|
template,
|
|
"template",
|
|
SyntaxConfig {},
|
|
WhitespaceConfig::default(),
|
|
) {
|
|
Ok(ast) => ast,
|
|
Err(_) => return Some(ChatTemplateContentFormat::String),
|
|
};
|
|
|
|
let flags = Detector::new(&ast).run();
|
|
Some(if flags.any() {
|
|
ChatTemplateContentFormat::OpenAI
|
|
} else {
|
|
ChatTemplateContentFormat::String
|
|
})
|
|
}
|
|
|
|
/// Parameters for chat template application
|
|
#[derive(Default)]
|
|
pub struct ChatTemplateParams<'a> {
|
|
pub add_generation_prompt: bool,
|
|
pub tools: Option<&'a [serde_json::Value]>,
|
|
pub documents: Option<&'a [serde_json::Value]>,
|
|
pub template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
|
|
}
|
|
|
|
/// Custom tojson filter compatible with HuggingFace transformers' implementation.
|
|
///
|
|
/// HuggingFace transformers registers a custom `tojson` filter that accepts additional
|
|
/// keyword arguments beyond what standard Jinja2 provides:
|
|
/// - `ensure_ascii` (bool): Whether to escape non-ASCII characters (ignored in Rust, always UTF-8)
|
|
/// - `indent` (int): Number of spaces for indentation (pretty-printing)
|
|
/// - `separators` (ignored): Custom separators for JSON output
|
|
/// - `sort_keys` (bool): Whether to sort dictionary keys
|
|
///
|
|
/// This is necessary for compatibility with chat templates from HuggingFace Hub models.
|
|
/// See: https://github.com/huggingface/transformers/blob/main/src/transformers/utils/chat_template_utils.py
|
|
fn tojson_filter(value: Value, kwargs: Kwargs) -> std::result::Result<Value, MinijinjaError> {
|
|
let _ensure_ascii: Option<bool> = kwargs.get("ensure_ascii")?;
|
|
let indent: Option<i64> = kwargs.get("indent")?;
|
|
let _separators: Option<Value> = kwargs.get("separators")?;
|
|
let sort_keys: Option<bool> = kwargs.get("sort_keys")?;
|
|
|
|
// Ensure all kwargs are consumed to avoid "unknown keyword argument" errors
|
|
kwargs.assert_all_used()?;
|
|
|
|
let json_value: serde_json::Value = serde_json::to_value(&value).map_err(|e| {
|
|
MinijinjaError::new(
|
|
ErrorKind::InvalidOperation,
|
|
format!("Failed to convert to JSON value: {}", e),
|
|
)
|
|
})?;
|
|
|
|
// Helper to serialize with custom indentation
|
|
fn serialize_with_indent<T: Serialize>(
|
|
value: &T,
|
|
spaces: usize,
|
|
) -> std::result::Result<String, MinijinjaError> {
|
|
let indent_str = vec![b' '; spaces];
|
|
let formatter = PrettyFormatter::with_indent(&indent_str);
|
|
let mut buf = Vec::new();
|
|
let mut serializer = serde_json::Serializer::with_formatter(&mut buf, formatter);
|
|
value.serialize(&mut serializer).map_err(|e| {
|
|
MinijinjaError::new(
|
|
ErrorKind::InvalidOperation,
|
|
format!("Failed to serialize JSON: {}", e),
|
|
)
|
|
})?;
|
|
String::from_utf8(buf).map_err(|e| {
|
|
MinijinjaError::new(
|
|
ErrorKind::InvalidOperation,
|
|
format!("Invalid UTF-8 in JSON output: {}", e),
|
|
)
|
|
})
|
|
}
|
|
|
|
// Serialize with options
|
|
let json_str: std::result::Result<String, MinijinjaError> = {
|
|
let sorted_json;
|
|
let value_to_serialize = if sort_keys.unwrap_or(false) {
|
|
sorted_json = sort_json_keys(&json_value);
|
|
&sorted_json
|
|
} else {
|
|
&json_value
|
|
};
|
|
|
|
if let Some(spaces) = indent {
|
|
if spaces < 0 {
|
|
return Err(MinijinjaError::new(
|
|
ErrorKind::InvalidOperation,
|
|
"indent cannot be negative",
|
|
));
|
|
}
|
|
serialize_with_indent(value_to_serialize, spaces as usize)
|
|
} else {
|
|
serde_json::to_string(value_to_serialize).map_err(|e| {
|
|
MinijinjaError::new(
|
|
ErrorKind::InvalidOperation,
|
|
format!("Failed to serialize JSON: {}", e),
|
|
)
|
|
})
|
|
}
|
|
};
|
|
|
|
json_str.map(Value::from_safe_string)
|
|
}
|
|
|
|
/// Recursively sort all object keys in a JSON value
|
|
fn sort_json_keys(value: &JsonValue) -> JsonValue {
|
|
match value {
|
|
JsonValue::Object(map) => {
|
|
let mut sorted: serde_json::Map<String, JsonValue> = serde_json::Map::new();
|
|
let mut keys: Vec<_> = map.keys().collect();
|
|
keys.sort();
|
|
for key in keys {
|
|
sorted.insert(key.clone(), sort_json_keys(&map[key]));
|
|
}
|
|
JsonValue::Object(sorted)
|
|
}
|
|
JsonValue::Array(arr) => JsonValue::Array(arr.iter().map(sort_json_keys).collect()),
|
|
_ => value.clone(),
|
|
}
|
|
}
|
|
|
|
/// Chat template processor using Jinja2 - simple wrapper like HuggingFace
|
|
pub struct ChatTemplateProcessor {
|
|
template: String,
|
|
}
|
|
|
|
impl ChatTemplateProcessor {
|
|
/// Create a new chat template processor
|
|
pub fn new(template: String) -> Self {
|
|
ChatTemplateProcessor { template }
|
|
}
|
|
|
|
/// Apply the chat template to a list of messages
|
|
///
|
|
/// This mimics the behavior of HuggingFace's apply_chat_template method
|
|
/// but returns the formatted string instead of token IDs.
|
|
/// Messages should be pre-processed into the format expected by the template.
|
|
pub fn apply_chat_template(
|
|
&self,
|
|
messages: &[serde_json::Value],
|
|
params: ChatTemplateParams,
|
|
) -> Result<String> {
|
|
let mut env = Environment::new();
|
|
|
|
// Register the template
|
|
env.add_template("chat", &self.template)
|
|
.map_err(|e| anyhow!("Failed to add template: {}", e))?;
|
|
|
|
// Enable Python method compatibility (e.g., str.startswith, str.endswith)
|
|
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
|
|
|
|
// Register custom tojson filter compatible with HuggingFace transformers
|
|
// This overrides minijinja's built-in tojson to support additional kwargs
|
|
// like ensure_ascii, separators, and sort_keys that HuggingFace templates use
|
|
env.add_filter("tojson", tojson_filter);
|
|
|
|
// Get the template
|
|
let tmpl = env
|
|
.get_template("chat")
|
|
.map_err(|e| anyhow!("Failed to get template: {}", e))?;
|
|
|
|
// Convert messages to minijinja::Value (messages already processed by router)
|
|
let minijinja_messages: Vec<Value> = messages.iter().map(Value::from_serialize).collect();
|
|
|
|
let base_context = context! {
|
|
messages => &minijinja_messages,
|
|
add_generation_prompt => params.add_generation_prompt,
|
|
tools => params.tools,
|
|
documents => params.documents,
|
|
};
|
|
|
|
// Merge with template_kwargs if provided
|
|
let ctx = if let Some(kwargs) = params.template_kwargs {
|
|
context! {
|
|
..base_context,
|
|
..Value::from_serialize(kwargs)
|
|
}
|
|
} else {
|
|
base_context
|
|
};
|
|
|
|
// Render the template
|
|
let rendered = tmpl
|
|
.render(&ctx)
|
|
.map_err(|e| anyhow!("Failed to render template: {}", e))?;
|
|
|
|
Ok(rendered)
|
|
}
|
|
}
|
|
|
|
/// Load chat template from tokenizer config JSON
|
|
pub fn load_chat_template_from_config(config_path: &str) -> Result<Option<String>> {
|
|
let content = fs::read_to_string(config_path)?;
|
|
let config: serde_json::Value = serde_json::from_str(&content)?;
|
|
|
|
// Look for chat_template in the config
|
|
if let Some(template) = config.get("chat_template") {
|
|
if let Some(template_str) = template.as_str() {
|
|
return Ok(Some(template_str.to_string()));
|
|
}
|
|
}
|
|
|
|
Ok(None)
|
|
}
|