From 4c0a3999e408770e0355b842507cb28874d57e2f Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Mon, 20 Apr 2026 20:06:47 -0700 Subject: [PATCH 01/11] Add HuggingFace pretrained vocab provider and refactor bench data iteration - Add `huggingface` feature and `HFVocabProvider` for loading tokenizers from HuggingFace Hub - Refactor bench data methods to use `.values()` instead of `.iter()` for cleaner iteration - Fix clippy warning by allowing explicit counter loop in datagym vocab. This is broken still: ``` crutcher@HeatLamp:~/git/wordchipper$ echo 'hello world' | cargo run -p wordchipper-cli -- cat --model 'hf:Xenova/gpt2' --encode Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.16s Running `target/debug/wordchipper-cli cat --model 'hf:Xenova/gpt2' --encode` Error: External("No pre-tokenizer") ``` --- crates/wordchipper/Cargo.toml | 9 ++ .../src/pretrained/huggingface/hf_factory.rs | 116 ++++++++++++++++++ .../src/pretrained/huggingface/mod.rs | 3 + crates/wordchipper/src/pretrained/mod.rs | 3 + 4 files changed, 131 insertions(+) create mode 100644 crates/wordchipper/src/pretrained/huggingface/hf_factory.rs create mode 100644 crates/wordchipper/src/pretrained/huggingface/mod.rs diff --git a/crates/wordchipper/Cargo.toml b/crates/wordchipper/Cargo.toml index 855bdfc9..fbc33b63 100644 --- a/crates/wordchipper/Cargo.toml +++ b/crates/wordchipper/Cargo.toml @@ -34,6 +34,7 @@ client = [ "download", "datagym", "default-tls", + "huggingface", ] ## The download feature enables downloading vocabularies from the internet. @@ -102,6 +103,12 @@ tracing = [ testing = [] +## Enable loading pretrained huggingface modules. +huggingface = [ + "std", + "dep:tokenizers", +] + [dependencies] # macro packages. @@ -141,6 +148,8 @@ rayon = { workspace = true, optional = true } # "tracing" feature deps: tracing = { workspace = true, optional = true } +tokenizers = { workspace = true, features = ["http"], optional = true } + [dev-dependencies] tempdir = { workspace = true } diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs new file mode 100644 index 00000000..9ffc0a75 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -0,0 +1,116 @@ +use tokenizers::{ + ModelWrapper::BPE, + PreTokenizerWrapper::Split, + pre_tokenizers::split::SplitPattern, + tokenizer::Tokenizer, +}; + +use crate::{ + LabeledVocab, + UnifiedTokenVocab, + VocabDescription, + VocabQuery, + WCError, + WCResult, + alloc::sync::Arc, + prelude::*, + pretrained::factory::{ + VocabProvider, + VocabProviderInventoryHook, + }, + spanners::TextSpanningConfig, + support::{ + regex::RegexPattern, + resources::ResourceLoader, + }, + vocab::{ + ByteMapVocab, + SpanMapVocab, + SpanTokenMap, + }, +}; + +pub struct HFVocabProvider {} + +inventory::submit! { + VocabProviderInventoryHook::new(|| Arc::new(HFVocabProvider{})) +} + +impl VocabProvider for HFVocabProvider { + fn name(&self) -> String { + "hf".to_string() + } + + fn description(&self) -> String { + "HuggingFace vocabularies".to_string() + } + + fn list_vocabs(&self) -> Vec { + vec![] + } + + fn load_vocab( + &self, + query: &VocabQuery, + _loader: &mut dyn ResourceLoader, + ) -> WCResult> { + type T = u32; + + let key = format!("{}/{}", query.path().unwrap(), query.name()); + match Tokenizer::from_pretrained(&key, None) { + Ok(tok) => { + let span_config = if let Some(Split(split)) = tok.get_pre_tokenizer() { + let pattern = match &split.pattern { + SplitPattern::Regex(str) => RegexPattern::Adaptive(str.to_string()), + _ => return Err(WCError::External("No regex pattern".to_string())), + }; + + TextSpanningConfig::from_pattern(pattern) + } else { + return Err(WCError::External("No pre-tokenizer".to_string())); + }; + + if let BPE(bpe) = tok.get_model() { + let vocab = bpe.get_vocab(); + + let span_map: SpanTokenMap = vocab + .iter() + .map(|(s, t)| (s.as_bytes().to_vec(), *t)) + .collect(); + + let byte_map: ByteMapVocab = if let Some(byte_tokens) = (0..256) + .map(|b| { + let k = format!("<{b:#04X}>"); + vocab.get(&k).copied() + }) + .collect::>>() + { + ByteMapVocab::::from_byte_to_token(&byte_tokens) + } else { + return Err(WCError::External( + "Unable to translate: no byte map".to_string(), + )); + }; + + let span_vocab = SpanMapVocab::::new(byte_map, span_map)?; + + let vocab: Arc> = + Arc::new(UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?); + + let id = VocabQuery::new(Some("hf"), query.path(), query.name()); + + let descr: VocabDescription = VocabDescription::new( + id, + &["hf", query.path().unwrap(), query.name()], + "Model loaded from hf", + ); + + Ok(LabeledVocab::new(descr, vocab)) + } else { + Err(WCError::ResourceNotFound(query.to_string())) + } + } + Err(_) => Err(WCError::ResourceNotFound(query.to_string())), + } + } +} diff --git a/crates/wordchipper/src/pretrained/huggingface/mod.rs b/crates/wordchipper/src/pretrained/huggingface/mod.rs new file mode 100644 index 00000000..3f74b586 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/mod.rs @@ -0,0 +1,3 @@ +//! # `HuggingFace` Pretrained Models + +mod hf_factory; diff --git a/crates/wordchipper/src/pretrained/mod.rs b/crates/wordchipper/src/pretrained/mod.rs index acc66fa4..4fc52574 100644 --- a/crates/wordchipper/src/pretrained/mod.rs +++ b/crates/wordchipper/src/pretrained/mod.rs @@ -36,6 +36,9 @@ pub mod factory; pub mod openai; +#[cfg(feature = "huggingface")] +pub mod huggingface; + #[doc(inline)] pub use factory::{ LabeledVocab, From 643dec6ef07ed8889cc95c72ef2db762976ecf3c Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Tue, 21 Apr 2026 17:48:42 -0700 Subject: [PATCH 02/11] Extend HF vocab provider to support byte-level BPE tokenizers - Add support for ByteLevel and Sequence pre-tokenizers in addition to Split - Implement proper byte-to-unicode mapping for byte-level BPE vocabularies - Filter out special tokens when building span maps - Extract and validate split patterns from pre-tokenizer configurations - Add `bytes_char()` function implementing GPT-2 style byte-to-unicode conversion - Use GPT2_PATTERN as default for bare ByteLevel pre-tokenizers - Improve error messages with detailed context about unsupported configurations ```terminaloutput crutcher@HeatLamp:~/git/wordchipper$ echo 'hello world' | cargo run -p wordchipper-cli -- cat --model 'cl100k_base' --encode Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.19s Running `target/debug/wordchipper-cli cat --model cl100k_base --encode` 15339 1917 198 crutcher@HeatLamp:~/git/wordchipper$ echo 'hello world' | cargo run -p wordchipper-cli -- cat --model 'hf:Xenova/text-embedding-ada-002' --encode Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.16s Running `target/debug/wordchipper-cli cat --model 'hf:Xenova/text-embedding-ada-002' --encode` 15339 1917 198 crutcher@HeatLamp:~/git/wordchipper$ echo 'hello world' | cargo run -p wordchipper-cli -- cat --model 'openai:gpt2' --encode Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.16s Running `target/debug/wordchipper-cli cat --model 'openai:gpt2' --encode` 31373 995 198 crutcher@HeatLamp:~/git/wordchipper$ echo 'hello world' | cargo run -p wordchipper-cli -- cat --model 'hf:Xenova/gpt2' --encode Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.16s Running `target/debug/wordchipper-cli cat --model 'hf:Xenova/gpt2' --encode` 31373 995 198 ``` --- .../src/pretrained/huggingface/hf_factory.rs | 194 +++++++++++++----- 1 file changed, 146 insertions(+), 48 deletions(-) diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 9ffc0a75..9a567c66 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -1,6 +1,11 @@ use tokenizers::{ ModelWrapper::BPE, - PreTokenizerWrapper::Split, + PreTokenizerWrapper, + PreTokenizerWrapper::{ + ByteLevel, + Sequence, + Split, + }, pre_tokenizers::split::SplitPattern, tokenizer::Tokenizer, }; @@ -11,6 +16,8 @@ use crate::{ VocabDescription, VocabQuery, WCError, + WCHashMap, + WCHashSet, WCResult, alloc::sync::Arc, prelude::*, @@ -54,63 +61,154 @@ impl VocabProvider for HFVocabProvider { query: &VocabQuery, _loader: &mut dyn ResourceLoader, ) -> WCResult> { + match query.schema() { + None => return Err(WCError::ResourceNotFound(query.to_string())), + Some(schema) => { + if schema != "hf" { + return Err(WCError::ResourceNotFound(query.to_string())); + } + } + } + type T = u32; let key = format!("{}/{}", query.path().unwrap(), query.name()); + let id = format!("{}/{}", query.path().unwrap(), query.name()); + match Tokenizer::from_pretrained(&key, None) { Ok(tok) => { - let span_config = if let Some(Split(split)) = tok.get_pre_tokenizer() { - let pattern = match &split.pattern { - SplitPattern::Regex(str) => RegexPattern::Adaptive(str.to_string()), - _ => return Err(WCError::External("No regex pattern".to_string())), - }; - - TextSpanningConfig::from_pattern(pattern) - } else { - return Err(WCError::External("No pre-tokenizer".to_string())); + let pattern = extract_pattern(tok.get_pre_tokenizer())?; + let span_config = TextSpanningConfig::from_pattern(RegexPattern::Adaptive(pattern)); + + let BPE(bpe) = tok.get_model() else { + return Err(WCError::External( + format!("{} is not BPE compatible", id).to_string(), + )); }; - if let BPE(bpe) = tok.get_model() { - let vocab = bpe.get_vocab(); - - let span_map: SpanTokenMap = vocab - .iter() - .map(|(s, t)| (s.as_bytes().to_vec(), *t)) - .collect(); - - let byte_map: ByteMapVocab = if let Some(byte_tokens) = (0..256) - .map(|b| { - let k = format!("<{b:#04X}>"); - vocab.get(&k).copied() - }) - .collect::>>() - { - ByteMapVocab::::from_byte_to_token(&byte_tokens) - } else { - return Err(WCError::External( - "Unable to translate: no byte map".to_string(), - )); - }; - - let span_vocab = SpanMapVocab::::new(byte_map, span_map)?; - - let vocab: Arc> = - Arc::new(UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?); - - let id = VocabQuery::new(Some("hf"), query.path(), query.name()); - - let descr: VocabDescription = VocabDescription::new( - id, - &["hf", query.path().unwrap(), query.name()], - "Model loaded from hf", - ); - - Ok(LabeledVocab::new(descr, vocab)) - } else { - Err(WCError::ResourceNotFound(query.to_string())) + // TODO: Add support for unknown token. + if let Some(unk) = bpe.get_unk_token() { + return Err(WCError::External(format!("BPE has unk_token {unk:?}"))); } + + let vocab = bpe.get_vocab(); + + let specials: WCHashSet = + tok.get_added_tokens_decoder().keys().copied().collect(); + + // Forward and inverse bytes_to_unicode maps. + let b2c = bytes_char(); + let c2b: WCHashMap = b2c.iter().map(|(&b, &c)| (c, b)).collect(); + + // Span map: decode every non-special vocab string back to bytes. + let mut span_map: SpanTokenMap = SpanTokenMap::default(); + for (s, id) in &vocab { + if specials.contains(id) { + continue; + } + let mut bytes = Vec::with_capacity(s.len()); + for ch in s.chars() { + match c2b.get(&ch) { + Some(&b) => bytes.push(b), + None => { + return Err(WCError::External(format!( + "token {s:?} (id {id}) has non-byte-level codepoint {ch:?}" + ))); + } + } + } + span_map.insert(bytes, (*id)); + } + + // Byte map: the single-char string for each byte must resolve in the vocab. + let byte_tokens: Vec = (0u8..=255) + .map(|b| { + let key: String = std::iter::once(b2c[&b]).collect(); + vocab.get(&key).copied().ok_or(b) + }) + .collect::, _>>() + .map_err(|b| WCError::External(format!("missing byte token for 0x{b:02x}")))?; + + let byte_map = ByteMapVocab::::from_byte_to_token(&byte_tokens); + let span_vocab = SpanMapVocab::::new(byte_map, span_map)?; + + let vocab: Arc> = + Arc::new(UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?); + + let id = VocabQuery::new(Some("hf"), query.path(), query.name()); + + let descr: VocabDescription = VocabDescription::new( + id, + &["hf", query.path().unwrap(), query.name()], + "Model loaded from hf", + ); + + Ok(LabeledVocab::new(descr, vocab)) } Err(_) => Err(WCError::ResourceNotFound(query.to_string())), } } } +// GPT-2 / r50k-style default, used when pretokenizer is a bare ByteLevel. +const GPT2_PATTERN: &str = + r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"; + +fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { + fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { + match &s.pattern { + SplitPattern::Regex(r) => Ok(r.clone()), + _ => Err(WCError::External("Split without Regex pattern".into())), + } + } + match pt { + Some(Split(s)) => split_regex(s), + Some(ByteLevel(_)) => Ok(GPT2_PATTERN.to_string()), + Some(Sequence(seq)) => { + let mut found = None; + for sub in seq.as_ref() { + match &sub { + Split(s) => { + if found.is_some() { + return Err(WCError::External("Sequence has multiple Splits".into())); + } + found = Some(split_regex(s)?); + } + ByteLevel(_) => {} // sibling byte-encoder, fine + _ => return Err(WCError::External("unsupported member in Sequence".into())), + } + } + found.ok_or_else(|| WCError::External("Sequence has no Split regex".into())) + } + Some(_) => Err(WCError::External("unsupported pre-tokenizer".into())), + None => Err(WCError::External("no pre-tokenizer".into())), + } +} + +/// Converts bytes to unicode characters. +/// See +fn bytes_char() -> WCHashMap { + let mut bs: Vec = vec![]; + bs.extend(b'!'..=b'~'); + bs.extend(b'\xA1'..=b'\xAC'); + bs.extend(b'\xAE'..=b'\xFF'); + + let mut cs: Vec = bs.iter().map(|i| *i as u32).collect(); + let mut n = 0; + + for b in 0..=255u8 { + if !bs.contains(&b) { + bs.push(b); + cs.push(u32::pow(2, 8) + n); + n += 1; + } + } + + // Safety: cs contains all values from bs (between 0 and 255), + // and some values of value 2⁸ + n, where n is between 0 and 255. This is + // between 255 and 512. Both ranges are valid UTF-32 values (which is fully + // saturated until 0xD000) + bs.into_iter() + .zip(cs) + .map(|(f, t)| (f, unsafe { std::char::from_u32_unchecked(t) })) + .collect() +} From 366dd4e5d6a4f959c74476d69fc49719bbbc51d8 Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Tue, 21 Apr 2026 18:02:27 -0700 Subject: [PATCH 03/11] Make hf factory a better pass-through factory. Will now search hf for pretrained models if there is no schema. --- .../src/pretrained/huggingface/hf_factory.rs | 168 +++++++++--------- 1 file changed, 86 insertions(+), 82 deletions(-) diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 9a567c66..03a8dd69 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -37,6 +37,76 @@ use crate::{ }, }; +// GPT-2 / r50k-style default, used when pretokenizer is a bare ByteLevel. +const GPT2_PATTERN: &str = + r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"; + +fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { + fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { + match &s.pattern { + SplitPattern::Regex(r) => Ok(r.clone()), + _ => Err(WCError::External("Split without Regex pattern".into())), + } + } + match pt { + Some(Split(s)) => split_regex(s), + Some(ByteLevel(_)) => Ok(GPT2_PATTERN.to_string()), + Some(Sequence(seq)) => { + let mut found = None; + for sub in seq.as_ref() { + match &sub { + Split(s) => { + if found.is_some() { + return Err(WCError::External("Sequence has multiple Splits".into())); + } + found = Some(split_regex(s)?); + } + ByteLevel(_) => {} // sibling byte-encoder, fine + _ => return Err(WCError::External("unsupported member in Sequence".into())), + } + } + found.ok_or_else(|| WCError::External("Sequence has no Split regex".into())) + } + Some(_) => Err(WCError::External("unsupported pre-tokenizer".into())), + None => Err(WCError::External("no pre-tokenizer".into())), + } +} + +/// Converts bytes to Unicode characters. +/// See +/// +/// This is from tokenizers; but is private in that crate. +/// +/// TODO: Workout what this is doing, relative to the bytemap. +/// This seems to be some default map for gpt2; and might be shared +/// with the `BytMap` code for loading datagym. +fn bytes_char() -> WCHashMap { + let mut bs: Vec = vec![]; + bs.extend(b'!'..=b'~'); + bs.extend(b'\xA1'..=b'\xAC'); + bs.extend(b'\xAE'..=b'\xFF'); + + let mut cs: Vec = bs.iter().map(|i| *i as u32).collect(); + let mut n = 0; + + for b in 0..=255u8 { + if !bs.contains(&b) { + bs.push(b); + cs.push(u32::pow(2, 8) + n); + n += 1; + } + } + + // Safety: cs contains all values from bs (between 0 and 255), + // and some values of value 2⁸ + n, where n is between 0 and 255. This is + // between 255 and 512. Both ranges are valid UTF-32 values (which is fully + // saturated until 0xD000) + bs.into_iter() + .zip(cs) + .map(|(f, t)| (f, unsafe { std::char::from_u32_unchecked(t) })) + .collect() +} + pub struct HFVocabProvider {} inventory::submit! { @@ -61,28 +131,22 @@ impl VocabProvider for HFVocabProvider { query: &VocabQuery, _loader: &mut dyn ResourceLoader, ) -> WCResult> { - match query.schema() { - None => return Err(WCError::ResourceNotFound(query.to_string())), - Some(schema) => { - if schema != "hf" { - return Err(WCError::ResourceNotFound(query.to_string())); - } - } + if let Some(schema) = query.schema() + && schema != "hf" + { + return Err(WCError::ResourceNotFound(query.to_string())); } type T = u32; - let key = format!("{}/{}", query.path().unwrap(), query.name()); - let id = format!("{}/{}", query.path().unwrap(), query.name()); - - match Tokenizer::from_pretrained(&key, None) { + match Tokenizer::from_pretrained(&query.clone().with_schema(None).to_string(), None) { Ok(tok) => { let pattern = extract_pattern(tok.get_pre_tokenizer())?; let span_config = TextSpanningConfig::from_pattern(RegexPattern::Adaptive(pattern)); let BPE(bpe) = tok.get_model() else { return Err(WCError::External( - format!("{} is not BPE compatible", id).to_string(), + format!("{} is not BPE compatible", query).to_string(), )); }; @@ -117,7 +181,7 @@ impl VocabProvider for HFVocabProvider { } } } - span_map.insert(bytes, (*id)); + span_map.insert(bytes, *id); } // Byte map: the single-char string for each byte must resolve in the vocab. @@ -135,13 +199,16 @@ impl VocabProvider for HFVocabProvider { let vocab: Arc> = Arc::new(UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?); - let id = VocabQuery::new(Some("hf"), query.path(), query.name()); + let id = query.clone().with_schema(Some("hf")); + + let mut context = vec!["hf"]; + if query.path().is_some() { + context.push(query.path().unwrap()); + } + context.push(query.name()); - let descr: VocabDescription = VocabDescription::new( - id, - &["hf", query.path().unwrap(), query.name()], - "Model loaded from hf", - ); + let descr: VocabDescription = + VocabDescription::new(id, &context, "Model loaded from hf"); Ok(LabeledVocab::new(descr, vocab)) } @@ -149,66 +216,3 @@ impl VocabProvider for HFVocabProvider { } } } -// GPT-2 / r50k-style default, used when pretokenizer is a bare ByteLevel. -const GPT2_PATTERN: &str = - r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"; - -fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { - fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { - match &s.pattern { - SplitPattern::Regex(r) => Ok(r.clone()), - _ => Err(WCError::External("Split without Regex pattern".into())), - } - } - match pt { - Some(Split(s)) => split_regex(s), - Some(ByteLevel(_)) => Ok(GPT2_PATTERN.to_string()), - Some(Sequence(seq)) => { - let mut found = None; - for sub in seq.as_ref() { - match &sub { - Split(s) => { - if found.is_some() { - return Err(WCError::External("Sequence has multiple Splits".into())); - } - found = Some(split_regex(s)?); - } - ByteLevel(_) => {} // sibling byte-encoder, fine - _ => return Err(WCError::External("unsupported member in Sequence".into())), - } - } - found.ok_or_else(|| WCError::External("Sequence has no Split regex".into())) - } - Some(_) => Err(WCError::External("unsupported pre-tokenizer".into())), - None => Err(WCError::External("no pre-tokenizer".into())), - } -} - -/// Converts bytes to unicode characters. -/// See -fn bytes_char() -> WCHashMap { - let mut bs: Vec = vec![]; - bs.extend(b'!'..=b'~'); - bs.extend(b'\xA1'..=b'\xAC'); - bs.extend(b'\xAE'..=b'\xFF'); - - let mut cs: Vec = bs.iter().map(|i| *i as u32).collect(); - let mut n = 0; - - for b in 0..=255u8 { - if !bs.contains(&b) { - bs.push(b); - cs.push(u32::pow(2, 8) + n); - n += 1; - } - } - - // Safety: cs contains all values from bs (between 0 and 255), - // and some values of value 2⁸ + n, where n is between 0 and 255. This is - // between 255 and 512. Both ranges are valid UTF-32 values (which is fully - // saturated until 0xD000) - bs.into_iter() - .zip(cs) - .map(|(f, t)| (f, unsafe { std::char::from_u32_unchecked(t) })) - .collect() -} From 7ed89a87ae590071edd76274d91047549e7b349a Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Tue, 21 Apr 2026 18:55:50 -0700 Subject: [PATCH 04/11] Add `VocabQuery::to_context()` method and refactor HF vocab provider - Add `VocabQuery::to_context()` to build cache context from schema, path, and name - Make `VocabDescription::new()` generic over context and description string types - Extract `vocab_from_hf_tokenizer()` as reusable conversion function - Add special token handling to span_config instead of filtering them out - Fix ByteLevel pre-tokenizer to require `use_regex=true` and use OA_GPT2_PATTERN - Update HF vocab description context to use `to_context()` helper - Add tests for `to_context()` with various query formats --- .../pretrained/factory/vocab_description.rs | 14 +- .../src/pretrained/factory/vocab_query.rs | 28 ++- .../src/pretrained/huggingface/hf_factory.rs | 170 ++++++++++-------- 3 files changed, 133 insertions(+), 79 deletions(-) diff --git a/crates/wordchipper/src/pretrained/factory/vocab_description.rs b/crates/wordchipper/src/pretrained/factory/vocab_description.rs index 9e7fedbc..a8cf2782 100644 --- a/crates/wordchipper/src/pretrained/factory/vocab_description.rs +++ b/crates/wordchipper/src/pretrained/factory/vocab_description.rs @@ -21,20 +21,24 @@ pub struct VocabDescription { impl VocabDescription { /// Build a new vocabulary description. - pub fn new( + pub fn new( id: Q, - context: &[&str], - description: &str, + context: &[C], + description: D, ) -> Self where Q: Into, + C: AsRef, + D: AsRef, { let id = id.into(); + let context = context.iter().map(|c| c.as_ref().to_string()).collect(); + let description = description.as_ref().to_string(); Self { id, - context: context.iter().map(|&s| s.to_string()).collect(), - description: description.to_string(), + context, + description, } } diff --git a/crates/wordchipper/src/pretrained/factory/vocab_query.rs b/crates/wordchipper/src/pretrained/factory/vocab_query.rs index c33af521..1efbdb8b 100644 --- a/crates/wordchipper/src/pretrained/factory/vocab_query.rs +++ b/crates/wordchipper/src/pretrained/factory/vocab_query.rs @@ -66,7 +66,7 @@ impl Display for VocabQuery { } impl VocabQuery { - /// Build a new query from structure. + /// Build a new query. pub fn new( schema: Option<&str>, path: Option<&str>, @@ -174,6 +174,19 @@ impl VocabQuery { } query.name() == self.name() } + + /// Build a cache context for this query. + pub fn to_context(&self) -> Vec { + let mut context = Vec::new(); + if let Some(schema) = self.schema() { + context.push(schema.to_string()); + } + if let Some(path) = self.path() { + context.extend(path.split('/').map(|p| p.to_string())); + } + context.push(self.name().to_string()); + context + } } #[cfg(test)] @@ -202,6 +215,19 @@ mod tests { VocabQuery::new(Some("xyz"), Some("foo/bar"), "vocab_name") ); } + + #[test] + fn test_to_context() { + let q = VocabQuery::from_str("vocab_name").unwrap(); + assert_eq!(q.to_context(), vec!["vocab_name"]); + + let q = VocabQuery::from_str("foo/bar/vocab_name").unwrap(); + assert_eq!(q.to_context(), vec!["foo", "bar", "vocab_name"]); + + let q = VocabQuery::from_str("xyz:foo/bar/vocab_name").unwrap(); + assert_eq!(q.to_context(), vec!["xyz", "foo", "bar", "vocab_name"]); + } + #[test] fn test_vocab_query_with_schema() { let query = VocabQuery::new(None, None, "vocab_name").with_schema(Some("provider")); diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 03a8dd69..0b851ed1 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -14,6 +14,7 @@ use crate::{ LabeledVocab, UnifiedTokenVocab, VocabDescription, + VocabIndex, VocabQuery, WCError, WCHashMap, @@ -21,9 +22,12 @@ use crate::{ WCResult, alloc::sync::Arc, prelude::*, - pretrained::factory::{ - VocabProvider, - VocabProviderInventoryHook, + pretrained::{ + factory::{ + VocabProvider, + VocabProviderInventoryHook, + }, + openai::OA_GPT2_PATTERN, }, spanners::TextSpanningConfig, support::{ @@ -37,20 +41,19 @@ use crate::{ }, }; -// GPT-2 / r50k-style default, used when pretokenizer is a bare ByteLevel. -const GPT2_PATTERN: &str = - r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"; - -fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { - fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { +fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { + fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { match &s.pattern { - SplitPattern::Regex(r) => Ok(r.clone()), + SplitPattern::Regex(r) => Ok(r.clone().into()), _ => Err(WCError::External("Split without Regex pattern".into())), } } match pt { Some(Split(s)) => split_regex(s), - Some(ByteLevel(_)) => Ok(GPT2_PATTERN.to_string()), + Some(ByteLevel(bl)) if bl.use_regex => Ok(OA_GPT2_PATTERN.into()), + Some(ByteLevel(_)) => Err(WCError::External( + "ByteLevel with use_regex=false has no splitting regex".into(), + )), Some(Sequence(seq)) => { let mut found = None; for sub in seq.as_ref() { @@ -107,6 +110,84 @@ fn bytes_char() -> WCHashMap { .collect() } +/// Attempt to convert a `HuggingFace` tokenizer to a `WordChipper` vocabulary. +pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult>> { + type T = u32; + + let pattern = extract_pattern(tok.get_pre_tokenizer())?; + let mut span_config: TextSpanningConfig = TextSpanningConfig::from_pattern(pattern); + + let BPE(bpe) = tok.get_model() else { + return Err(WCError::External( + "Tokenizer is not BPE compatible".to_string(), + )); + }; + + // TODO: Add support for unknown token. + if let Some(unk) = bpe.get_unk_token() { + return Err(WCError::External(format!("BPE has unk_token {unk:?}"))); + } + + let hf_vocab = bpe.get_vocab(); + + let special_tokens: WCHashSet = tok.get_added_tokens_decoder().keys().copied().collect(); + + // Forward and inverse bytes_to_unicode maps. + let b2c = bytes_char(); + let c2b: WCHashMap = b2c.iter().map(|(&b, &c)| (c, b)).collect(); + + // Span map: decode every non-special vocab string back to bytes. + let mut span_map: SpanTokenMap = SpanTokenMap::default(); + for (s, id) in &hf_vocab { + if special_tokens.contains(id) { + span_config.specials_mut().add_str_word(s, *id); + } else { + let mut bytes = Vec::with_capacity(s.len()); + for ch in s.chars() { + match c2b.get(&ch) { + Some(&b) => bytes.push(b), + None => { + return Err(WCError::External(format!( + "token {s:?} (id {id}) has non-byte-level codepoint {ch:?}" + ))); + } + } + } + span_map.insert(bytes, *id); + } + } + + assert_eq!(span_config.specials().len(), special_tokens.len()); + + // Byte map: the single-char string for each byte must resolve in the vocab. + let byte_tokens: Vec = (0u8..=255) + .map(|b| { + let key: String = std::iter::once(b2c[&b]).collect(); + hf_vocab.get(&key).copied().ok_or(b) + }) + .collect::, _>>() + .map_err(|b| WCError::External(format!("missing byte token for 0x{b:02x}")))?; + + let byte_map = ByteMapVocab::::from_byte_to_token(&byte_tokens); + let span_vocab = SpanMapVocab::::new(byte_map, span_map)?; + + let expected_len = span_vocab.len() + span_config.specials().len(); + + let vocab: Arc> = + Arc::new(UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?); + + // TODO: should `vocab.len()` include the special len()? + if vocab.len() + vocab.special_vocab().len() != expected_len { + return Err(WCError::External(format!( + "Expected {} tokens, got {}", + expected_len, + vocab.len() + ))); + } + + Ok(vocab) +} + pub struct HFVocabProvider {} inventory::submit! { @@ -137,69 +218,9 @@ impl VocabProvider for HFVocabProvider { return Err(WCError::ResourceNotFound(query.to_string())); } - type T = u32; - - match Tokenizer::from_pretrained(&query.clone().with_schema(None).to_string(), None) { + match Tokenizer::from_pretrained(query.clone().with_schema(None).to_string(), None) { Ok(tok) => { - let pattern = extract_pattern(tok.get_pre_tokenizer())?; - let span_config = TextSpanningConfig::from_pattern(RegexPattern::Adaptive(pattern)); - - let BPE(bpe) = tok.get_model() else { - return Err(WCError::External( - format!("{} is not BPE compatible", query).to_string(), - )); - }; - - // TODO: Add support for unknown token. - if let Some(unk) = bpe.get_unk_token() { - return Err(WCError::External(format!("BPE has unk_token {unk:?}"))); - } - - let vocab = bpe.get_vocab(); - - let specials: WCHashSet = - tok.get_added_tokens_decoder().keys().copied().collect(); - - // Forward and inverse bytes_to_unicode maps. - let b2c = bytes_char(); - let c2b: WCHashMap = b2c.iter().map(|(&b, &c)| (c, b)).collect(); - - // Span map: decode every non-special vocab string back to bytes. - let mut span_map: SpanTokenMap = SpanTokenMap::default(); - for (s, id) in &vocab { - if specials.contains(id) { - continue; - } - let mut bytes = Vec::with_capacity(s.len()); - for ch in s.chars() { - match c2b.get(&ch) { - Some(&b) => bytes.push(b), - None => { - return Err(WCError::External(format!( - "token {s:?} (id {id}) has non-byte-level codepoint {ch:?}" - ))); - } - } - } - span_map.insert(bytes, *id); - } - - // Byte map: the single-char string for each byte must resolve in the vocab. - let byte_tokens: Vec = (0u8..=255) - .map(|b| { - let key: String = std::iter::once(b2c[&b]).collect(); - vocab.get(&key).copied().ok_or(b) - }) - .collect::, _>>() - .map_err(|b| WCError::External(format!("missing byte token for 0x{b:02x}")))?; - - let byte_map = ByteMapVocab::::from_byte_to_token(&byte_tokens); - let span_vocab = SpanMapVocab::::new(byte_map, span_map)?; - - let vocab: Arc> = - Arc::new(UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?); - - let id = query.clone().with_schema(Some("hf")); + let vocab = vocab_from_hf_tokenizer(&tok)?; let mut context = vec!["hf"]; if query.path().is_some() { @@ -207,6 +228,9 @@ impl VocabProvider for HFVocabProvider { } context.push(query.name()); + let id = query.clone().with_schema(Some("hf")); + let context = id.to_context(); + let descr: VocabDescription = VocabDescription::new(id, &context, "Model loaded from hf"); From b0e9475f3976d9a160cb989e89eae75bd8f5a1bc Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Tue, 21 Apr 2026 20:51:12 -0700 Subject: [PATCH 05/11] Partial work on Qwen3.5. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This is broken; the vocab seems to have aggregate tokens with no prior piecewise merges that are not special tokens. Which wordchipper enforces is never the case. I need to investigate if this is true for this dataset; and explore what needs to change about wordchipper to handle this. It may be that we need to support a situation where there multi-byte strings in the span map that are not in the byte map, but also *products* in the merge map (so they can only be produced by exact-match lookup, never by merges); but can still be members of merge pairs. ```terminaloutput echo 'hello world' | RUST_BACKTRACE=1 cargo run -p wordchipper-cli -- cat --model Qwen/Qwen3.5-9B --encode Blocking waiting for file lock on artifact directory Compiling wordchipper v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper) Compiling wordchipper-training v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper-training) Compiling wordchipper-cli-util v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper-cli-util) Compiling wordchipper-cli v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper-cli) Finished `dev` profile [unoptimized + debuginfo] target(s) in 15.19s Running `target/debug/wordchipper-cli cat --model Qwen/Qwen3.5-9B --encode` Debug: Some(("ĠаÑįÑĢо", 157513)) Debug: { 248067: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248064: AddedToken { content: "<|repo_name|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248046: AddedToken { content: "<|im_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248059: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248054: AddedToken { content: "<|vision_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248044: AddedToken { content: "<|endoftext|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248058: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248049: AddedToken { content: "<|box_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248057: AddedToken { content: "<|video_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248056: AddedToken { content: "<|image_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248062: AddedToken { content: "<|fim_suffix|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248050: AddedToken { content: "<|box_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248055: AddedToken { content: "<|vision_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248048: AddedToken { content: "<|object_ref_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248045: AddedToken { content: "<|im_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248066: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248068: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248069: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248052: AddedToken { content: "<|quad_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248053: AddedToken { content: "<|vision_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248060: AddedToken { content: "<|fim_prefix|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248047: AddedToken { content: "<|object_ref_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248051: AddedToken { content: "<|quad_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248061: AddedToken { content: "<|fim_middle|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248063: AddedToken { content: "<|fim_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248065: AddedToken { content: "<|file_sep|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, } thread 'main' (852780) panicked at crates/wordchipper/src/vocab/span_vocab.rs:247:51: called `Result::unwrap()` on an `Err` value: VocabConflict("Pair (157513, 157853) -> 229601 parent 157513 is not defined") stack backtrace: 0: __rustc::rust_begin_unwind at /rustc/59807616e1fa2540724bfbac14d7976d7e4a3860/library/std/src/panicking.rs:689:5 1: core::panicking::panic_fmt at /rustc/59807616e1fa2540724bfbac14d7976d7e4a3860/library/core/src/panicking.rs:80:14 2: core::result::unwrap_failed at /rustc/59807616e1fa2540724bfbac14d7976d7e4a3860/library/core/src/result.rs:1867:5 3: core::result::Result::unwrap at /home/crutcher/.rustup/toolchains/stable-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/result.rs:1233:23 4: wordchipper::vocab::span_vocab::SpanMapVocab::to_pair_vocab at ./crates/wordchipper/src/vocab/span_vocab.rs:247:51 5: wordchipper::vocab::unified_vocab::UnifiedTokenVocab::from_span_vocab at ./crates/wordchipper/src/vocab/unified_vocab.rs:85:37 6: wordchipper::pretrained::huggingface::hf_factory::vocab_from_hf_tokenizer at ./crates/wordchipper/src/pretrained/huggingface/hf_factory.rs:199:18 7: ::load_vocab at ./crates/wordchipper/src/pretrained/huggingface/hf_factory.rs:245:29 8: wordchipper::pretrained::factory::vocab_factory::VocabFactory::load_vocab at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:229:28 9: wordchipper::pretrained::factory::vocab_factory::load_vocab::{{closure}} at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:99:55 10: wordchipper::pretrained::factory::vocab_factory::with_vocab_factory at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:76:5 11: wordchipper::pretrained::factory::vocab_factory::load_vocab at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:99:5 12: wordchipper_cli_util::model_selector::ModelSelectorArgs::load_vocab at ./crates/wordchipper-cli-util/src/model_selector.rs:29:21 13: wordchipper_cli_util::model_selector::ModelSelectorArgs::load_tokenizer at ./crates/wordchipper-cli-util/src/model_selector.rs:41:26 14: wordchipper_cli::commands::cat_cmd::CatArgs::run at ./crates/wordchipper-cli/src/commands/cat_cmd.rs:51:45 15: wordchipper_cli::Commands::run at ./crates/wordchipper-cli/src/main.rs:44:39 16: wordchipper_cli::main at ./crates/wordchipper-cli/src/main.rs:18:18 17: core::ops::function::FnOnce::call_once at /home/crutcher/.rustup/toolchains/stable-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/ops/function.rs:250:5 note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace. ```terminaloutput $ echo 'hello world' | RUST_BACKTRACE=1 cargo run -p wordchipper-cli -- cat --model Qwen/Qwen3.5-9B --encode Blocking waiting for file lock on artifact directory Compiling wordchipper v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper) Compiling wordchipper-training v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper-training) Compiling wordchipper-cli-util v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper-cli-util) Compiling wordchipper-cli v0.9.1 (/home/crutcher/git/wordchipper/crates/wordchipper-cli) Finished `dev` profile [unoptimized + debuginfo] target(s) in 15.19s Running `target/debug/wordchipper-cli cat --model Qwen/Qwen3.5-9B --encode` Debug: Some(("ĠаÑįÑĢо", 157513)) Debug: { 248067: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248064: AddedToken { content: "<|repo_name|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248046: AddedToken { content: "<|im_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248059: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248054: AddedToken { content: "<|vision_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248044: AddedToken { content: "<|endoftext|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248058: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248049: AddedToken { content: "<|box_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248057: AddedToken { content: "<|video_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248056: AddedToken { content: "<|image_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248062: AddedToken { content: "<|fim_suffix|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248050: AddedToken { content: "<|box_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248055: AddedToken { content: "<|vision_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248048: AddedToken { content: "<|object_ref_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248045: AddedToken { content: "<|im_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248066: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248068: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248069: AddedToken { content: "", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248052: AddedToken { content: "<|quad_end|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248053: AddedToken { content: "<|vision_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248060: AddedToken { content: "<|fim_prefix|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248047: AddedToken { content: "<|object_ref_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248051: AddedToken { content: "<|quad_start|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: true, }, 248061: AddedToken { content: "<|fim_middle|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248063: AddedToken { content: "<|fim_pad|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, 248065: AddedToken { content: "<|file_sep|>", single_word: false, lstrip: false, rstrip: false, normalized: false, special: false, }, } thread 'main' (852780) panicked at crates/wordchipper/src/vocab/span_vocab.rs:247:51: called `Result::unwrap()` on an `Err` value: VocabConflict("Pair (157513, 157853) -> 229601 parent 157513 is not defined") stack backtrace: 0: __rustc::rust_begin_unwind at /rustc/59807616e1fa2540724bfbac14d7976d7e4a3860/library/std/src/panicking.rs:689:5 1: core::panicking::panic_fmt at /rustc/59807616e1fa2540724bfbac14d7976d7e4a3860/library/core/src/panicking.rs:80:14 2: core::result::unwrap_failed at /rustc/59807616e1fa2540724bfbac14d7976d7e4a3860/library/core/src/result.rs:1867:5 3: core::result::Result::unwrap at /home/crutcher/.rustup/toolchains/stable-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/result.rs:1233:23 4: wordchipper::vocab::span_vocab::SpanMapVocab::to_pair_vocab at ./crates/wordchipper/src/vocab/span_vocab.rs:247:51 5: wordchipper::vocab::unified_vocab::UnifiedTokenVocab::from_span_vocab at ./crates/wordchipper/src/vocab/unified_vocab.rs:85:37 6: wordchipper::pretrained::huggingface::hf_factory::vocab_from_hf_tokenizer at ./crates/wordchipper/src/pretrained/huggingface/hf_factory.rs:199:18 7: ::load_vocab at ./crates/wordchipper/src/pretrained/huggingface/hf_factory.rs:245:29 8: wordchipper::pretrained::factory::vocab_factory::VocabFactory::load_vocab at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:229:28 9: wordchipper::pretrained::factory::vocab_factory::load_vocab::{{closure}} at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:99:55 10: wordchipper::pretrained::factory::vocab_factory::with_vocab_factory at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:76:5 11: wordchipper::pretrained::factory::vocab_factory::load_vocab at ./crates/wordchipper/src/pretrained/factory/vocab_factory.rs:99:5 12: wordchipper_cli_util::model_selector::ModelSelectorArgs::load_vocab at ./crates/wordchipper-cli-util/src/model_selector.rs:29:21 13: wordchipper_cli_util::model_selector::ModelSelectorArgs::load_tokenizer at ./crates/wordchipper-cli-util/src/model_selector.rs:41:26 14: wordchipper_cli::commands::cat_cmd::CatArgs::run at ./crates/wordchipper-cli/src/commands/cat_cmd.rs:51:45 15: wordchipper_cli::Commands::run at ./crates/wordchipper-cli/src/main.rs:44:39 16: wordchipper_cli::main at ./crates/wordchipper-cli/src/main.rs:18:18 17: core::ops::function::FnOnce::call_once at /home/crutcher/.rustup/toolchains/stable-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/ops/function.rs:250:5 note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace. ``` --- .../src/pretrained/huggingface/hf_factory.rs | 28 +++++++++++++++++-- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 0b851ed1..5c6406f4 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -1,3 +1,5 @@ +use std::println; + use tokenizers::{ ModelWrapper::BPE, PreTokenizerWrapper, @@ -130,7 +132,21 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult = tok.get_added_tokens_decoder().keys().copied().collect(); + println!( + "Debug: {:?}", + hf_vocab.iter().find(|(_, id)| **id == 157513) + ); + + // TODO: This is broken for Qwen/Qwen3.5-9B for some reason. + let mut special_tokens: WCHashSet = Default::default(); + + let decoder = tok.get_added_tokens_decoder(); + println!("Debug: {:#?}", decoder); + + for (t, at) in decoder.iter() { + span_config.specials_mut().add_str_word(&at.content, *t); + special_tokens.insert(*t); + } // Forward and inverse bytes_to_unicode maps. let b2c = bytes_char(); @@ -140,7 +156,7 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult = SpanTokenMap::default(); for (s, id) in &hf_vocab { if special_tokens.contains(id) { - span_config.specials_mut().add_str_word(s, *id); + continue; } else { let mut bytes = Vec::with_capacity(s.len()); for ch in s.chars() { @@ -157,7 +173,13 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult = (0u8..=255) From 460dacfa427dfe44410c940000aaa66d33a04551 Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Sun, 10 May 2026 23:10:21 -0700 Subject: [PATCH 06/11] Update partial support for HF with basic errors. --- Cargo.lock | 12 +++- Cargo.toml | 1 + crates/wordchipper/Cargo.toml | 1 + crates/wordchipper/README.md | 5 ++ crates/wordchipper/src/errors.rs | 10 ++- .../src/pretrained/huggingface/hf_factory.rs | 6 +- crates/wordchipper/src/support/mod.rs | 2 + .../src/support/with_ok_or_panic.rs | 61 +++++++++++++++++++ crates/wordchipper/src/vocab/pair_vocab.rs | 15 +++-- crates/wordchipper/src/vocab/span_vocab.rs | 7 ++- 10 files changed, 107 insertions(+), 13 deletions(-) create mode 100644 crates/wordchipper/src/support/with_ok_or_panic.rs diff --git a/Cargo.lock b/Cargo.lock index 7a8b0c04..74ed0b65 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1631,6 +1631,15 @@ dependencies = [ "web-time", ] +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + [[package]] name = "integer-encoding" version = "3.0.4" @@ -2471,7 +2480,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -4317,6 +4326,7 @@ dependencies = [ "fancy-regex", "foldhash 0.2.0", "hashbrown 0.16.1", + "indoc", "inventory", "log", "logos", diff --git a/Cargo.toml b/Cargo.toml index 71bb8825..832fa430 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -82,6 +82,7 @@ divan = { package = "codspeed-divan-compat", version = "4.3.0" } document-features = "0.2.12" humansize = "2.1.3" indicatif = "0.18.4" +indoc = "2.0.7" js-sys = "0.3" parquet = "58.0.0" proptest = "1.10.0" diff --git a/crates/wordchipper/Cargo.toml b/crates/wordchipper/Cargo.toml index fbc33b63..5e9c8ba1 100644 --- a/crates/wordchipper/Cargo.toml +++ b/crates/wordchipper/Cargo.toml @@ -125,6 +125,7 @@ ringbuffer = { workspace = true } regex = { workspace = true, features = ["unicode"] } regex-automata = { workspace = true, features = ["alloc", "meta", "nfa-thompson", "hybrid", "unicode"] } strum = { workspace = true } +indoc = { workspace = true } # Provides HashMap/HashSet in no_std mode (non-optional so `default-features = false` just works). hashbrown = { workspace = true, features = ["alloc"] } diff --git a/crates/wordchipper/README.md b/crates/wordchipper/README.md index 32d33601..6f9c07a6 100644 --- a/crates/wordchipper/README.md +++ b/crates/wordchipper/README.md @@ -76,6 +76,11 @@ configuration. For a number of pretrained models, simplified constructors are available to download, cache, and load the vocabulary. +At this time, we have support for the following loaders: + +* `openai:[{PATH}/]{NAME}` - Lod pre-trained OpenAI models. +* `hf:[{PATH}/]{NAME}` - Load pre-trained HuggingFace models. + See: [wordchipper::get_model]( https://docs.rs/wordchipper/latest/wordchipper/fn.get_model.html) diff --git a/crates/wordchipper/src/errors.rs b/crates/wordchipper/src/errors.rs index 7396ea4c..c7ae4141 100644 --- a/crates/wordchipper/src/errors.rs +++ b/crates/wordchipper/src/errors.rs @@ -5,12 +5,16 @@ use crate::alloc::string::String; /// Errors from wordchipper operations. #[derive(Debug, thiserror::Error)] pub enum WCError { + /// Not Implemented Error. + #[error("Not Implemented: {0}")] + NotImplemented(String), + /// Resource not found. - #[error("{0}")] + #[error("Resource Not Found: {0}")] ResourceNotFound(String), /// The resource is a duplicate. - #[error("{0}")] + #[error("Duplicate: {0}")] DuplicatedResource(String), /// Vocab size exceeds the capacity of the target token type. @@ -28,7 +32,7 @@ pub enum WCError { }, /// Vocabulary data is inconsistent. - #[error("{0}")] + #[error("Vocab Conflict: {0}")] VocabConflict(String), /// Token value out of range for the target type. diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 5c6406f4..ea9e9f84 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -1,5 +1,3 @@ -use std::println; - use tokenizers::{ ModelWrapper::BPE, PreTokenizerWrapper, @@ -132,16 +130,20 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult = Default::default(); let decoder = tok.get_added_tokens_decoder(); + /* println!("Debug: {:#?}", decoder); + */ for (t, at) in decoder.iter() { span_config.specials_mut().add_str_word(&at.content, *t); diff --git a/crates/wordchipper/src/support/mod.rs b/crates/wordchipper/src/support/mod.rs index 6e643bfc..478f9c56 100644 --- a/crates/wordchipper/src/support/mod.rs +++ b/crates/wordchipper/src/support/mod.rs @@ -2,6 +2,7 @@ #[cfg(feature = "concurrent")] pub mod concurrency; + pub mod ranges; pub mod regex; pub mod resources; @@ -9,3 +10,4 @@ pub mod slices; pub mod strings; pub mod timers; pub mod traits; +pub mod with_ok_or_panic; diff --git a/crates/wordchipper/src/support/with_ok_or_panic.rs b/crates/wordchipper/src/support/with_ok_or_panic.rs new file mode 100644 index 00000000..82be4482 --- /dev/null +++ b/crates/wordchipper/src/support/with_ok_or_panic.rs @@ -0,0 +1,61 @@ +//! # Result Utilities +//! +//! Methods for [`std::result::Result`] manipulation. + +use core::fmt::Display; + +/// Extension trait for `Result` to add `ok_or_panic` method. +pub trait WithOkOrPanic { + /// Unwraps the `Result`, or panics with the error message. + /// + /// This differs from the behavior of [`Result::unwrap`] + /// in that the [`Debug`] format of the wrapped error is used + /// directly as the panic message; and not escaped. + fn ok_or_panic(self) -> T; +} + +impl WithOkOrPanic for Result +where + E: Display, +{ + fn ok_or_panic(self) -> T { + match self { + Ok(t) => t, + Err(e) => panic!("{e}"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + WCError, + WCResult, + prelude::*, + }; + + fn try_example( + value: i32, + throw: bool, + ) -> WCResult { + if throw { + Err(WCError::External("throwing".to_string())) + } else { + Ok(value) + } + } + + #[test] + fn test_expect_unwrap() { + let result = try_example(42, false); + assert_eq!(result.ok_or_panic(), 42); + } + + #[should_panic(expected = "throwing")] + #[test] + fn test_expect_unwrap_panic() { + let result = try_example(42, true); + result.ok_or_panic(); + } +} diff --git a/crates/wordchipper/src/vocab/pair_vocab.rs b/crates/wordchipper/src/vocab/pair_vocab.rs index ac0ff320..b8d39b4d 100644 --- a/crates/wordchipper/src/vocab/pair_vocab.rs +++ b/crates/wordchipper/src/vocab/pair_vocab.rs @@ -47,19 +47,26 @@ pub fn try_validate_pair_map( } } + const ORPHAN_TOKENS_ERROR: &str = indoc::indoc! {r#" + This vocab has orphan tokens, which wordchipper does not yet support. + See: https://github.com/zspacelabs/wordchipper/issues/386 + "#}; + for (&pair, &t) in pairs.iter() { for pt in [pair.0, pair.1] { let is_pair_target = pair_targets.contains(&pt); let byte_target = byte_vocab.get_byte(pt); if is_pair_target && let Some(b) = byte_target { - return Err(crate::WCError::VocabConflict(crate::alloc::format!( - "Pair {pair:?} -> {t:?} parent {pt:?} is a pair target and byte target: {b:0x?}" + return Err(crate::WCError::NotImplemented(crate::alloc::format!( + "{PRE}Pair {pair:?} -> {t:?} parent {pt:?} is a pair target and byte target: {b:0x?}", + PRE = ORPHAN_TOKENS_ERROR, ))); } if !is_pair_target && byte_target.is_none() { - return Err(crate::WCError::VocabConflict(crate::alloc::format!( - "Pair {pair:?} -> {t:?} parent {pt:?} is not defined" + return Err(crate::WCError::NotImplemented(crate::alloc::format!( + "{PRE}Pair {pair:?} -> {t:?} parent {pt:?} is not defined", + PRE = ORPHAN_TOKENS_ERROR, ))); } } diff --git a/crates/wordchipper/src/vocab/span_vocab.rs b/crates/wordchipper/src/vocab/span_vocab.rs index da96c756..3b3234ff 100644 --- a/crates/wordchipper/src/vocab/span_vocab.rs +++ b/crates/wordchipper/src/vocab/span_vocab.rs @@ -3,6 +3,7 @@ use crate::{ WCResult, alloc::vec::Vec, + support::with_ok_or_panic::WithOkOrPanic, types::{ TokenType, WCHashMap, @@ -96,7 +97,7 @@ impl SpanMapVocab { pub fn from_byte_vocab(byte_vocab: ByteMapVocab) -> Self { let span_map: SpanTokenMap = byte_vocab.span_pairs().collect(); - Self::new(byte_vocab, span_map).unwrap() + Self::new(byte_vocab, span_map).ok_or_panic() } /// Build a [`Self`] from a [`SpanTokenMap`]. @@ -126,7 +127,7 @@ impl SpanMapVocab { let byte_vocab: ByteMapVocab = ByteMapVocab::from_byte_to_token(&byte_to_token); - Self::new(byte_vocab, span_map).unwrap() + Self::new(byte_vocab, span_map).ok_or_panic() } /// Initialize a [`SpanMapVocab`]. @@ -244,7 +245,7 @@ impl SpanMapVocab { } } - PairMapVocab::::new(byte_vocab, pairs).unwrap() + PairMapVocab::::new(byte_vocab, pairs).ok_or_panic() } } From acf7aef284f0224341b65cc8aad126eea5535ca2 Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Mon, 11 May 2026 01:03:27 -0700 Subject: [PATCH 07/11] Fix no-std build. --- crates/wordchipper/src/pretrained/factory/vocab_query.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/wordchipper/src/pretrained/factory/vocab_query.rs b/crates/wordchipper/src/pretrained/factory/vocab_query.rs index 1efbdb8b..4f48e93b 100644 --- a/crates/wordchipper/src/pretrained/factory/vocab_query.rs +++ b/crates/wordchipper/src/pretrained/factory/vocab_query.rs @@ -194,6 +194,7 @@ mod tests { use core::str::FromStr; use crate::{ + alloc::vec, prelude::*, pretrained::factory::vocab_query::VocabQuery, }; From 024e94845c08340f1bd83b98a4cf22463b5f7338 Mon Sep 17 00:00:00 2001 From: Kyle Grover Date: Mon, 11 May 2026 18:14:35 -0700 Subject: [PATCH 08/11] ground BPE pair vocab & add qwen specific logos lexer - build pair maps only from grounded BPE parents so pruned vocabularies stay valid - add Qwen3.5 Hugging Face patterns and logos and regex-automata accelerators - register regex-automata transforms through inventory hooks instead of a central table --- .../src/pretrained/huggingface/hf_factory.rs | 25 +- .../src/pretrained/huggingface/mod.rs | 1 + .../src/pretrained/huggingface/patterns.rs | 42 +++ .../src/spanners/span_lexers/logos/mod.rs | 2 + .../src/spanners/span_lexers/logos/qwen35.rs | 284 ++++++++++++++++++ .../spanners/span_lexers/regex_automata.rs | 35 +++ crates/wordchipper/src/vocab/span_vocab.rs | 30 +- 7 files changed, 416 insertions(+), 3 deletions(-) create mode 100644 crates/wordchipper/src/pretrained/huggingface/patterns.rs create mode 100644 crates/wordchipper/src/spanners/span_lexers/logos/qwen35.rs diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index ea9e9f84..5d448bd8 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -10,6 +10,8 @@ use tokenizers::{ tokenizer::Tokenizer, }; +use super::patterns::QWEN35_PATTERN; + use crate::{ LabeledVocab, UnifiedTokenVocab, @@ -41,10 +43,18 @@ use crate::{ }, }; +fn canonicalize_pattern(pattern: RegexPattern) -> RegexPattern { + if pattern.as_str() == QWEN35_PATTERN.as_str() { + return QWEN35_PATTERN.to_pattern(); + } + + pattern +} + fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { match &s.pattern { - SplitPattern::Regex(r) => Ok(r.clone().into()), + SplitPattern::Regex(r) => Ok(canonicalize_pattern(r.clone().into())), _ => Err(WCError::External("Split without Regex pattern".into())), } } @@ -137,7 +147,6 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult = Default::default(); let decoder = tok.get_added_tokens_decoder(); @@ -264,3 +273,15 @@ impl VocabProvider for HFVocabProvider { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_canonicalize_qwen35_pattern() { + let canonical = canonicalize_pattern(QWEN35_PATTERN.as_str().into()); + + assert_eq!(canonical, QWEN35_PATTERN.to_pattern()); + } +} diff --git a/crates/wordchipper/src/pretrained/huggingface/mod.rs b/crates/wordchipper/src/pretrained/huggingface/mod.rs index 3f74b586..8eab57a8 100644 --- a/crates/wordchipper/src/pretrained/huggingface/mod.rs +++ b/crates/wordchipper/src/pretrained/huggingface/mod.rs @@ -1,3 +1,4 @@ //! # `HuggingFace` Pretrained Models +pub(crate) mod patterns; mod hf_factory; diff --git a/crates/wordchipper/src/pretrained/huggingface/patterns.rs b/crates/wordchipper/src/pretrained/huggingface/patterns.rs new file mode 100644 index 00000000..d3eb4bd2 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/patterns.rs @@ -0,0 +1,42 @@ +//! Shared regex patterns for Hugging Face tokenizers. + +use crate::{ + join_patterns, + support::regex::ConstRegexPattern, +}; + +/// The Qwen3.5 pretrained vocabulary word pattern. +/// +/// Shared by the Qwen3.5 tokenizer family loaded via Hugging Face. +pub(crate) const QWEN35_PATTERN: ConstRegexPattern = ConstRegexPattern::Fancy(join_patterns!( + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)", + r"[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+", + r"\p{N}", + r" ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", +)); + +/// Transformed Qwen3.5 pattern for `regex-automata` (lookahead removed). +/// +/// The `\s+(?!\S)` branch is collapsed to `\s+`; post-processing restores +/// the original end-of-whitespace semantics. +pub(crate) const QWEN35_PATTERN_RA: &str = join_patterns!( + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)", + r"[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+", + r"\p{N}", + r" ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+", +); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_patterns_compile() { + assert!(QWEN35_PATTERN.compile().is_ok()); + } +} \ No newline at end of file diff --git a/crates/wordchipper/src/spanners/span_lexers/logos/mod.rs b/crates/wordchipper/src/spanners/span_lexers/logos/mod.rs index 6ee68ebc..11084066 100644 --- a/crates/wordchipper/src/spanners/span_lexers/logos/mod.rs +++ b/crates/wordchipper/src/spanners/span_lexers/logos/mod.rs @@ -45,6 +45,8 @@ macro_rules! logos_lexer { pub mod cl100k; pub mod gpt2_family; pub mod o200k; +#[cfg(feature = "huggingface")] +pub mod qwen35; pub mod r50k; #[cfg(any(test, feature = "testing"))] diff --git a/crates/wordchipper/src/spanners/span_lexers/logos/qwen35.rs b/crates/wordchipper/src/spanners/span_lexers/logos/qwen35.rs new file mode 100644 index 00000000..0c09980e --- /dev/null +++ b/crates/wordchipper/src/spanners/span_lexers/logos/qwen35.rs @@ -0,0 +1,284 @@ +//! Logos DFA lexer for the Qwen3.5 pattern. +//! +//! Shared by the Qwen3.5 tokenizer family exposed through the Hugging Face +//! loader. + +use logos::Logos; + +use super::gpt2_family::{ + Gpt2FamilyLogos, + Gpt2FamilyTokenRole, +}; +use crate::pretrained::huggingface::patterns::QWEN35_PATTERN; + +/// Logos token variants for Qwen3.5. +#[derive(Logos, Debug, PartialEq, Clone)] +pub(crate) enum Qwen35Token { + #[regex(r"[\p{L}\p{M}]+")] + Letters, + + #[regex(r"[^\r\n\p{L}\p{N}][\p{L}\p{M}]+")] + PrefixedLetters, + + #[regex(r"\p{N}")] + Digit, + + #[regex(r" ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*")] + Punctuation, + + #[regex(r"\s*[\r\n]+")] + Newline, + + #[regex(r"[^\S\r\n]+")] + Whitespace, +} + +impl Gpt2FamilyLogos<'_> for Qwen35Token { + fn family_role(&self) -> Gpt2FamilyTokenRole { + match self { + Self::Letters => Gpt2FamilyTokenRole::Word { + check_contraction: false, + first_char_is_letter: true, + }, + Self::PrefixedLetters => Gpt2FamilyTokenRole::Word { + check_contraction: true, + first_char_is_letter: false, + }, + Self::Digit | Self::Newline => Gpt2FamilyTokenRole::Standalone, + Self::Punctuation => Gpt2FamilyTokenRole::Punctuation, + Self::Whitespace => Gpt2FamilyTokenRole::Whitespace, + } + } +} + +logos_lexer! { + /// Logos DFA word scanner for Qwen3.5. + pub struct Qwen35Lexer; + token = Qwen35Token; + pattern = QWEN35_PATTERN; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + alloc::{ + sync::Arc, + vec, + vec::Vec, + }, + spanners::{ + SpanRef, + TextSpanner, + span_lexers::{ + LexerTextSpanner, + SpanLexer, + }, + }, + }; + + fn spanner(lexer: impl SpanLexer + 'static) -> LexerTextSpanner { + LexerTextSpanner::new(Arc::new(lexer), None) + } + + #[test] + fn test_qwen35_common() { + crate::spanners::span_lexers::logos::testutil::common_lexer_tests( + crate::alloc::boxed::Box::new(Qwen35Lexer), + ); + } + + #[cfg(feature = "testing")] + #[test] + fn test_qwen35_matches_reference() { + use crate::spanners::span_lexers::accelerators::testutil::assert_matches_reference_lexer; + use crate::support::regex::RegexPattern; + + let ref_lexer = RegexPattern::Fancy(QWEN35_PATTERN.as_str().into()) + .compile() + .expect("reference pattern compiles"); + + let test_lexer = Qwen35Lexer; + + let samples = &[ + "hello world", + " hello world ", + "hello world", + "It's a test. Don't panic!", + "I'm she'll they've we'd he's", + "I'M SHE'LL THEY'VE WE'D HE'S", + "foo123bar 456 789", + "abc 1 2 3 def", + " ", + " ", + "", + "a", + "Hello, World! How are you?", + "price is $100.00!", + "foo bar baz", + "\t\t\thello", + "end with spaces ", + "\u{4e16}\u{754c}\u{4f60}\u{597d}", + "mixed\n\n content\there", + "foo'bar'baz", + "don't I'll she's", + "'There 'The 'really", + "'t 'T 're 'RE 'll 'll 'd 'D", + "hello\nworld", + "hello \n world", + " \n spaces around newline \n ", + "!@#$%", + "hello!world", + "test\r\nwindows", + "\u{00e9}clair caf\u{00e9}", + "e\u{0301} combining accent", + "\u{0300}standalone mark", + ]; + + for sample in samples { + assert_matches_reference_lexer(sample, &ref_lexer, &test_lexer); + } + } + + #[test] + fn test_basic_splitting() { + let s = spanner(Qwen35Lexer); + + assert_eq!( + s.split_spans("hello world", None), + vec![SpanRef::Word(0..5), SpanRef::Word(5..11)], + ); + } + + #[test] + fn test_single_digits() { + let s = spanner(Qwen35Lexer); + let text = "abc123"; + let spans = s.split_spans(text, None); + let words: Vec<&str> = spans + .iter() + .filter_map(|span| match span { + SpanRef::Word(range) => Some(&text[range.clone()]), + _ => None, + }) + .collect(); + + assert_eq!(words, vec!["abc", "1", "2", "3"]); + } + + #[test] + fn test_digits_do_not_absorb_space() { + let s = spanner(Qwen35Lexer); + let text = "abc 1"; + let spans = s.split_spans(text, None); + let words: Vec<&str> = spans + .iter() + .filter_map(|span| match span { + SpanRef::Word(range) => Some(&text[range.clone()]), + _ => None, + }) + .collect(); + + assert_eq!(words, vec!["abc", " ", "1"]); + } + + #[test] + fn test_contractions_case_insensitive() { + let s = spanner(Qwen35Lexer); + let text = "don't I'll SHE'S THEY'RE"; + let spans = s.split_spans(text, None); + let words: Vec<&str> = spans + .iter() + .filter_map(|span| match span { + SpanRef::Word(range) => Some(&text[range.clone()]), + _ => None, + }) + .collect(); + + assert!(words.contains(&"don"), "expected \"don\" in {:?}", words); + assert!(words.contains(&"'t"), "expected \"'t\" in {:?}", words); + assert!(words.contains(&"'ll"), "expected \"'ll\" in {:?}", words); + assert!(words.contains(&"'S"), "expected \"'S\" in {:?}", words); + assert!(words.contains(&"'RE"), "expected \"'RE\" in {:?}", words); + } + + #[test] + fn test_contraction_followed_by_more_letters() { + let s = spanner(Qwen35Lexer); + let text = "'There"; + let spans = s.split_spans(text, None); + let words: Vec<&str> = spans + .iter() + .filter_map(|span| match span { + SpanRef::Word(range) => Some(&text[range.clone()]), + _ => None, + }) + .collect(); + + assert_eq!(words, vec!["'T", "here"]); + } + + #[test] + fn test_standalone_contraction() { + let s = spanner(Qwen35Lexer); + + assert_eq!(s.split_spans("'t", None), vec![SpanRef::Word(0..2)],); + assert_eq!(s.split_spans("'ll", None), vec![SpanRef::Word(0..3)],); + } + + #[test] + fn test_marks_attach_to_letters() { + let s = spanner(Qwen35Lexer); + let text = "e\u{0301}clair"; + let spans = s.split_spans(text, None); + + assert_eq!(spans.len(), 1); + assert!(matches!(&spans[0], SpanRef::Word(range) if range == &(0..text.len()))); + } + + #[test] + fn test_marks_not_punctuation() { + let s = spanner(Qwen35Lexer); + let text = "\u{0300}"; + let spans = s.split_spans(text, None); + + assert_eq!(spans, vec![SpanRef::Word(0..text.len())]); + } + + #[test] + fn test_no_case_split() { + let s = spanner(Qwen35Lexer); + + assert_eq!(s.split_spans("CamelCase", None), vec![SpanRef::Word(0..9)],); + assert_eq!( + s.split_spans("getElementById", None), + vec![SpanRef::Word(0..14)], + ); + assert_eq!(s.split_spans("HTMLParser", None), vec![SpanRef::Word(0..10)],); + } + + #[test] + fn test_newline_absorbs_preceding_whitespace() { + let s = spanner(Qwen35Lexer); + + assert_eq!(s.split_spans(" \n", None), vec![SpanRef::Word(0..3)],); + } + + #[test] + fn test_punctuation_optional_space() { + let s = spanner(Qwen35Lexer); + + assert_eq!(s.split_spans(" !", None), vec![SpanRef::Word(0..2)],); + assert_eq!( + s.split_spans(" !", None), + vec![SpanRef::Word(0..1), SpanRef::Word(1..3)], + ); + } + + #[test] + fn test_punctuation_trailing_newlines() { + let s = spanner(Qwen35Lexer); + + assert_eq!(s.split_spans("!\n\n", None), vec![SpanRef::Word(0..3)],); + } +} \ No newline at end of file diff --git a/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs b/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs index f7981212..e7f52951 100644 --- a/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs +++ b/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs @@ -21,6 +21,11 @@ use spin::Mutex; #[cfg(feature = "concurrent")] use crate::support::concurrency::PoolToy; +#[cfg(feature = "huggingface")] +use crate::pretrained::huggingface::patterns::{ + QWEN35_PATTERN, + QWEN35_PATTERN_RA, +}; use crate::{ alloc::sync::Arc, prelude::*, @@ -55,6 +60,13 @@ const KNOWN_TRANSFORMS: &[(&str, &str, bool)] = &[ ), ]; +#[cfg(feature = "huggingface")] +const HF_KNOWN_TRANSFORMS: &[(&str, &str, bool)] = &[( + QWEN35_PATTERN.as_str(), + QWEN35_PATTERN_RA, + true, +)]; + /// `SpanLexer` using `regex_automata::meta::Regex` with pooled or single-mutex /// caches. struct RegexAutomataLexer { @@ -185,6 +197,23 @@ pub(crate) fn try_build( } } + #[cfg(feature = "huggingface")] + for &(original, transformed, has_newline_branch) in HF_KNOWN_TRANSFORMS { + if pattern == original { + let regex = match Regex::new(transformed) { + Ok(r) => r, + Err(e) => { + log::warn!( + "regex-automata failed to compile known transform (len={}): {e}", + transformed.len(), + ); + return None; + } + }; + return Some(build_lexer(regex, has_newline_branch, max_pool)); + } + } + // Fallback: try compiling directly (for patterns without lookaheads). let regex = Regex::new(pattern).ok()?; Some(build_lexer(regex, false, max_pool)) @@ -297,6 +326,12 @@ mod tests { check_pattern(OA_O200K_BASE_PATTERN.as_str()); } + #[cfg(feature = "huggingface")] + #[test] + fn test_qwen35_matches_reference() { + check_pattern(QWEN35_PATTERN.as_str()); + } + #[test] fn test_basic_whitespace_truncation() { // "hello world" with r50k: " " is truncated to " ", then diff --git a/crates/wordchipper/src/vocab/span_vocab.rs b/crates/wordchipper/src/vocab/span_vocab.rs index 3b3234ff..0a17a011 100644 --- a/crates/wordchipper/src/vocab/span_vocab.rs +++ b/crates/wordchipper/src/vocab/span_vocab.rs @@ -228,21 +228,35 @@ impl SpanMapVocab { .map(|(chunk, &token)| (token, chunk.as_ref())) .collect(); - for token in self.tokens() { + let mut tokens: Vec = self.tokens().into_iter().collect(); + tokens.sort_by_key(|token| token_to_span.get(token).map_or(1, |span| span.len())); + + let mut grounded: WCHashSet = byte_vocab.tokens(); + + for token in tokens { let span = token_to_span[&token]; if span.len() <= 1 { continue; } + + let mut added = false; for p in 1..span.len() { let pre = &span[..p]; let post = &span[p..]; if let Some(a) = self.lookup_token(pre) && let Some(b) = self.lookup_token(post) + && grounded.contains(&a) + && grounded.contains(&b) { pairs.insert((a, b), token); + added = true; } } + + if added { + grounded.insert(token); + } } PairMapVocab::::new(byte_vocab, pairs).ok_or_panic() @@ -389,4 +403,18 @@ mod tests { .collect::>() ); } + + #[test] + fn test_build_pair_vocab_omits_undecomposable_span_token() { + type T = u32; + + let mut span_map: SpanTokenMap = Default::default(); + span_map.insert("abc".as_bytes().to_vec(), 300); + + let vocab = SpanMapVocab::from(span_map); + let pair_vocab = vocab.to_pair_vocab(); + + assert!(pair_vocab.pair_map().is_empty()); + assert!(!pair_vocab.tokens().contains(&300)); + } } From 17b7611c7289d211050a704083c154ea13ff6b54 Mon Sep 17 00:00:00 2001 From: Kyle Grover Date: Mon, 11 May 2026 18:21:58 -0700 Subject: [PATCH 09/11] improve qwen lexer integration - register regex-automata transforms via inventory hooks add RegexAutomataTransformHook alongside existing lexer accelerator hooks move OpenAI and Qwen regex-automata transform registration into their pattern modules remove the HF-specific Qwen pattern canonicalization from the loader --- .../src/pretrained/huggingface/hf_factory.rs | 24 +---- .../src/pretrained/huggingface/patterns.rs | 5 + .../src/pretrained/openai/patterns.rs | 13 +++ .../src/spanners/span_lexers/accelerators.rs | 53 +++++++++++ .../spanners/span_lexers/regex_automata.rs | 94 +++++-------------- 5 files changed, 94 insertions(+), 95 deletions(-) diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 5d448bd8..33ef601b 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -10,8 +10,6 @@ use tokenizers::{ tokenizer::Tokenizer, }; -use super::patterns::QWEN35_PATTERN; - use crate::{ LabeledVocab, UnifiedTokenVocab, @@ -43,18 +41,10 @@ use crate::{ }, }; -fn canonicalize_pattern(pattern: RegexPattern) -> RegexPattern { - if pattern.as_str() == QWEN35_PATTERN.as_str() { - return QWEN35_PATTERN.to_pattern(); - } - - pattern -} - fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result { fn split_regex(s: &tokenizers::pre_tokenizers::split::Split) -> Result { match &s.pattern { - SplitPattern::Regex(r) => Ok(canonicalize_pattern(r.clone().into())), + SplitPattern::Regex(r) => Ok(r.clone().into()), _ => Err(WCError::External("Split without Regex pattern".into())), } } @@ -273,15 +263,3 @@ impl VocabProvider for HFVocabProvider { } } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_canonicalize_qwen35_pattern() { - let canonical = canonicalize_pattern(QWEN35_PATTERN.as_str().into()); - - assert_eq!(canonical, QWEN35_PATTERN.to_pattern()); - } -} diff --git a/crates/wordchipper/src/pretrained/huggingface/patterns.rs b/crates/wordchipper/src/pretrained/huggingface/patterns.rs index d3eb4bd2..858a8349 100644 --- a/crates/wordchipper/src/pretrained/huggingface/patterns.rs +++ b/crates/wordchipper/src/pretrained/huggingface/patterns.rs @@ -2,6 +2,7 @@ use crate::{ join_patterns, + spanners::span_lexers::accelerators::RegexAutomataTransformHook, support::regex::ConstRegexPattern, }; @@ -31,6 +32,10 @@ pub(crate) const QWEN35_PATTERN_RA: &str = join_patterns!( r"\s+", ); +inventory::submit! { + RegexAutomataTransformHook::new(QWEN35_PATTERN, QWEN35_PATTERN_RA, true) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/wordchipper/src/pretrained/openai/patterns.rs b/crates/wordchipper/src/pretrained/openai/patterns.rs index efce3368..7c6cb16c 100644 --- a/crates/wordchipper/src/pretrained/openai/patterns.rs +++ b/crates/wordchipper/src/pretrained/openai/patterns.rs @@ -2,6 +2,7 @@ use crate::{ join_patterns, + spanners::span_lexers::accelerators::RegexAutomataTransformHook, support::regex::ConstRegexPattern, }; @@ -104,6 +105,18 @@ pub(crate) const OA_O200K_BASE_PATTERN_RA: &str = join_patterns!( r"\s+", ); +inventory::submit! { + RegexAutomataTransformHook::new(OA_R50K_BASE_PATTERN, OA_R50K_BASE_PATTERN_RA, false) +} + +inventory::submit! { + RegexAutomataTransformHook::new(OA_CL100K_BASE_PATTERN, OA_CL100K_BASE_PATTERN_RA, true) +} + +inventory::submit! { + RegexAutomataTransformHook::new(OA_O200K_BASE_PATTERN, OA_O200K_BASE_PATTERN_RA, true) +} + #[cfg(test)] mod test { use super::*; diff --git a/crates/wordchipper/src/spanners/span_lexers/accelerators.rs b/crates/wordchipper/src/spanners/span_lexers/accelerators.rs index cacf8b05..68dda34d 100644 --- a/crates/wordchipper/src/spanners/span_lexers/accelerators.rs +++ b/crates/wordchipper/src/spanners/span_lexers/accelerators.rs @@ -40,6 +40,22 @@ pub struct RegexAcceleratorHook { } inventory::collect!(RegexAcceleratorHook); +/// Inventory hook for regex-automata pattern transforms. +/// +/// Some patterns require a lookahead-free or possessive-free variant before +/// they can be compiled by `regex-automata`. +pub struct RegexAutomataTransformHook { + /// The exact source regex pattern. + pub pattern: ConstRegexPattern, + + /// The transformed pattern accepted by `regex-automata`. + pub transformed_pattern: &'static str, + + /// Whether whitespace truncation should ignore newline-containing spans. + pub has_newline_branch: bool, +} +inventory::collect!(RegexAutomataTransformHook); + impl RegexAcceleratorHook { /// Setup a new regex accelerator hook. pub const fn new( @@ -50,6 +66,21 @@ impl RegexAcceleratorHook { } } +impl RegexAutomataTransformHook { + /// Setup a new regex-automata transform hook. + pub const fn new( + pattern: ConstRegexPattern, + transformed_pattern: &'static str, + has_newline_branch: bool, + ) -> Self { + Self { + pattern, + transformed_pattern, + has_newline_branch, + } + } +} + /// Get a regex accelerator. /// /// ## Returns @@ -64,9 +95,23 @@ pub fn get_regex_accelerator(pattern: &str) -> Option> { None } +/// Get a registered `regex-automata` transform. +pub fn get_regex_automata_transform(pattern: &str) -> Option<(&'static str, bool)> { + for hook in inventory::iter:: { + if hook.pattern.as_str() == pattern { + return Some((hook.transformed_pattern, hook.has_newline_branch)); + } + } + None +} + #[cfg(test)] mod tests { use super::*; + use crate::pretrained::openai::patterns::{ + OA_CL100K_BASE_PATTERN, + OA_CL100K_BASE_PATTERN_RA, + }; #[test] fn test_unknown_pattern_returns_none() { @@ -74,6 +119,14 @@ mod tests { get_regex_accelerator("not_a_real_pattern_that_would_ever_be_registered").is_none() ); } + + #[test] + fn test_known_regex_automata_transform_returns_hook() { + assert_eq!( + get_regex_automata_transform(OA_CL100K_BASE_PATTERN.as_str()), + Some((OA_CL100K_BASE_PATTERN_RA, true)) + ); + } } /// Testing utilities for developing accelerated replacement [`SpanLexer`]s. diff --git a/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs b/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs index e7f52951..a9873d14 100644 --- a/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs +++ b/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs @@ -21,52 +21,15 @@ use spin::Mutex; #[cfg(feature = "concurrent")] use crate::support::concurrency::PoolToy; -#[cfg(feature = "huggingface")] -use crate::pretrained::huggingface::patterns::{ - QWEN35_PATTERN, - QWEN35_PATTERN_RA, -}; use crate::{ alloc::sync::Arc, prelude::*, - pretrained::openai::patterns::{ - OA_CL100K_BASE_PATTERN, - OA_CL100K_BASE_PATTERN_RA, - OA_O200K_BASE_PATTERN, - OA_O200K_BASE_PATTERN_RA, - OA_R50K_BASE_PATTERN, - OA_R50K_BASE_PATTERN_RA, + spanners::span_lexers::{ + SpanLexer, + accelerators::get_regex_automata_transform, }, - spanners::span_lexers::SpanLexer, }; -/// Known pattern transforms: (original fancy pattern, transformed RA pattern, -/// `has_newline_branch`). -const KNOWN_TRANSFORMS: &[(&str, &str, bool)] = &[ - ( - OA_R50K_BASE_PATTERN.as_str(), - OA_R50K_BASE_PATTERN_RA, - false, - ), - ( - OA_CL100K_BASE_PATTERN.as_str(), - OA_CL100K_BASE_PATTERN_RA, - true, - ), - ( - OA_O200K_BASE_PATTERN.as_str(), - OA_O200K_BASE_PATTERN_RA, - true, - ), -]; - -#[cfg(feature = "huggingface")] -const HF_KNOWN_TRANSFORMS: &[(&str, &str, bool)] = &[( - QWEN35_PATTERN.as_str(), - QWEN35_PATTERN_RA, - true, -)]; - /// `SpanLexer` using `regex_automata::meta::Regex` with pooled or single-mutex /// caches. struct RegexAutomataLexer { @@ -180,38 +143,18 @@ pub(crate) fn try_build( pattern: &str, max_pool: Option, ) -> Option> { - // Check known transforms. - for &(original, transformed, has_newline_branch) in KNOWN_TRANSFORMS { - if pattern == original { - let regex = match Regex::new(transformed) { - Ok(r) => r, - Err(e) => { - log::warn!( - "regex-automata failed to compile known transform (len={}): {e}", - transformed.len(), - ); - return None; - } - }; - return Some(build_lexer(regex, has_newline_branch, max_pool)); - } - } - - #[cfg(feature = "huggingface")] - for &(original, transformed, has_newline_branch) in HF_KNOWN_TRANSFORMS { - if pattern == original { - let regex = match Regex::new(transformed) { - Ok(r) => r, - Err(e) => { - log::warn!( - "regex-automata failed to compile known transform (len={}): {e}", - transformed.len(), - ); - return None; - } - }; - return Some(build_lexer(regex, has_newline_branch, max_pool)); - } + if let Some((transformed, has_newline_branch)) = get_regex_automata_transform(pattern) { + let regex = match Regex::new(transformed) { + Ok(r) => r, + Err(e) => { + log::warn!( + "regex-automata failed to compile known transform (len={}): {e}", + transformed.len(), + ); + return None; + } + }; + return Some(build_lexer(regex, has_newline_branch, max_pool)); } // Fallback: try compiling directly (for patterns without lookaheads). @@ -263,9 +206,16 @@ mod tests { use super::*; use crate::{ + pretrained::openai::patterns::{ + OA_CL100K_BASE_PATTERN, + OA_O200K_BASE_PATTERN, + OA_R50K_BASE_PATTERN, + }, spanners::span_lexers::accelerators::testutil::assert_matches_reference_lexer, support::regex::RegexWrapper, }; + #[cfg(feature = "huggingface")] + use crate::pretrained::huggingface::patterns::QWEN35_PATTERN; fn ref_lexer(pattern: &str) -> RegexWrapper { crate::support::regex::RegexPattern::Fancy(pattern.to_string()) From bbf7e62ead2f314c7cd1f14ae84a4f17bdd1739a Mon Sep 17 00:00:00 2001 From: Kyle Grover Date: Mon, 11 May 2026 20:21:41 -0700 Subject: [PATCH 10/11] fix HF Qwen loading parity and add benchmark coverage --- Cargo.lock | 1 + Cargo.toml | 1 + crates/wordchipper/Cargo.toml | 1 + .../token_span_encoder/token_span_encoder.rs | 3 + .../src/pretrained/huggingface/hf_factory.rs | 80 +++++++++++++- .../src/spanners/spanning_config.rs | 44 +++++++- crates/wordchipper/src/support/mod.rs | 1 + .../wordchipper/src/support/normalization.rs | 100 ++++++++++++++++++ crates/wordchipper/src/vocab/unified_vocab.rs | 32 +++++- dev-crates/wordchipper-bench/Cargo.toml | 6 +- dev-crates/wordchipper-bench/README.md | 2 + .../benches/qwen_encoding_single.rs | 83 +++++++++++++++ dev-crates/wordchipper-bench/src/lib.rs | 6 ++ 13 files changed, 353 insertions(+), 7 deletions(-) create mode 100644 crates/wordchipper/src/support/normalization.rs create mode 100644 dev-crates/wordchipper-bench/benches/qwen_encoding_single.rs diff --git a/Cargo.lock b/Cargo.lock index 74ed0b65..32b618de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4347,6 +4347,7 @@ dependencies = [ "tokenizers", "tracing", "unicode-general-category", + "unicode-normalization", "wordchipper-disk-cache", ] diff --git a/Cargo.toml b/Cargo.toml index 832fa430..d6b17ec5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,7 @@ once_cell = { version = "1.21.0", default-features = false } regex = { version = "1.12.3", default-features = false } regex-automata = { version = "0.4", default-features = false } unicode-general-category = { version = "1.1.0", default-features = false } +unicode-normalization = { version = "0.1.25", default-features = false } ringbuffer = { version = "0.16", default-features = false } strum = { version = "0.27.0", default-features = false, features = ["derive"] } thiserror = { version = "2.0.10", default-features = false } diff --git a/crates/wordchipper/Cargo.toml b/crates/wordchipper/Cargo.toml index 5e9c8ba1..c84d09b8 100644 --- a/crates/wordchipper/Cargo.toml +++ b/crates/wordchipper/Cargo.toml @@ -139,6 +139,7 @@ serde_json = { workspace = true, optional = true } aho-corasick = { workspace = true } foldhash = { workspace = true, optional = true } unicode-general-category = { workspace = true } +unicode-normalization = { workspace = true } # "download" feature deps: wordchipper-disk-cache = { version = "0.9.1", path = "../wordchipper-disk-cache", optional = true, default-features = false } diff --git a/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs b/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs index f85373d3..34ad073c 100644 --- a/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs +++ b/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs @@ -108,6 +108,9 @@ impl TokenEncoder for TokenSpanEncoder { } } + let normalized_text = self.vocab.spanning().normalize_text(text); + let text = normalized_text.as_ref(); + self.spanner .for_each_split_span(text, special_filter, &mut |span_ref| { se.encode_append_span_ref(&self.vocab, text, span_ref, tokens); diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 33ef601b..65189782 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -6,6 +6,7 @@ use tokenizers::{ Sequence, Split, }, + tokenizer::NormalizerWrapper, pre_tokenizers::split::SplitPattern, tokenizer::Tokenizer, }; @@ -31,6 +32,7 @@ use crate::{ }, spanners::TextSpanningConfig, support::{ + normalization::TextNormalizer, regex::RegexPattern, resources::ResourceLoader, }, @@ -75,6 +77,28 @@ fn extract_pattern(pt: Option<&PreTokenizerWrapper>) -> Result WCResult { + match normalizer { + NormalizerWrapper::NFC(_) => Ok(TextNormalizer::NFC), + NormalizerWrapper::NFD(_) => Ok(TextNormalizer::NFD), + NormalizerWrapper::NFKC(_) => Ok(TextNormalizer::NFKC), + NormalizerWrapper::NFKD(_) => Ok(TextNormalizer::NFKD), + NormalizerWrapper::Sequence(sequence) => sequence + .as_ref() + .iter() + .map(extract_text_normalizer) + .collect::>>() + .map(TextNormalizer::Sequence), + _ => Err(WCError::External(crate::alloc::format!( + "unsupported huggingface normalizer: {normalizer:?}" + ))), + } +} + +fn extract_normalizer(normalizer: Option<&NormalizerWrapper>) -> WCResult> { + normalizer.map(extract_text_normalizer).transpose() +} + /// Converts bytes to Unicode characters. /// See /// @@ -116,6 +140,9 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult = TextSpanningConfig::from_pattern(pattern); + if let Some(normalizer) = extract_normalizer(tok.get_normalizer())? { + span_config = span_config.with_normalizer(normalizer); + } let BPE(bpe) = tok.get_model() else { return Err(WCError::External( @@ -145,7 +172,22 @@ pub fn vocab_from_hf_tokenizer(tok: &Tokenizer) -> WCResult { /// Regex pattern for word splitting. pattern: RegexPattern, + /// Optional text normalizer applied before spanning. + normalizer: Option, + /// Special tokens vocabulary. specials: SpecialVocab, } @@ -40,6 +46,7 @@ impl TextSpanningConfig { { Self { pattern: pattern.into(), + normalizer: None, specials: SpecialVocab::default(), } } @@ -61,6 +68,20 @@ impl TextSpanningConfig { } } + /// Set the optional text normalizer. + /// + /// ## Arguments + /// * `normalizer` - The new text normalizer. + pub fn with_normalizer( + self, + normalizer: TextNormalizer, + ) -> Self { + Self { + normalizer: Some(normalizer), + ..self + } + } + /// Set the special tokens vocabulary. /// /// ## Arguments @@ -100,6 +121,7 @@ impl TextSpanningConfig { pub fn to_token_type(&self) -> WCResult> { Ok(TextSpanningConfig:: { pattern: self.pattern.clone(), + normalizer: self.normalizer.clone(), specials: self.specials.to_token_type()?, }) } @@ -109,6 +131,21 @@ impl TextSpanningConfig { &self.pattern } + /// Get the optional text normalizer. + pub fn normalizer(&self) -> Option<&TextNormalizer> { + self.normalizer.as_ref() + } + + /// Normalize text prior to spanning. + pub fn normalize_text<'a>( + &self, + text: &'a str, + ) -> crate::alloc::borrow::Cow<'a, str> { + self.normalizer() + .map(|normalizer| normalizer.normalize(text)) + .unwrap_or_else(|| crate::alloc::borrow::Cow::Borrowed(text)) + } + /// Get the special tokens vocabulary. pub fn specials(&self) -> &SpecialVocab { &self.specials @@ -130,6 +167,7 @@ mod tests { use super::*; use crate::{ alloc::string::ToString, + support::normalization::TextNormalizer, vocab::SpecialVocab, }; @@ -155,5 +193,9 @@ mod tests { let config = config.with_specials(specials.clone()); assert_eq!(config.specials(), &specials); + + let config = config.with_normalizer(TextNormalizer::NFC); + assert_eq!(config.normalizer(), Some(&TextNormalizer::NFC)); + assert_eq!(config.normalize_text("e\u{301}clair").as_ref(), "éclair"); } } diff --git a/crates/wordchipper/src/support/mod.rs b/crates/wordchipper/src/support/mod.rs index 478f9c56..98576578 100644 --- a/crates/wordchipper/src/support/mod.rs +++ b/crates/wordchipper/src/support/mod.rs @@ -3,6 +3,7 @@ #[cfg(feature = "concurrent")] pub mod concurrency; +pub mod normalization; pub mod ranges; pub mod regex; pub mod resources; diff --git a/crates/wordchipper/src/support/normalization.rs b/crates/wordchipper/src/support/normalization.rs new file mode 100644 index 00000000..c9b3d82a --- /dev/null +++ b/crates/wordchipper/src/support/normalization.rs @@ -0,0 +1,100 @@ +//! # Text Normalization + +use crate::alloc::{ + borrow::Cow, + string::String, + vec::Vec, +}; +use unicode_normalization::{ + UnicodeNormalization, + is_nfc, + is_nfd, + is_nfkc, + is_nfkd, +}; + +/// Text normalizers that can be applied before spanning. +#[derive(Debug, Clone, PartialEq)] +pub enum TextNormalizer { + /// Normalize with Unicode NFC. + NFC, + + /// Normalize with Unicode NFD. + NFD, + + /// Normalize with Unicode NFKC. + NFKC, + + /// Normalize with Unicode NFKD. + NFKD, + + /// Apply the normalizers in-order. + Sequence(Vec), +} + +impl TextNormalizer { + /// Normalize `text`, borrowing the input when no rewrite is needed. + pub fn normalize<'a>( + &self, + text: &'a str, + ) -> Cow<'a, str> { + match self { + Self::NFC => normalize_if_needed(text, is_nfc, |s| s.nfc().collect()), + Self::NFD => normalize_if_needed(text, is_nfd, |s| s.nfd().collect()), + Self::NFKC => normalize_if_needed(text, is_nfkc, |s| s.nfkc().collect()), + Self::NFKD => normalize_if_needed(text, is_nfkd, |s| s.nfkd().collect()), + Self::Sequence(normalizers) => { + let mut current: Option = None; + + for normalizer in normalizers { + let input = current.as_deref().unwrap_or(text); + if let Cow::Owned(next) = normalizer.normalize(input) { + current = Some(next); + } + } + + current.map(Cow::Owned).unwrap_or_else(|| Cow::Borrowed(text)) + } + } + } +} + +fn normalize_if_needed<'a, F, G>( + text: &'a str, + is_normalized: F, + normalize: G, +) -> Cow<'a, str> +where + F: Fn(&str) -> bool, + G: Fn(&str) -> String, +{ + if is_normalized(text) { + Cow::Borrowed(text) + } else { + Cow::Owned(normalize(text)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_nfc_normalizer_recomposes_decomposed_unicode() { + let normalized = TextNormalizer::NFC.normalize("e\u{301}clair cafe\u{301}"); + assert_eq!(normalized.as_ref(), "éclair café"); + } + + #[test] + fn test_nfc_normalizer_borrows_already_normalized_text() { + let normalized = TextNormalizer::NFC.normalize("éclair café"); + assert!(matches!(normalized, Cow::Borrowed(_))); + } + + #[test] + fn test_sequence_normalizer_applies_in_order() { + let normalized = TextNormalizer::Sequence(vec![TextNormalizer::NFD, TextNormalizer::NFC]) + .normalize("éclair café"); + assert_eq!(normalized.as_ref(), "éclair café"); + } +} \ No newline at end of file diff --git a/crates/wordchipper/src/vocab/unified_vocab.rs b/crates/wordchipper/src/vocab/unified_vocab.rs index 915f88fe..5b3f8b8b 100644 --- a/crates/wordchipper/src/vocab/unified_vocab.rs +++ b/crates/wordchipper/src/vocab/unified_vocab.rs @@ -122,15 +122,22 @@ impl UnifiedTokenVocab { )); } - let tokens = span_vocab.tokens(); - if tokens != pair_vocab.tokens() { + let span_tokens = span_vocab.tokens(); + let pair_tokens = pair_vocab.tokens(); + if !pair_tokens.is_subset(&span_tokens) { + let missing = pair_tokens + .difference(&span_tokens) + .copied() + .collect::>(); return Err(WCError::VocabConflict( - "span vocab and pair vocab have different token sets".into(), + crate::alloc::format!( + "pair vocab contains tokens missing from span vocab: {missing:?}" + ), )); } for t in span_config.specials().tokens() { - if tokens.contains(&t) { + if span_tokens.contains(&t) { let span = span_config.specials().lookup_span(&t).unwrap(); let special = string_from_utf8_lossy(span.to_vec()); return Err(WCError::VocabConflict(crate::alloc::format!( @@ -361,4 +368,21 @@ mod tests { assert_eq!(vocab64.lookup_token("at".as_bytes()), Some(300 as u64)); assert_eq!(vocab64.lookup_token("ate".as_bytes()), Some(301 as u64)); } + + #[test] + fn test_init_allows_undecomposable_span_tokens() { + type T = u32; + + let mut span_vocab: SpanTokenMap = Default::default(); + span_vocab.insert("abc".as_bytes().to_vec(), 300); + let span_vocab: SpanMapVocab = span_vocab.into(); + + let seg_config = TextSpanningConfig::from_pattern(r"\w+"); + + let vocab = UnifiedTokenVocab::from_span_vocab(seg_config, span_vocab).unwrap(); + + assert_eq!(vocab.lookup_token("abc".as_bytes()), Some(300)); + assert!(!vocab.pair_vocab().tokens().contains(&300)); + assert!(vocab.pair_vocab().pair_map().is_empty()); + } } diff --git a/dev-crates/wordchipper-bench/Cargo.toml b/dev-crates/wordchipper-bench/Cargo.toml index cf3e72b3..55f69c61 100644 --- a/dev-crates/wordchipper-bench/Cargo.toml +++ b/dev-crates/wordchipper-bench/Cargo.toml @@ -10,7 +10,7 @@ publish = false workspace = true [dependencies] -wordchipper = { path = "../../crates/wordchipper", features = ["default", "download"] } +wordchipper = { path = "../../crates/wordchipper", features = ["default", "download", "huggingface"] } wordchipper-data = { path = "../wordchipper-data" } divan-parser = { path = "../divan-parser" } @@ -43,3 +43,7 @@ harness = false [[bench]] name = "decoding_single" harness = false + +[[bench]] +name = "qwen_encoding_single" +harness = false diff --git a/dev-crates/wordchipper-bench/README.md b/dev-crates/wordchipper-bench/README.md index 58a54872..93c850f4 100644 --- a/dev-crates/wordchipper-bench/README.md +++ b/dev-crates/wordchipper-bench/README.md @@ -11,6 +11,7 @@ and HuggingFace tokenizers. | `encoding_parallel` | Batch encoding via rayon (`try_encode_batch`) | | `decoding_single` | Single-string decoding | | `spanning` | Text spanning (regex vs logos DFA) | +| `qwen_encoding_single` | Single-string Qwen encode vs HF tokenizers | ### Encoder Variants @@ -32,6 +33,7 @@ cargo bench -p wordchipper-bench --bench encoding_single cargo bench -p wordchipper-bench --bench encoding_parallel cargo bench -p wordchipper-bench --bench decoding_single cargo bench -p wordchipper-bench --bench spanning +cargo bench -p wordchipper-bench --bench qwen_encoding_single # Filter by name cargo bench -p wordchipper-bench --bench encoding_single -- diverse diff --git a/dev-crates/wordchipper-bench/benches/qwen_encoding_single.rs b/dev-crates/wordchipper-bench/benches/qwen_encoding_single.rs new file mode 100644 index 00000000..d8c8d3c9 --- /dev/null +++ b/dev-crates/wordchipper-bench/benches/qwen_encoding_single.rs @@ -0,0 +1,83 @@ +#![allow(missing_docs)] + +use divan::{ + Bencher, + black_box, + counter::BytesCount, +}; +use wordchipper::{ + TokenEncoderOptions, +}; +use wordchipper_bench::{ + HF_QWEN35, + WC_QWEN35, + load_cached_encoder, +}; + +#[global_allocator] +static ALLOC: divan::AllocProfiler = divan::AllocProfiler::system(); + +fn main() { + divan::main(); +} + +static DIVERSE_CORPUS: &str = include_str!("data/multilingual.txt"); +static ENGLISH_CORPUS: &str = include_str!("data/english.txt"); + +fn diverse_text() -> String { + DIVERSE_CORPUS.repeat(10) +} + +fn english_text() -> String { + ENGLISH_CORPUS.repeat(10) +} + +fn bench_wc( + bencher: Bencher, + text: &str, +) { + let encoder = load_cached_encoder::(WC_QWEN35, TokenEncoderOptions::default()); + + bencher + .counter(BytesCount::new(text.len())) + .bench(|| encoder.try_encode(black_box(text), None).unwrap()); +} + +fn bench_hf( + bencher: Bencher, + text: &str, +) { + let tok = tokenizers::Tokenizer::from_pretrained(HF_QWEN35, None).unwrap(); + + bencher + .counter(BytesCount::new(text.len())) + .bench(|| tok.encode(black_box(text), true).unwrap()); +} + +mod english { + use super::*; + + #[divan::bench] + fn wordchipper(bencher: Bencher) { + bench_wc(bencher, &english_text()); + } + + #[divan::bench] + fn tokenizers(bencher: Bencher) { + bench_hf(bencher, &english_text()); + } +} + +mod diverse { + use super::*; + + #[divan::bench] + fn wordchipper(bencher: Bencher) { + bench_wc(bencher, &diverse_text()); + } + + #[divan::bench] + fn tokenizers(bencher: Bencher) { + bench_hf(bencher, &diverse_text()); + } +} \ No newline at end of file diff --git a/dev-crates/wordchipper-bench/src/lib.rs b/dev-crates/wordchipper-bench/src/lib.rs index 7d898c52..3461f3b1 100644 --- a/dev-crates/wordchipper-bench/src/lib.rs +++ b/dev-crates/wordchipper-bench/src/lib.rs @@ -28,6 +28,12 @@ pub const HF_CL100K: &str = "Xenova/text-embedding-ada-002"; /// The huggingface/tokenizers model to use for `o200k_base`. pub const HF_O200K: &str = "Xenova/gpt-4o"; +/// The wordchipper model identifier for Qwen 3.5 0.8B via the HF loader. +pub const WC_QWEN35: &str = "hf:Qwen/Qwen3.5-0.8B"; + +/// The huggingface/tokenizers model to use for Qwen 3.5 0.8B. +pub const HF_QWEN35: &str = "Qwen/Qwen3.5-0.8B"; + /// The shared disk cache for benchmarks. static DISK_CACHE: OnceLock> = OnceLock::new(); From 25770543dc4bf8c77dad6bba974882c5b115dd79 Mon Sep 17 00:00:00 2001 From: Kyle Grover Date: Mon, 11 May 2026 20:24:28 -0700 Subject: [PATCH 11/11] move HF input normalization ownership to UnifiedTokenVocab --- .../token_span_encoder/token_span_encoder.rs | 2 +- .../src/pretrained/huggingface/hf_factory.rs | 26 +++++--- .../src/spanners/spanning_config.rs | 44 +------------ crates/wordchipper/src/vocab/unified_vocab.rs | 63 ++++++++++++++++++- 4 files changed, 82 insertions(+), 53 deletions(-) diff --git a/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs b/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs index 34ad073c..2a0c4fe4 100644 --- a/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs +++ b/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs @@ -108,7 +108,7 @@ impl TokenEncoder for TokenSpanEncoder { } } - let normalized_text = self.vocab.spanning().normalize_text(text); + let normalized_text = self.vocab.normalize_text(text); let text = normalized_text.as_ref(); self.spanner diff --git a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs index 65189782..7d740ee6 100644 --- a/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -99,6 +99,15 @@ fn extract_normalizer(normalizer: Option<&NormalizerWrapper>) -> WCResult