Render chat templates via hf-chat-template for transformers parity

models/kalosm-llama/Cargo.toml

@@ -23,11 +23,13 @@ kalosm-language-model = { workspace = true, features = ["serde"] }
 kalosm-model-types.workspace = true
 kalosm-common = { workspace = true }
 thiserror.workspace = true
-minijinja = { version = "2.5.0", features = ["loader"] }
-minijinja-contrib = { version = "2.5.0", features = ["pycompat"] }
-chrono = { version = "0.4.41", default-features = false, features = ["now", "std"] }
+# Byte-identical Hugging Face chat-template rendering (vs Python transformers). Renders the
+# template; bos/eos and special tokens are supplied per call. Brings minijinja transitively.
+hf-chat-template = "0.2"
+# Kept only for the `minijinja::Error` type carried by this crate's public error enums.
+minijinja = "2.5.0"
 serde = { version = "1", features = ["derive"] }
-serde_json = { version = "1.0.139", optional = true }
+serde_json = "1.0.139"
 futures-util = { workspace = true }
 futures-channel = { workspace = true }
 web-time.workspace = true
@@ -42,6 +44,7 @@ getrandom = { version = "0.3", features = ["wasm_js"] }
 
 [dev-dependencies]
 ahash = "0.8.12"
+chrono = { version = "0.4.41", default-features = false, features = ["now", "std"] }
 pollster = "0.4"
 tracing-subscriber = "0.3.18"
 pretty_assertions = "1.4.1"
@@ -59,7 +62,7 @@ default = ["gpu"]
 cpu = ["fusor/cpu"]
 gpu = ["fusor/gpu"]
 chat-template-json = ["minijinja/json"]
-hf-config-json = ["dep:serde_json"]
+hf-config-json = []
 vision = ["dep:image", "dep:pollster", "kalosm-language-model/media"]
 structured = ["dep:kalosm-sample", "kalosm-language-model/structured", "dep:rayon"]
 hf-tokenizer-json = ["dep:tokenizers"]

models/kalosm-llama/src/chat_template.rs

@@ -1,46 +1,33 @@
 use std::fmt::Display;
 
-use kalosm_language_model::{ChatMessage, ContentChunk};
-use minijinja::{context, Environment, ErrorKind, Value};
-use minijinja_contrib::pycompat;
+use hf_chat_template::{ChatTemplate, Content, Message, RenderInput};
+use kalosm_language_model::{ChatMessage, ContentChunk, MessageType};
+use minijinja::ErrorKind;
+use serde_json::{json, Value as Json};
 
-#[cfg(test)]
-use kalosm_language_model::MessageType;
 #[cfg(test)]
 use pretty_assertions::assert_eq;
 
+/// Renders Hugging Face chat templates by delegating to `hf-chat-template`, which matches
+/// `transformers.apply_chat_template(..., tokenize=False)` byte-for-byte (Python string methods,
+/// `strftime_now`, `trim_blocks`/`lstrip_blocks`, the transformers-compatible `tojson`, ...).
+///
+/// The error type stays `minijinja::Error` so the surrounding error enums that carry
+/// `#[from] minijinja::Error` are unchanged.
 pub(crate) struct HuggingFaceChatTemplate {
-    environment: Environment<'static>,
+    template: ChatTemplate,
+}
+
+/// Map `hf-chat-template`'s error into the `minijinja::Error` the rest of the crate expects.
+fn to_minijinja_error(err: impl Display) -> minijinja::Error {
+    minijinja::Error::new(ErrorKind::InvalidOperation, err.to_string())
 }
 
 impl HuggingFaceChatTemplate {
     pub(crate) fn create(chat_template: impl Display) -> Result<Self, minijinja::Error> {
-        let chat_template = chat_template.to_string();
-        let mut environment = Environment::new();
-
-        // enable python compatibility methods because most models are tested with python
-        environment.set_unknown_method_callback(pycompat::unknown_method_callback);
-
-        // add the raise_exception function from huggingface templates to the environment
-        let raise_exception = |err_text: String| -> Result<String, minijinja::Error> {
-            Err(minijinja::Error::new(
-                ErrorKind::InvalidOperation,
-                format!("The template raised an exception: {err_text}"),
-            ))
-        };
-        // add the strftime_now function from huggingface templates to the environment
-        let strftime_now = |format: String| -> Result<String, minijinja::Error> {
-            let now = chrono::Utc::now();
-            let formatted_time = now.format(&format).to_string();
-            Ok(formatted_time)
-        };
-        environment.add_function("raise_exception", raise_exception);
-        environment.add_function("strftime_now", strftime_now);
-
-        // compile the template expression in the environment
-        environment.add_template_owned("main", chat_template)?;
-
-        Ok(Self { environment })
+        let template =
+            ChatTemplate::from_str(&chat_template.to_string()).map_err(to_minijinja_error)?;
+        Ok(Self { template })
     }
 
     pub(crate) fn format(
@@ -50,36 +37,51 @@ impl HuggingFaceChatTemplate {
         messages: &[ChatMessage],
         add_generation_prompt: bool,
     ) -> Result<String, minijinja::Error> {
-        let tools: Option<()> = None;
         let messages = messages
             .iter()
             .map(|message| {
-                let role = message.role();
+                let role = match message.role() {
+                    MessageType::SystemPrompt => "system",
+                    MessageType::UserMessage => "user",
+                    MessageType::ModelAnswer => "assistant",
+                };
                 let content = message.content();
-                let content: Value = if let Some(content) = content.as_str() {
-                    content.into()
+                let content = if let Some(text) = content.as_str() {
+                    Content::Text(text.to_string())
                 } else {
-                    let chunks = content
+                    let parts = content
                         .chunks()
                         .iter()
                         .map(|chunk| match chunk {
-                            ContentChunk::Text(text) => {
-                                context! { text }
-                            }
-                            ContentChunk::Media(_) => {
-                                context! { image => "" }
-                            }
+                            ContentChunk::Text(text) => json!({ "text": text }),
+                            ContentChunk::Media(_) => json!({ "image": "" }),
                         })
-                        .collect::<Vec<_>>();
-                    chunks.into()
+                        .collect::<Vec<Json>>();
+                    Content::Parts(parts)
                 };
-                context! { role, content }
+                Message {
+                    role: role.to_string(),
+                    content: Some(content),
+                    ..Default::default()
+                }
             })
             .collect::<Vec<_>>();
-        let ctx = context! { bos_token, eos_token, messages, add_generation_prompt, tools };
-        let template = self.environment.get_template("main")?;
-        let result = template.render(&ctx)?;
-        Ok(result)
+
+        // The template is compiled without special tokens, so pass bos/eos through `extra`,
+        // which the renderer injects as context globals (`{{ bos_token }}`, `{{ eos_token }}`).
+        let mut input = RenderInput {
+            messages,
+            add_generation_prompt,
+            ..Default::default()
+        };
+        input
+            .extra
+            .insert("bos_token".to_string(), Json::String(bos_token.to_string()));
+        input
+            .extra
+            .insert("eos_token".to_string(), Json::String(eos_token.to_string()));
+
+        self.template.render(&input).map_err(to_minijinja_error)
     }
 }