diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d54c2be5b..5ddd480fe 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -59,7 +59,8 @@ repos: python/sglang/srt/grpc/.*_pb2\.py| python/sglang/srt/grpc/.*_pb2_grpc\.py| python/sglang/srt/grpc/.*_pb2\.pyi| - python/sglang/srt/grpc/.*_pb2_grpc\.pyi + python/sglang/srt/grpc/.*_pb2_grpc\.pyi| + sgl-router/src/tokenizer/chat_template\.rs )$ - repo: https://github.com/pre-commit/mirrors-clang-format rev: v18.1.8 diff --git a/sgl-router/src/tokenizer/chat_template.rs b/sgl-router/src/tokenizer/chat_template.rs index 9f33921c5..0bc169b0b 100644 --- a/sgl-router/src/tokenizer/chat_template.rs +++ b/sgl-router/src/tokenizer/chat_template.rs @@ -13,9 +13,11 @@ use minijinja::{ parse, WhitespaceConfig, }, syntax::SyntaxConfig, - Environment, Value, + value::Kwargs, + Environment, Error as MinijinjaError, ErrorKind, Value, }; -use serde_json; +use serde::Serialize; +use serde_json::{self, ser::PrettyFormatter, Value as JsonValue}; /// Chat template content format #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -349,6 +351,104 @@ pub struct ChatTemplateParams<'a> { pub template_kwargs: Option<&'a HashMap>, } +/// 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 { + let _ensure_ascii: Option = kwargs.get("ensure_ascii")?; + let indent: Option = kwargs.get("indent")?; + let _separators: Option = kwargs.get("separators")?; + let sort_keys: Option = 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( + value: &T, + spaces: usize, + ) -> std::result::Result { + 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 = { + 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 = 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, @@ -379,6 +479,11 @@ impl ChatTemplateProcessor { // 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") diff --git a/sgl-router/tests/chat_template_integration.rs b/sgl-router/tests/chat_template_integration.rs index 3ca166ed7..805175755 100644 --- a/sgl-router/tests/chat_template_integration.rs +++ b/sgl-router/tests/chat_template_integration.rs @@ -265,6 +265,77 @@ fn test_empty_messages_template() { assert_eq!(result, ""); } +/// Test that tojson filter accepts ensure_ascii kwarg (HuggingFace compatibility) +/// This is the fix for: "unknown keyword argument 'ensure_ascii'" +#[test] +fn test_tojson_with_ensure_ascii() { + // Template that uses tojson(ensure_ascii=False) like HuggingFace templates do + let template = r#" +{%- for message in messages -%} +{{ message.role }}: {{ message.content }} +{%- if message.tool_calls is defined and message.tool_calls -%} +Tools: {{ message.tool_calls|tojson(ensure_ascii=False) }} +{%- endif -%} +{% endfor -%} +"#; + + let processor = ChatTemplateProcessor::new(template.to_string()); + + let messages = [ChatMessage::User { + content: MessageContent::Text("Test with Unicode: 日本語".to_string()), + name: None, + }]; + + // Convert to JSON values + let json_messages: Vec = messages + .iter() + .map(|msg| serde_json::to_value(msg).unwrap()) + .collect(); + + // This should NOT fail with "unknown keyword argument 'ensure_ascii'" + let result = processor + .apply_chat_template(&json_messages, ChatTemplateParams::default()) + .unwrap(); + + assert!(result.contains("user: Test with Unicode: 日本語")); +} + +/// Test tojson with all HuggingFace kwargs +#[test] +fn test_tojson_with_all_huggingface_kwargs() { + // Template using all the kwargs that HuggingFace's custom tojson accepts + let template = r#" +{%- set data = {"z_key": 1, "a_key": 2, "m_key": 3} -%} +Unsorted: {{ data|tojson }} +Sorted: {{ data|tojson(sort_keys=True) }} +Indented: {{ data|tojson(indent=2) }} +All: {{ data|tojson(ensure_ascii=False, sort_keys=True, indent=2) }} +"#; + + let processor = ChatTemplateProcessor::new(template.to_string()); + let messages: Vec = vec![]; + + // This should NOT fail - all kwargs should be accepted + let result = processor + .apply_chat_template(&messages, ChatTemplateParams::default()) + .unwrap(); + + // Verify sorted output contains keys in alphabetical order + assert!(result.contains("Sorted:")); + // The sorted output should have a_key before m_key before z_key + let sorted_line = result.lines().find(|l| l.starts_with("Sorted:")).unwrap(); + let a_pos = sorted_line.find("a_key").unwrap(); + let m_pos = sorted_line.find("m_key").unwrap(); + let z_pos = sorted_line.find("z_key").unwrap(); + assert!(a_pos < m_pos && m_pos < z_pos, "Keys should be sorted"); + + // Verify indented output is pretty-printed with newlines + assert!( + result.contains("Indented: {\n"), + "Indented JSON should be pretty-printed with newlines" + ); +} + #[test] fn test_content_format_detection() { let string_template = r#"