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 855bdfc9..5e9c8ba1 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"] } @@ -141,6 +149,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/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..ea9e9f84 --- /dev/null +++ b/crates/wordchipper/src/pretrained/huggingface/hf_factory.rs @@ -0,0 +1,266 @@ +use tokenizers::{ + ModelWrapper::BPE, + PreTokenizerWrapper, + PreTokenizerWrapper::{ + ByteLevel, + Sequence, + Split, + }, + 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::{ + 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())), + } +} + +/// 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 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) + ); + */ + + // 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(); + 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: 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! { + 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())), + } + } +} 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, 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() } }