diff --git a/Cargo.lock b/Cargo.lock index 79ac05f6b..89b7e4323 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -217,7 +217,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -228,7 +228,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2032,7 +2032,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.2", ] [[package]] @@ -2053,7 +2053,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -4048,7 +4048,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4404,7 +4404,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5681,6 +5681,18 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" +[[package]] +name = "hf-chat-template" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15e938b1f75afdc8eed9071ff130befb7d014f9b09dbd8d81877ade94ee9f9fe" +dependencies = [ + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", +] + [[package]] name = "hmac" version = "0.12.1" @@ -6203,7 +6215,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi 0.5.2", "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6496,6 +6508,7 @@ dependencies = [ "futures-util", "getrandom 0.3.4", "half", + "hf-chat-template", "image 0.25.10", "kalosm", "kalosm-common", @@ -6505,7 +6518,6 @@ dependencies = [ "kalosm-streams", "kalosm-tokenizer", "minijinja", - "minijinja-contrib", "pollster", "pretty_assertions", "rand 0.9.4", @@ -7239,6 +7251,7 @@ version = "2.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2929e494b2280e1e18959bb2e121da03347ae896896fdfaceaab43c88a02803f" dependencies = [ + "indexmap 2.14.0", "memo-map", "serde", "serde_json", @@ -7611,7 +7624,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9622,7 +9635,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -10302,7 +10315,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -10403,7 +10416,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -10932,7 +10945,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -12495,7 +12508,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/models/kalosm-llama/Cargo.toml b/models/kalosm-llama/Cargo.toml index 592bf3e87..1589f4063 100644 --- a/models/kalosm-llama/Cargo.toml +++ b/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"] diff --git a/models/kalosm-llama/src/chat_template.rs b/models/kalosm-llama/src/chat_template.rs index bfb11b54a..bff4f4e2f 100644 --- a/models/kalosm-llama/src/chat_template.rs +++ b/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 { - 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 { - 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 { - 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 { - 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::>(); - chunks.into() + .collect::>(); + Content::Parts(parts) }; - context! { role, content } + Message { + role: role.to_string(), + content: Some(content), + ..Default::default() + } }) .collect::>(); - 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) } }