diff --git a/src/added_tokens.rs b/src/added_tokens.rs index 8ddbd1f..dd3d429 100644 --- a/src/added_tokens.rs +++ b/src/added_tokens.rs @@ -16,6 +16,11 @@ pub struct AddedTokens { /// Token lengths (in bytes) indexed by token ID, for matched tokens only. /// Non-added token IDs map to 0. token_lens: Vec, + /// Per-ID `lstrip`/`rstrip` flags. When set, a match absorbs adjacent + /// Unicode whitespace (left/right) into the token span, matching HF's + /// `AddedToken` behavior. Indexed by token ID; `false` for non-added IDs. + lstrip: Vec, + rstrip: Vec, /// Distinct first bytes of all added token strings. Used to quickly skip /// positions that cannot start any token via SIMD memchr. start_bytes: Vec, @@ -59,6 +64,8 @@ impl AddedTokens { let max_id = configs.iter().map(|c| c.id).max().unwrap_or(0); let mut token_lens = vec![0usize; (max_id + 1) as usize]; + let mut lstrip = vec![false; (max_id + 1) as usize]; + let mut rstrip = vec![false; (max_id + 1) as usize]; let mut id_to_content = HashMap::with_capacity(configs.len()); let mut special_ids = HashSet::new(); @@ -69,6 +76,8 @@ impl AddedTokens { .iter() .map(|c| { token_lens[c.id as usize] = c.content.len(); + lstrip[c.id as usize] = c.lstrip; + rstrip[c.id as usize] = c.rstrip; id_to_content.insert(c.id, c.content.clone()); content_to_id.insert(c.content.clone(), c.id); if c.special { @@ -102,6 +111,8 @@ impl AddedTokens { Ok(Some(Self { daac, token_lens, + lstrip, + rstrip, start_bytes, max_token_len, id_to_content, @@ -204,11 +215,12 @@ impl AddedTokens { if let Some(m) = self.daac.leftmost_find_iter(window).next() && m.start() == 0 { - if pos > prev_end { - segments.push(Segment::Text(&input[prev_end..pos])); + let (start, end) = self.strip_bounds(input, m.value(), pos, pos + m.end(), prev_end); + if start > prev_end { + segments.push(Segment::Text(&input[prev_end..start])); } segments.push(Segment::Token(m.value())); - prev_end = pos + m.end(); + prev_end = end; } } @@ -222,17 +234,57 @@ impl AddedTokens { segments } + /// Expand a match span to absorb adjacent Unicode whitespace per the + /// token's `lstrip`/`rstrip` flags, matching HuggingFace `AddedToken` + /// behavior. Absorbed whitespace is excluded from the surrounding text. + /// + /// `lstrip` extends the start left over whitespace, bounded by `floor` + /// (the end of the previous segment) so it never reclaims already-emitted + /// text. `rstrip` extends the end right over whitespace. When two strip + /// tokens share a whitespace run, the left token's `rstrip` consumes it + /// first (via the advanced `floor`), so the right token's `lstrip` finds + /// none — mirroring HF's left-to-right resolution. + fn strip_bounds( + &self, + input: &str, + id: u32, + mut start: usize, + mut end: usize, + floor: usize, + ) -> (usize, usize) { + if self.lstrip[id as usize] { + for (rel_i, c) in input[floor..start].char_indices().rev() { + if c.is_whitespace() { + start = floor + rel_i; + } else { + break; + } + } + } + if self.rstrip[id as usize] { + for c in input[end..].chars() { + if c.is_whitespace() { + end += c.len_utf8(); + } else { + break; + } + } + } + (start, end) + } + /// Full-scan fallback for >3 distinct start bytes. fn split_full_scan<'a>(&self, input: &'a str) -> Vec> { let mut segments = Vec::new(); let mut prev_end = 0; for m in self.daac.leftmost_find_iter(input) { - if m.start() > prev_end { - segments.push(Segment::Text(&input[prev_end..m.start()])); + let (start, end) = self.strip_bounds(input, m.value(), m.start(), m.end(), prev_end); + if start > prev_end { + segments.push(Segment::Text(&input[prev_end..start])); } segments.push(Segment::Token(m.value())); - prev_end = m.end(); + prev_end = end; } if prev_end < input.len() { @@ -567,4 +619,117 @@ mod tests { vec![Segment::Token(9), Segment::Token(9), Segment::Token(9)] ); } + + // ── lstrip / rstrip whitespace absorption ─────────────────────────── + + fn make_strip_config(id: u32, content: &str, lstrip: bool, rstrip: bool) -> AddedTokenConfig { + AddedTokenConfig { + id, + content: content.to_string(), + single_word: false, + lstrip, + rstrip, + normalized: false, + special: true, + } + } + + #[test] + fn lstrip_absorbs_leading_whitespace() { + let configs = vec![make_strip_config(1, "", true, false)]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + // The spaces before are absorbed into the token span; "ab" remains, + // trailing "cd" is untouched. + assert_eq!( + at.split("ab cd"), + vec![Segment::Text("ab"), Segment::Token(1), Segment::Text("cd")] + ); + } + + #[test] + fn rstrip_absorbs_trailing_whitespace() { + let configs = vec![make_strip_config(1, "", false, true)]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + assert_eq!( + at.split("ab cd"), + vec![Segment::Text("ab"), Segment::Token(1), Segment::Text("cd")] + ); + } + + #[test] + fn strip_both_sides() { + let configs = vec![make_strip_config(1, "", true, true)]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + assert_eq!( + at.split("ab \t \n cd"), + vec![Segment::Text("ab"), Segment::Token(1), Segment::Text("cd")] + ); + } + + #[test] + fn no_strip_keeps_whitespace() { + let configs = vec![make_strip_config(1, "", false, false)]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + assert_eq!( + at.split("ab cd"), + vec![ + Segment::Text("ab "), + Segment::Token(1), + Segment::Text(" cd"), + ] + ); + } + + #[test] + fn adjacent_strip_tokens_share_whitespace() { + // The <|im_end|>\n<|im_start|> case from Phi-4: both tokens strip both + // sides. The single \n between them must be absorbed exactly once and + // produce no text token. + let configs = vec![ + make_strip_config(1, "<|im_end|>", true, true), + make_strip_config(2, "<|im_start|>", true, true), + ]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + assert_eq!( + at.split("hi<|im_end|>\n<|im_start|>user"), + vec![ + Segment::Text("hi"), + Segment::Token(1), + Segment::Token(2), + Segment::Text("user"), + ] + ); + } + + #[test] + fn lstrip_bounded_by_previous_token() { + // A preceding token's span must not be reclaimed by the next token's + // lstrip: there is no whitespace between the tokens here. + let configs = vec![ + make_strip_config(1, "", false, false), + make_strip_config(2, "", true, false), + ]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + assert_eq!( + at.split(""), + vec![Segment::Token(1), Segment::Token(2)] + ); + } + + #[test] + fn strip_via_full_scan_path() { + // >3 distinct start bytes forces the full-scan path; rstrip must still + // absorb trailing whitespace there. + let configs = vec![ + make_strip_config(1, "", false, true), + make_config(2, "[SEP]"), + make_config(3, "{pad}"), + make_config(4, "|mask|"), + ]; + let at = AddedTokens::from_configs(&configs).unwrap().unwrap(); + assert_eq!( + at.split(" x"), + vec![Segment::Token(1), Segment::Text("x")] + ); + } } diff --git a/src/decoders.rs b/src/decoders.rs index 5bf7e46..4daade8 100644 --- a/src/decoders.rs +++ b/src/decoders.rs @@ -1,20 +1,42 @@ +mod byte_fallback; mod byte_level; +mod replace; use crate::json_structs::{DecoderConfig, DecoderKind}; +pub use self::byte_fallback::ByteFallbackDecoder; pub use self::byte_level::ByteLevelDecoder; +pub use self::replace::ReplaceDecoder; /// Errors from constructing or running a decoder. #[derive(Debug, thiserror::Error)] pub enum Error { + #[error("invalid config value: {0}")] + Json(#[from] serde_json::Error), + + #[error("regex error: {0}")] + Regex(#[from] fancy_regex::Error), + #[error("unsupported decoder type: {0}")] Unsupported(String), } +impl From for Error { + fn from(e: crate::normalizers::Error) -> Self { + match e { + crate::normalizers::Error::Json(j) => Self::Json(j), + crate::normalizers::Error::Regex(r) => Self::Regex(r), + crate::normalizers::Error::Unsupported(s) => Self::Unsupported(s), + } + } +} + /// A compiled decoder ready for use. #[derive(Debug)] pub enum Decoder { + ByteFallback(ByteFallbackDecoder), ByteLevel(ByteLevelDecoder), + Replace(ReplaceDecoder), Sequence(Vec), } @@ -22,7 +44,11 @@ impl Decoder { /// Build a decoder from its JSON configuration. pub fn from_config(config: DecoderConfig) -> Result { match config { + DecoderConfig::ByteFallback => Ok(Self::ByteFallback(ByteFallbackDecoder)), DecoderConfig::ByteLevel => Ok(Self::ByteLevel(ByteLevelDecoder)), + DecoderConfig::Replace { pattern, content } => Ok(Self::Replace( + ReplaceDecoder::from_config(pattern, content)?, + )), DecoderConfig::Fuse => Ok(Self::Sequence(vec![])), // identity/no-op DecoderConfig::Sequence { decoders } => { let steps = decoders @@ -49,7 +75,9 @@ impl Decoder { /// transforms the token list, and the final result is joined. pub fn decode_chain(&self, tokens: Vec) -> Result, Error> { match self { + Self::ByteFallback(bf) => Ok(bf.decode_chain(tokens)), Self::ByteLevel(bl) => Ok(bl.decode_chain(tokens)), + Self::Replace(repl) => Ok(repl.decode_chain(tokens)), Self::Sequence(steps) => { let mut current = tokens; for step in steps { diff --git a/src/decoders/byte_fallback.rs b/src/decoders/byte_fallback.rs new file mode 100644 index 0000000..758531b --- /dev/null +++ b/src/decoders/byte_fallback.rs @@ -0,0 +1,101 @@ +/// ByteFallback decoder. + +#[derive(Debug)] +pub struct ByteFallbackDecoder; + +impl ByteFallbackDecoder { + pub fn decode_chain(&self, tokens: Vec) -> Vec { + let mut out: Vec = Vec::new(); + let mut byte_run: Vec = Vec::new(); + + for token in tokens { + if let Some(b) = parse_byte_token(&token) { + byte_run.push(b); + continue; + } + + flush_byte_run(&mut out, &mut byte_run); + out.push(token); + } + + flush_byte_run(&mut out, &mut byte_run); + out + } +} + +fn parse_byte_token(token: &str) -> Option { + let bytes = token.as_bytes(); + if bytes.len() != 6 + || bytes[0] != b'<' + || bytes[1] != b'0' + || bytes[2] != b'x' + || bytes[5] != b'>' + { + return None; + } + + fn hex_value(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } + } + + Some(hex_value(bytes[3])? << 4 | hex_value(bytes[4])?) +} + +fn flush_byte_run(out: &mut Vec, byte_run: &mut Vec) { + if byte_run.is_empty() { + return; + } + + let bytes = std::mem::take(byte_run); + + match String::from_utf8(bytes) { + Ok(s) => out.push(s), + Err(err) => { + for _ in 0..err.into_bytes().len() { + out.push("\u{FFFD}".to_string()); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decode() { + let decoder = ByteFallbackDecoder; + + let res = decoder.decode_chain(vec!["Hey".into(), "friend!".into()]); + assert_eq!(res, vec!["Hey", "friend!"]); + + let res = decoder.decode_chain(vec!["<0x61>".into()]); + assert_eq!(res, vec!["a"]); + + let res = decoder.decode_chain(vec!["<0xE5>".into()]); + assert_eq!(res, vec!["�"]); + + let res = decoder.decode_chain(vec!["<0xE5>".into(), "<0x8f>".into()]); + assert_eq!(res, vec!["�", "�"]); + + // 叫 + let res = decoder.decode_chain(vec!["<0xE5>".into(), "<0x8f>".into(), "<0xab>".into()]); + assert_eq!(res, vec!["叫"]); + + let res = decoder.decode_chain(vec![ + "<0xE5>".into(), + "<0x8f>".into(), + "<0xab>".into(), + "a".into(), + ]); + assert_eq!(res, vec!["叫", "a"]); + + let res = decoder.decode_chain(vec!["<0xE5>".into(), "<0x8f>".into(), "a".into()]); + assert_eq!(res, vec!["�", "�", "a"]); + } +} \ No newline at end of file diff --git a/src/decoders/replace.rs b/src/decoders/replace.rs new file mode 100644 index 0000000..2317c90 --- /dev/null +++ b/src/decoders/replace.rs @@ -0,0 +1,66 @@ +use serde_json::Value; + +use super::Error; +use crate::normalizers::Replace; + +/// Replace decoder. +/// +/// Applies the [`Replace`](crate::normalizers::Replace) transform to each token +/// in the decode chain. The replacement semantics (literal vs. regex, literal +/// replacement content) are identical to the normalizer; the only difference is +/// that the decoder maps the transform over a list of tokens. +#[derive(Clone, Debug)] +pub struct ReplaceDecoder { + inner: Replace, +} + +impl ReplaceDecoder { + pub fn from_config(pattern: Value, content: String) -> Result { + Ok(Self { + inner: Replace::from_config(pattern, content)?, + }) + } + + pub fn decode_chain(&self, tokens: Vec) -> Vec { + tokens + .into_iter() + .map(|token| self.inner.normalize(&token).into_owned()) + .collect() + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn literal_replace_decoder() { + let dec = ReplaceDecoder::from_config(json!("▁"), " ".to_string()).unwrap(); + let out = dec.decode_chain(vec!["▁Hello".to_string(), "▁world".to_string()]); + assert_eq!(out, vec![" Hello", " world"]); + } + + #[test] + fn regex_replace_decoder() { + let dec = ReplaceDecoder::from_config(json!({"Regex": "[0-9]+"}), "#".to_string()).unwrap(); + let out = dec.decode_chain(vec!["a12b".to_string(), "34".to_string()]); + assert_eq!(out, vec!["a#b", "#"]); + } + + #[test] + fn literal_empty_pattern_is_not_noop() { + let dec = ReplaceDecoder::from_config(json!(""), "-".to_string()).unwrap(); + let out = dec.decode_chain(vec!["ab".to_string()]); + assert_eq!(out, vec!["-a-b-"]); + } + + #[test] + fn regex_replacement_content_is_literal() { + let dec = ReplaceDecoder::from_config(json!({"Regex": "([a-z]+)"}), "$1".to_string()) + .unwrap(); + let out = dec.decode_chain(vec!["abc".to_string()]); + assert_eq!(out, vec!["$1"]); + } +} diff --git a/src/lib.rs b/src/lib.rs index b931528..5ce99e8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,7 +20,7 @@ pub use self::{ PreTokenizerKind, TokenizerJson, }, models::Model, - normalizers::{Nfc, Normalizer}, + normalizers::{Nfc, Normalizer, Replace}, post_processors::PostProcessor, pre_tokenizers::{ByteLevel, PreTokenizer, Split, SplitBehavior}, }; @@ -125,6 +125,12 @@ pub struct Tokenizer { /// When the pre-tokenizer is `Sequence([Split, ByteLevel(bulk)])`, /// we store a Split-only pre-tokenizer and fuse ByteLevel into BPE. split_only: Option, + /// Whether to run vocab-aware (unbridgeable-bigram) splitting on encode. + /// True for metaspace models (e.g. Gemma) whose pre-tokenizer is a no-op + /// after normalization; false for ByteLevel models, whose regex `Split` + /// already produces word-level chunks. The pass is output-preserving, so + /// this flag only affects performance, never correctness. + needs_vocab_splitting: bool, } impl Tokenizer { @@ -146,6 +152,11 @@ impl Tokenizer { // Detect Sequence([Split, ByteLevel(bulk)]) for fused byte-level+BPE. let split_only = Self::detect_fused_byte_level(&pre_tokenizer); + // ByteLevel pipelines already chunk at word boundaries via their regex + // Split, so vocab-aware splitting adds cost without benefit. Only run it + // when no ByteLevel step is present (metaspace models like Gemma). + let needs_vocab_splitting = !Self::pre_tokenizer_contains_byte_level(&pre_tokenizer); + Ok(Self { added_tokens, normalizer, @@ -154,9 +165,23 @@ impl Tokenizer { post_processor, decoder, split_only, + needs_vocab_splitting, }) } + /// Recursively check whether a pre-tokenizer pipeline contains a `ByteLevel` + /// step (including inside a `Sequence`). + fn pre_tokenizer_contains_byte_level(pt: &Option) -> bool { + fn contains(pt: &PreTokenizer) -> bool { + match pt { + PreTokenizer::ByteLevel(_) => true, + PreTokenizer::Split(_) => false, + PreTokenizer::Sequence(steps) => steps.iter().any(contains), + } + } + pt.as_ref().is_some_and(contains) + } + /// If `pt` is `Sequence([Split, ByteLevel(bulk)])`, return a Split-only /// pre-tokenizer for fused mode. fn detect_fused_byte_level(pt: &Option) -> Option { @@ -288,6 +313,17 @@ impl Tokenizer { pt.pre_tokenize(&mut pts)?; } + // 2b. Break each text split at unbridgeable byte-pair boundaries. + // Split at positions where adjacent bytes never appear together in + // any vocab token. This is provably output-preserving and provides + // fine-grained word-level chunking for models that don't use + // ByteLevel (whose regex Split already chunks at word boundaries). + if self.needs_vocab_splitting { + if let Some(table) = self.model.bigram_bridge_table() { + split_on_unbridgeable_bigrams(&mut pts, table); + } + } + // 3. Tokenize each text split with the model. let ids = pts .tokenize(|text, out| self.model.tokenize_into(text, out)) @@ -315,6 +351,11 @@ impl Tokenizer { self.post_processor = pp; } + /// Replace the normalizer. + pub fn set_normalizer(&mut self, normalizer: Option) { + self.normalizer = normalizer; + } + pub fn post_process(&self, ids: Vec, add_special_tokens: bool) -> Vec { match &self.post_processor { Some(pp) => pp.post_process_single(ids, add_special_tokens), @@ -485,6 +526,60 @@ impl Tokenizer { } } + +/// Split each text chunk at unbridgeable byte-pair boundaries using the +/// vocab-derived bigram bridge table. +/// +/// A byte pair (prev, cur) is "unbridgeable" if no vocabulary token contains +/// that adjacent byte sequence. Splitting at such boundaries is provably +/// output-preserving: any BPE merge that spans the boundary would produce a +/// token containing that byte pair, which cannot exist in the vocabulary. +/// +/// This generalizes newline splitting and enables fine-grained word-level +/// chunking even in metaspace tokenizers like Gemma, where the pre-tokenizer +/// is effectively a no-op after normalization. +fn split_on_unbridgeable_bigrams( + pts: &mut PreTokenizedString, + bigram_table: &models::bpe::BigramBridgeTable, +) { + let bytes = pts.buffer().as_bytes(); + let mut new_splits = Vec::with_capacity(pts.splits().len() * 2); + + for split in pts.splits() { + if split.token_id.is_some() || split.range.is_empty() { + new_splits.push(split.clone()); + continue; + } + + let end = split.range.end; + let mut start = split.range.start; + + for i in (start + 1)..end { + let prev = bytes[i - 1]; + let cur = bytes[i]; + + // Split here if: + // 1. This byte pair never appears in vocab, AND + // 2. Position i is a UTF-8 char boundary (cur is not a continuation byte) + if !bigram_table.is_bridgeable(prev, cur) && (cur & 0xC0) != 0x80 { + new_splits.push(PtSplit { + range: start..i, + token_id: None, + }); + start = i; + } + } + + // Push the final segment + new_splits.push(PtSplit { + range: start..end, + token_id: None, + }); + } + + pts.refine_splits(new_splits); +} + // --------------------------------------------------------------------------- // Streaming decode // --------------------------------------------------------------------------- diff --git a/src/models.rs b/src/models.rs index 6111078..f0cd3d6 100644 --- a/src/models.rs +++ b/src/models.rs @@ -86,4 +86,11 @@ impl Model { Self::Bpe(bpe) => bpe.vocab_size(), } } + + /// Access the bigram bridge table for vocab-aware splitting. + pub fn bigram_bridge_table(&self) -> Option<&bpe::BigramBridgeTable> { + match self { + Self::Bpe(bpe) => Some(&bpe.bigram_bridge_table), + } + } } diff --git a/src/models/bpe.rs b/src/models/bpe.rs index ad81eee..0470676 100644 --- a/src/models/bpe.rs +++ b/src/models/bpe.rs @@ -92,6 +92,43 @@ impl MergeMap { } } +/// Bigram bridgeability table for vocab-aware safe splitting. +/// +/// For each of 256×256 possible byte pairs, records whether that pair +/// appears in any vocabulary token. Used to identify split points that +/// cannot be crossed by BPE merges. +#[derive(Clone, PartialEq)] +pub struct BigramBridgeTable { + /// Flat array: bridgeable[prev * 256 + cur] == true if some vocab + /// token contains adjacent bytes (prev, cur). + bridgeable: Box<[bool; 65536]>, +} + +impl BigramBridgeTable { + /// Check if a byte pair can be bridged by some vocab token. + #[inline(always)] + pub fn is_bridgeable(&self, prev: u8, cur: u8) -> bool { + self.bridgeable[prev as usize * 256 + cur as usize] + } +} + +/// Build a bigram bridge table by scanning all vocab tokens. +fn build_bigram_bridge_table(id_to_token: &[String]) -> BigramBridgeTable { + let mut bridgeable = Box::new([false; 65536]); + + for token_str in id_to_token { + let bytes = token_str.as_bytes(); + // Mark all adjacent byte pairs in this token as bridgeable + for window in bytes.windows(2) { + let prev = window[0] as usize; + let cur = window[1] as usize; + bridgeable[prev * 256 + cur] = true; + } + } + + BigramBridgeTable { bridgeable } +} + #[inline(always)] fn pack_pair(t1: u32, t2: u32) -> u64 { (t1 as u64) << 32 | t2 as u64 @@ -361,7 +398,6 @@ struct RawBpe { #[allow(dead_code)] fuse_unk: bool, #[serde(default)] - #[allow(dead_code)] byte_fallback: bool, #[serde(default)] ignore_merges: bool, @@ -590,10 +626,17 @@ pub struct Bpe { id_to_token: Vec, token_to_id: HashMap, byte_to_initial_token: [u32; 256], + byte_fallback_token_ids: [u32; 256], + /// Token id for each single ASCII-character string (`INVALID_TOKEN` when + /// absent). Fast path for the char-based merge engine, avoiding a HashMap + /// probe per character. + single_char_token: [u32; 128], ranked_merge_map: RankedMergeMap, byte_pair_initial: Vec<(u32, u32)>, merge_adj: MergeAdjacency, ignore_merges: bool, + byte_fallback: bool, + pub bigram_bridge_table: BigramBridgeTable, } impl TryFrom for Bpe { @@ -603,6 +646,7 @@ impl TryFrom for Bpe { let merge_map = parse_merges(&raw.vocab, &raw.merges)?; let mut bpe = Self::new(&raw.vocab, merge_map)?; bpe.ignore_merges = raw.ignore_merges; + bpe.byte_fallback = raw.byte_fallback; Ok(bpe) } } @@ -771,6 +815,14 @@ impl Bpe { } } + let mut byte_fallback_token_ids = [INVALID_TOKEN; 256]; + for byte_val in 0u16..256 { + let token = format!("<0x{byte_val:02X}>"); + if let Some(&id) = vocab.get(token.as_str()) { + byte_fallback_token_ids[byte_val as usize] = id; + } + } + // Pre-compute initial byte-pair merges (256×256 table). let mut byte_pair_initial = vec![(u32::MAX, 0u32); 65536]; for b1 in 0u16..256 { @@ -789,9 +841,20 @@ impl Bpe { } } + let mut single_char_token = [INVALID_TOKEN; 128]; + for (byte, slot) in single_char_token.iter_mut().enumerate() { + let ch = byte as u8 as char; + let mut buf = [0u8; 1]; + if let Some(&id) = vocab.get(ch.encode_utf8(&mut buf) as &str) { + *slot = id; + } + } + let vocab_size = id_to_token.len(); let merge_adj = MergeAdjacency::from_parsed(&merge_map, vocab_size); + let bigram_bridge_table = build_bigram_bridge_table(&id_to_token); + Ok(Self { id: BPE_ID_COUNTER.fetch_add(1, Ordering::Relaxed), daac, @@ -804,10 +867,14 @@ impl Bpe { id_to_token, token_to_id: vocab.clone(), byte_to_initial_token, + byte_fallback_token_ids, + single_char_token, ranked_merge_map, byte_pair_initial, merge_adj, ignore_merges: false, + byte_fallback: false, + bigram_bridge_table, }) } @@ -928,20 +995,44 @@ impl Bpe { for ch in input.chars() { let mut buf = [0u8; 4]; let s = ch.encode_utf8(&mut buf); - let id = self - .token_to_id - .get(s) - .copied() - .ok_or_else(|| format!("character {ch:?} not in vocabulary"))?; - scratch.symbols.push(MergeSymbol { - c: id, - prev: if n == 0 { -1 } else { (n - 1) as i32 }, - next: -1, - }); - if n > 0 { - scratch.symbols[n - 1].next = n as i32; + let found = if ch.is_ascii() { + let id = self.single_char_token[ch as usize]; + (id != INVALID_TOKEN).then_some(id) + } else { + self.token_to_id.get(s).copied() + }; + if let Some(id) = found { + scratch.symbols.push(MergeSymbol { + c: id, + prev: if n == 0 { -1 } else { (n - 1) as i32 }, + next: -1, + }); + if n > 0 { + scratch.symbols[n - 1].next = n as i32; + } + n += 1; + continue; + } + + if !self.byte_fallback { + return Err(format!("character {ch:?} not in vocabulary")); + } + + for &byte in s.as_bytes() { + let id = self.byte_fallback_token_ids[byte as usize]; + if id == INVALID_TOKEN { + return Err(format!("byte fallback token <0x{byte:02X}> not in vocabulary")); + } + scratch.symbols.push(MergeSymbol { + c: id, + prev: if n == 0 { -1 } else { (n - 1) as i32 }, + next: -1, + }); + if n > 0 { + scratch.symbols[n - 1].next = n as i32; + } + n += 1; } - n += 1; } if n == 1 { @@ -1025,7 +1116,7 @@ impl Bpe { })); } - #[inline(always)] +#[inline(always)] fn run_merge_loop(&self, scratch: &mut MergeScratch, out: &mut Vec) { let symbols = &mut scratch.symbols; let heap = &mut scratch.heap; @@ -1236,10 +1327,14 @@ impl Clone for Bpe { id_to_token: self.id_to_token.clone(), token_to_id: self.token_to_id.clone(), byte_to_initial_token: self.byte_to_initial_token, + byte_fallback_token_ids: self.byte_fallback_token_ids, + single_char_token: self.single_char_token, ranked_merge_map: self.ranked_merge_map.clone(), byte_pair_initial: self.byte_pair_initial.clone(), merge_adj: self.merge_adj.clone(), ignore_merges: self.ignore_merges, + byte_fallback: self.byte_fallback, + bigram_bridge_table: self.bigram_bridge_table.clone(), } } } @@ -1261,6 +1356,7 @@ impl PartialEq for Bpe { && self.next_prefix_map == other.next_prefix_map && self.token_lens == other.token_lens && self.ignore_merges == other.ignore_merges + && self.byte_fallback == other.byte_fallback } } diff --git a/src/normalizers.rs b/src/normalizers.rs index 5a1c8b6..d2a0e6a 100644 --- a/src/normalizers.rs +++ b/src/normalizers.rs @@ -1,13 +1,21 @@ mod nfc; +mod replace; use std::borrow::Cow; pub use self::nfc::Nfc; +pub use self::replace::Replace; use crate::json_structs::{NormalizerConfig, NormalizerKind}; /// Errors from constructing a normalizer. #[derive(Debug, thiserror::Error)] pub enum Error { + #[error("invalid config value: {0}")] + Json(#[from] serde_json::Error), + + #[error("regex error: {0}")] + Regex(#[from] fancy_regex::Error), + #[error("unsupported normalizer type: {0}")] Unsupported(String), } @@ -16,6 +24,7 @@ pub enum Error { #[derive(Debug)] pub enum Normalizer { Nfc(Nfc), + Replace(Replace), Sequence(Vec), } @@ -24,6 +33,9 @@ impl Normalizer { pub fn from_config(config: NormalizerConfig) -> Result { match config { NormalizerConfig::Nfc => Ok(Self::Nfc(Nfc)), + NormalizerConfig::Replace { pattern, content } => { + Ok(Self::Replace(Replace::from_config(pattern, content)?)) + } NormalizerConfig::Sequence { normalizers } => { let steps = normalizers .into_iter() @@ -46,6 +58,7 @@ impl Normalizer { pub fn normalize<'a>(&self, input: &'a str) -> Cow<'a, str> { match self { Self::Nfc(nfc) => nfc.normalize(input), + Self::Replace(replace) => replace.normalize(input), Self::Sequence(steps) => { let mut current = Cow::Borrowed(input); for step in steps { diff --git a/src/normalizers/replace.rs b/src/normalizers/replace.rs new file mode 100644 index 0000000..9809a54 --- /dev/null +++ b/src/normalizers/replace.rs @@ -0,0 +1,138 @@ +use std::borrow::Cow; + +use fancy_regex::{NoExpand, Regex}; +use serde::de::Error as _; +use serde_json::Value; + +use super::Error; + +#[derive(Clone, Debug)] +enum Pattern { + Literal(String), + Regex(Regex), +} + +impl Pattern { + fn from_json(value: Value) -> Result { + if let Some(s) = value.as_str() { + return Ok(Self::Literal(s.to_string())); + } + + let obj = value.as_object().ok_or_else(|| { + serde_json::Error::custom("Replace.pattern must be a string or an object") + })?; + + if let Some(literal) = obj.get("String").and_then(Value::as_str) { + return Ok(Self::Literal(literal.to_string())); + } + + if let Some(regex) = obj.get("Regex").and_then(Value::as_str) { + return Ok(Self::Regex(Regex::new(regex)?)); + } + + Err(serde_json::Error::custom("Replace.pattern object must contain String or Regex").into()) + } +} + +/// Replace normalizer. +/// +/// Replaces each occurrence of a literal string or regex pattern with +/// `content`. +#[derive(Clone, Debug)] +pub struct Replace { + pattern: Pattern, + content: String, +} + +impl Replace { + pub fn from_config(pattern: Value, content: String) -> Result { + Ok(Self { + pattern: Pattern::from_json(pattern)?, + content, + }) + } + + pub fn normalize<'a>(&self, input: &'a str) -> Cow<'a, str> { + match &self.pattern { + Pattern::Literal(needle) => replace_literal(input, needle, &self.content), + Pattern::Regex(re) => re.replace_all(input, NoExpand(&self.content)), + } + } +} + +fn replace_literal<'a>(input: &'a str, needle: &str, replacement: &str) -> Cow<'a, str> { + if needle.is_empty() { + let mut output = String::new(); + output.push_str(replacement); + for ch in input.chars() { + output.push(ch); + output.push_str(replacement); + } + return Cow::Owned(output); + } + + let Some(first_match) = input.find(needle) else { + return Cow::Borrowed(input); + }; + + let mut output = String::with_capacity(input.len()); + output.push_str(&input[..first_match]); + output.push_str(replacement); + + let mut tail = &input[first_match + needle.len()..]; + while let Some(next_match) = tail.find(needle) { + output.push_str(&tail[..next_match]); + output.push_str(replacement); + tail = &tail[next_match + needle.len()..]; + } + + output.push_str(tail); + Cow::Owned(output) +} + +#[cfg(test)] +mod tests { + use std::borrow::Cow; + + use serde_json::json; + + use super::*; + + #[test] + fn literal_replace() { + let repl = Replace::from_config(json!({"String": " "}), "▁".to_string()).unwrap(); + assert_eq!(repl.normalize("a b c"), "a▁b▁c"); + } + + #[test] + fn literal_no_change_borrowed() { + let repl = Replace::from_config(json!({"String": "x"}), "y".to_string()).unwrap(); + let out = repl.normalize("abc"); + assert_eq!(out, "abc"); + assert!(matches!(out, Cow::Borrowed(_))); + } + + #[test] + fn regex_replace() { + let repl = Replace::from_config(json!({"Regex": "\\s+"}), " ".to_string()).unwrap(); + assert_eq!(repl.normalize("hello world"), "hello world"); + } + + #[test] + fn accepts_plain_string_pattern() { + let repl = Replace::from_config(json!("."), " ".to_string()).unwrap(); + assert_eq!(repl.normalize("hello.world"), "hello world"); + } + + #[test] + fn literal_empty_pattern_is_not_noop() { + let repl = Replace::from_config(json!(""), "-".to_string()).unwrap(); + assert_eq!(repl.normalize("ab"), "-a-b-"); + } + + #[test] + fn regex_replacement_content_is_literal() { + let repl = Replace::from_config(json!({"Regex": "([a-z]+)"}), "$1".to_string()).unwrap(); + assert_eq!(repl.normalize("abc"), "$1"); + } +}