diff --git a/Cargo.lock b/Cargo.lock index 7a8b0c04..32b618de 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", @@ -4337,6 +4347,7 @@ dependencies = [ "tokenizers", "tracing", "unicode-general-category", + "unicode-normalization", "wordchipper-disk-cache", ] diff --git a/Cargo.toml b/Cargo.toml index 71bb8825..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 } @@ -82,6 +83,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 855bdfc9..c84d09b8 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. @@ -118,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"] } @@ -131,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 } @@ -141,6 +150,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/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/encoders/token_span_encoder/token_span_encoder.rs b/crates/wordchipper/src/encoders/token_span_encoder/token_span_encoder.rs index f85373d3..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,6 +108,9 @@ impl TokenEncoder for TokenSpanEncoder { } } + let normalized_text = self.vocab.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/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/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..4f48e93b 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)] @@ -181,6 +194,7 @@ mod tests { use core::str::FromStr; use crate::{ + alloc::vec, prelude::*, pretrained::factory::vocab_query::VocabQuery, }; @@ -202,6 +216,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 new file mode 100644 index 00000000..7d740ee6 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -0,0 +1,355 @@ +use tokenizers::{ + ModelWrapper::BPE, + PreTokenizerWrapper, + PreTokenizerWrapper::{ + ByteLevel, + Sequence, + Split, + }, + tokenizer::NormalizerWrapper, + pre_tokenizers::split::SplitPattern, + tokenizer::Tokenizer, +}; + +use crate::{ + LabeledVocab, + UnifiedTokenVocab, + VocabDescription, + VocabIndex, + VocabQuery, + WCError, + WCHashMap, + WCHashSet, + WCResult, + alloc::sync::Arc, + prelude::*, + pretrained::{ + factory::{ + VocabProvider, + VocabProviderInventoryHook, + }, + openai::OA_GPT2_PATTERN, + }, + spanners::TextSpanningConfig, + support::{ + normalization::TextNormalizer, + regex::RegexPattern, + resources::ResourceLoader, + }, + vocab::{ + ByteMapVocab, + SpanMapVocab, + SpanTokenMap, + }, +}; + +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()), + _ => Err(WCError::External("Split without Regex pattern".into())), + } + } + match pt { + Some(Split(s)) => split_regex(s), + 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() { + 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())), + } +} + +fn extract_text_normalizer(normalizer: &NormalizerWrapper) -> 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() +} + +fn normalize_input<'a>( + normalizer: Option<&TextNormalizer>, + text: &'a str, +) -> crate::alloc::borrow::Cow<'a, str> { + normalizer + .map(|normalizer| normalizer.normalize(text)) + .unwrap_or_else(|| crate::alloc::borrow::Cow::Borrowed(text)) +} + +/// 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() +} + +/// 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 input_normalizer = extract_normalizer(tok.get_normalizer())?; + 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(); + + /* + println!( + "Debug: {:?}", + hf_vocab.iter().find(|(_, id)| **id == 157513) + ); + */ + + let mut special_tokens: WCHashSet = Default::default(); + + let decoder = tok.get_added_tokens_decoder(); + /* + println!("Debug: {:#?}", decoder); + */ + + for (t, at) in decoder.iter() { + let special_content = if at.normalized { + normalize_input(input_normalizer.as_ref(), &at.content).into_owned() + } else { + let normalized = normalize_input(input_normalizer.as_ref(), &at.content); + if normalized.as_ref() != at.content { + return Err(WCError::External(crate::alloc::format!( + "unsupported non-normalized special token under text normalizer: {:?}", + at.content + ))); + } + at.content.clone() + }; + + span_config + .specials_mut() + .add_str_word(&special_content, *t); + special_tokens.insert(*t); + } + + // 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) { + continue; + } 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); + } + } + + if span_config.specials().len() != special_tokens.len() { + return Err(WCError::External(format!( + "hf vocab identifies {} special tokens, but only {} special tokens found in span_config", + special_tokens.len(), + span_config.specials().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 = UnifiedTokenVocab::from_span_vocab(span_config, span_vocab)?; + let vocab = if let Some(normalizer) = input_normalizer { + vocab.with_input_normalizer(normalizer) + } else { + vocab + }; + let vocab: Arc> = Arc::new(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! { + 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> { + if let Some(schema) = query.schema() + && schema != "hf" + { + return Err(WCError::ResourceNotFound(query.to_string())); + } + + match Tokenizer::from_pretrained(query.clone().with_schema(None).to_string(), None) { + Ok(tok) => { + let vocab = vocab_from_hf_tokenizer(&tok)?; + + let mut context = vec!["hf"]; + if query.path().is_some() { + context.push(query.path().unwrap()); + } + 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"); + + Ok(LabeledVocab::new(descr, vocab)) + } + Err(_) => Err(WCError::ResourceNotFound(query.to_string())), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokenizers::normalizers::{ + Lowercase, + NFC, + NormalizerWrapper, + Sequence, + }; + + #[test] + fn test_extract_normalizer_maps_nfc() { + assert_eq!( + extract_normalizer(Some(&NormalizerWrapper::NFC(NFC))).unwrap(), + Some(TextNormalizer::NFC) + ); + } + + #[test] + fn test_extract_normalizer_maps_sequence() { + let sequence = Sequence::new(vec![NormalizerWrapper::NFC(NFC)]); + + assert_eq!( + extract_normalizer(Some(&NormalizerWrapper::Sequence(sequence))).unwrap(), + Some(TextNormalizer::Sequence(vec![TextNormalizer::NFC])) + ); + } + + #[test] + fn test_extract_normalizer_rejects_unsupported_wrapper() { + let error = extract_normalizer(Some(&NormalizerWrapper::Lowercase(Lowercase))).unwrap_err(); + + assert!(matches!(error, WCError::External(_))); + } +} diff --git a/crates/wordchipper/src/pretrained/huggingface/mod.rs b/crates/wordchipper/src/pretrained/huggingface/mod.rs new file mode 100644 index 00000000..8eab57a8 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/mod.rs @@ -0,0 +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..858a8349 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/patterns.rs @@ -0,0 +1,47 @@ +//! Shared regex patterns for Hugging Face tokenizers. + +use crate::{ + join_patterns, + spanners::span_lexers::accelerators::RegexAutomataTransformHook, + 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+", +); + +inventory::submit! { + RegexAutomataTransformHook::new(QWEN35_PATTERN, QWEN35_PATTERN_RA, true) +} + +#[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/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, 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/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..a9873d14 100644 --- a/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs +++ b/crates/wordchipper/src/spanners/span_lexers/regex_automata.rs @@ -24,37 +24,12 @@ use crate::support::concurrency::PoolToy; 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, - ), -]; - /// `SpanLexer` using `regex_automata::meta::Regex` with pooled or single-mutex /// caches. struct RegexAutomataLexer { @@ -168,21 +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)); - } + 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). @@ -234,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()) @@ -297,6 +276,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/support/mod.rs b/crates/wordchipper/src/support/mod.rs index 6e643bfc..98576578 100644 --- a/crates/wordchipper/src/support/mod.rs +++ b/crates/wordchipper/src/support/mod.rs @@ -2,6 +2,8 @@ #[cfg(feature = "concurrent")] pub mod concurrency; + +pub mod normalization; pub mod ranges; pub mod regex; pub mod resources; @@ -9,3 +11,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/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/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..0a17a011 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`]. @@ -227,24 +228,38 @@ 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).unwrap() + PairMapVocab::::new(byte_vocab, pairs).ok_or_panic() } } @@ -388,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)); + } } diff --git a/crates/wordchipper/src/vocab/unified_vocab.rs b/crates/wordchipper/src/vocab/unified_vocab.rs index 915f88fe..0fcd0aac 100644 --- a/crates/wordchipper/src/vocab/unified_vocab.rs +++ b/crates/wordchipper/src/vocab/unified_vocab.rs @@ -6,9 +6,15 @@ use crate::{ WCError, WCHashSet, WCResult, - alloc::vec::Vec, + alloc::{ + borrow::Cow, + vec::Vec, + }, spanners::TextSpanningConfig, - support::strings::string_from_utf8_lossy, + support::{ + normalization::TextNormalizer, + strings::string_from_utf8_lossy, + }, vocab::{ ByteMapVocab, PairMapVocab, @@ -62,6 +68,9 @@ pub struct UnifiedTokenVocab { /// Text Spanning Configuration spanning: TextSpanningConfig, + /// Optional text normalizer applied before encoding. + input_normalizer: Option, + /// ``{ Vec -> T }`` vocabulary. span_vocab: SpanMapVocab, @@ -122,15 +131,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!( @@ -141,6 +157,7 @@ impl UnifiedTokenVocab { Ok(Self { spanning: span_config, + input_normalizer: None, span_vocab, pair_vocab, }) @@ -154,16 +171,41 @@ impl UnifiedTokenVocab { pub fn to_token_type(&self) -> WCResult> { Ok(UnifiedTokenVocab:: { spanning: self.spanning.to_token_type::()?, + input_normalizer: self.input_normalizer.clone(), span_vocab: self.span_vocab.to_token_type::()?, pair_vocab: self.pair_vocab.to_token_type::()?, }) } + /// Attach an input normalizer used before encoding. + pub fn with_input_normalizer( + mut self, + normalizer: TextNormalizer, + ) -> Self { + self.input_normalizer = Some(normalizer); + self + } + /// Get the [`TextSpanningConfig`]. pub fn spanning(&self) -> &TextSpanningConfig { &self.spanning } + /// Get the optional input normalizer. + pub fn input_normalizer(&self) -> Option<&TextNormalizer> { + self.input_normalizer.as_ref() + } + + /// Normalize text prior to encoding. + pub fn normalize_text<'a>( + &self, + text: &'a str, + ) -> Cow<'a, str> { + self.input_normalizer() + .map(|normalizer| normalizer.normalize(text)) + .unwrap_or_else(|| Cow::Borrowed(text)) + } + /// Get the `{ (T, T) -> T }` [`PairMapVocab`]. pub fn pair_vocab(&self) -> &PairMapVocab { &self.pair_vocab @@ -278,6 +320,7 @@ mod tests { use super::*; use crate::{ spanners::TextSpanningConfig, + support::normalization::TextNormalizer, vocab::{ PairTokenMap, SpanMapVocab, @@ -361,4 +404,44 @@ 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_input_normalizer() { + 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 vocab = UnifiedTokenVocab::from_span_vocab( + TextSpanningConfig::from_pattern(r"\w+"), + span_vocab, + ) + .unwrap() + .with_input_normalizer(TextNormalizer::NFC); + + assert_eq!(vocab.input_normalizer(), Some(&TextNormalizer::NFC)); + assert_eq!(vocab.normalize_text("e\u{301}clair").as_ref(), "éclair"); + assert_eq!( + vocab.to_token_type::().unwrap().input_normalizer(), + Some(&TextNormalizer::NFC) + ); + } + + #[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();