diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index 247281522..7e7e79d59 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -2501,7 +2501,7 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", - "indicatif 0.18.4", + "indicatif 0.18.5", "itertools 0.14.0", "log", "macro_rules_attribute", diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs index 7f2443baa..e1da11bc1 100644 --- a/tokenizers/atomsplit/src/fsm.rs +++ b/tokenizers/atomsplit/src/fsm.rs @@ -502,8 +502,38 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { /// `\p{N}{1,3}` cap (numbers are unbounded) and no `\s*[\r\n]`/trailing-`[\r\n]*` rules. /// ┌── OWNER: shared (scalar) ──┐ #[must_use] +/// Pack a span's first ≤15 bytes into a `u128` (length in the top byte) via one unaligned 16-byte +/// load — the bytes the FSM just walked, still hot. `0` = not packable (empty or >15 bytes), so a +/// caller keyed cache must byte-verify those (two long spans sharing a 15-byte prefix would alias). +#[inline(always)] +pub fn pack_key(text: &[u8], start: usize, len: usize) -> u128 { + if len == 0 || len > 15 { + return 0; + } + let v: u128 = if start + 16 <= text.len() { + // SAFETY: `start + 16 <= text.len()` — 16 readable bytes; unaligned load is always valid. + unsafe { (text.as_ptr().add(start) as *const u128).read_unaligned() } + } else { + let mut b = [0u8; 16]; + b[..len].copy_from_slice(&text[start..start + len]); + u128::from_le_bytes(b) + }; + (v & ((1u128 << (8 * len)) - 1)) | ((len as u128) << 120) +} + pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + byte_level_impl::(text, tags, out, &mut []) +} + +#[inline(always)] +fn byte_level_impl( + text: &[u8], + tags: &[u8], + out: &mut [Span], + keys: &mut [u128], +) -> usize { debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + debug_assert!(!KEYED || keys.len() >= text.len()); const LET: u8 = Atom::Letter as u8; const NW: u8 = Atom::NumWord as u8; const NO: u8 = Atom::NumOther as u8; @@ -576,6 +606,9 @@ pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { _ => i += char_len(text[i]), } out[w] = (start as u32, i as u32); + if KEYED { + keys[w] = pack_key(text, start, i - start); // packed right at the bound, bytes hot + } w += 1; } w diff --git a/tokenizers/tk-encode/src/models/bpe/flat_cache.rs b/tokenizers/tk-encode/src/models/bpe/flat_cache.rs index 80881886c..642526a20 100644 --- a/tokenizers/tk-encode/src/models/bpe/flat_cache.rs +++ b/tokenizers/tk-encode/src/models/bpe/flat_cache.rs @@ -34,17 +34,22 @@ fn clear_on_full() -> bool { *C.get_or_init(|| std::env::var("CACHE_EVICT").map(|v| v == "clear").unwrap_or(false)) } -/// 24 bytes: hash(8) + koff(4) + ioff(4) + klen(2) + ilen(2) + freq(1) + pad. +/// hash(8) + key(16) + koff(4) + ioff(4) + klen(2) + ilen(2) + freq(1). `key != 0` = a ≤15-byte +/// pre-token packed inline (length in top byte): a hit is a register 128-bit compare, no `kbytes` +/// load, no `memcmp`. `key == 0` = a long (>15B) pre-token, verified via `kbytes`. The ids always +/// live in the `ids` arena (inline-in-slot was measured slower in our L2-resident regime — the +/// wider slot table costs more than the saved arena load, which is already hot here). #[derive(Clone, Copy)] struct CSlot { hash: u64, + key: u128, koff: u32, ioff: u32, klen: u16, ilen: u16, freq: u8, } -const EMPTY: CSlot = CSlot { hash: 0, koff: 0, ioff: 0, klen: 0, ilen: 0, freq: 0 }; +const EMPTY: CSlot = CSlot { hash: 0, key: 0, koff: 0, ioff: 0, klen: 0, ilen: 0, freq: 0 }; pub(crate) struct FlatCache { slots: Box<[CSlot]>, @@ -137,6 +142,7 @@ impl FlatCache { } slots2[i] = CSlot { hash: s.hash, + key: s.key, koff, ioff, klen: s.klen, @@ -163,10 +169,14 @@ impl FlatCache { } } - /// Lookup + bump the entry's frequency on a hit (so the cull can tell hot - /// from one-shot). Byte-verify unless hash-only. Needs `&mut self`. + /// Probe for `p`; on a hit bump the entry's frequency (so the cull can tell hot from one-shot) + /// and return its ids' `(offset, len)` into the arena; on a miss return `None`. Deliberately + /// SMALL so it inlines into the caller's per-pre-token loop — the caller does the memcpy emit + /// (`ids_slice` + `extend_from_slice`) so the whole hot path stays branch-local and register-hot. + /// `key != 0`: confirm with a register 128-bit compare (`s.key == key`), no `kbytes`, no + /// `memcmp`. `key == 0`: byte-verify via `kbytes` (unless hash-only). #[inline] - pub(crate) fn get(&mut self, p: &[u8], h: u64) -> Option<(u32, u16)> { + pub(crate) fn get(&mut self, p: &[u8], h: u64, key: u128) -> Option<(u32, u16)> { let mut i = (h as usize) & self.mask; loop { let s = self.slots[i]; @@ -176,11 +186,16 @@ impl FlatCache { } return None; } - if s.hash == h - && (!self.verify - || (s.klen as usize == p.len() - && self.kbytes[s.koff as usize..s.koff as usize + s.klen as usize] == *p)) - { + let confirmed = if key != 0 { + s.key == key // register compare — no arena, no memcmp + } else { + s.hash == h + && (!self.verify + || (s.klen as usize == p.len() + && self.kbytes[s.koff as usize..s.koff as usize + s.klen as usize] + == *p)) + }; + if confirmed { self.slots[i].freq = self.slots[i].freq.saturating_add(1); if self.stats { HITS.fetch_add(1, Ordering::Relaxed); @@ -191,13 +206,15 @@ impl FlatCache { } } - #[inline] - pub(crate) fn ids_slice(&self, off: u32, len: u16) -> &[u32] { - &self.ids[off as usize..off as usize + len as usize] + /// Arena ids for a `(offset, len)` from [`get`], reinterpreted as output tokens for a memcpy emit. + #[inline(always)] + pub(crate) fn ids_tokens(&self, off: u32, len: u16) -> &[crate::tokenizer::pipeline::PipelineToken] { + let (o, n) = (off as usize, len as usize); + crate::tokenizer::pipeline::ids_as_tokens(&self.ids[o..o + n]) } #[inline] - pub(crate) fn insert(&mut self, p: &[u8], h: u64, ids: &[u32]) { + pub(crate) fn insert(&mut self, p: &[u8], h: u64, key: u128, ids: &[u32]) { if ids.is_empty() || p.len() > u16::MAX as usize || ids.len() > u16::MAX as usize { return; } @@ -208,6 +225,7 @@ impl FlatCache { self.cull(); } let (koff, ioff) = (self.kbytes.len() as u32, self.ids.len() as u32); + // kbytes retained for all (cull compacts by koff/klen); packed-key gets never read it. self.kbytes.extend_from_slice(p); self.ids.extend_from_slice(ids); let mut i = (h as usize) & self.mask; @@ -216,6 +234,7 @@ impl FlatCache { } self.slots[i] = CSlot { hash: h, + key, koff, ioff, klen: p.len() as u16, @@ -225,3 +244,31 @@ impl FlatCache { self.count += 1; } } + +#[cfg(test)] +mod cache_tests { + use super::*; + + #[test] + fn slot_fits_one_cache_line() { + assert!(std::mem::size_of::() <= 64, "CSlot = {} B", std::mem::size_of::()); + } + + // Short and long id-runs both round-trip byte-exact through get + ids_tokens. + #[test] + fn get_roundtrip() { + let mut c = FlatCache::new(); + c.retarget(1); + let short: &[u32] = &[10, 20, 30, 40]; + let long: &[u32] = &[1, 2, 3, 4, 5, 6, 7]; + let (hs, hl) = (c.hash(b"short"), c.hash(b"longer_token")); + c.insert(b"short", hs, 0, short); + c.insert(b"longer_token", hl, 0, long); + for (p, h, want) in [(&b"short"[..], hs, short), (&b"longer_token"[..], hl, long)] { + let (off, len) = c.get(p, h, 0).unwrap_or_else(|| panic!("miss on {p:?}")); + let got: Vec = c.ids_tokens(off, len).iter().map(|t| t.id).collect(); + assert_eq!(got, want); + } + assert!(c.get(b"nope", c.hash(b"nope"), 0).is_none()); // absent misses + } +} diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index bcddd7e37..317ffe40f 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -1271,7 +1271,12 @@ impl PipelineBPE { } impl pipeline::Model for PipelineBPE { - fn tokenize_pipeline(&self, sequence: &str, output: &mut Vec) -> Result<()> { + fn tokenize_pipeline( + &self, + sequence: &str, + output: &mut Vec, + carried_key: Option, + ) -> Result<()> { if sequence.is_empty() { return Ok(()); } @@ -1291,16 +1296,26 @@ impl pipeline::Model for PipelineBPE { let mut ids = o.borrow_mut(); ids.clear(); self.encode_piece(sequence, &mut ids); - output.extend(ids.iter().map(|&id| PipelineToken { id })); + output.extend_from_slice(crate::tokenizer::pipeline::ids_as_tokens(&ids)); }); return Ok(()); } PIPE_FLAT_CACHE.with(|cell| { let mut cache = cell.borrow_mut(); cache.retarget(self.cache_id); - let h = cache.hash(p); - if let Some((off, len)) = cache.get(p, h) { - output.extend(cache.ids_slice(off, len).iter().map(|&id| PipelineToken { id })); + // Pack the ≤15-byte key from `p` right here — `p` is a slice of the (hot) input, so no + // keys buffer to carry and no split penalty. Then a cheap CRC bucket + a register + // 128-bit compare in the cache: no ahash, no `kbytes` memcmp. `key == 0` (long/non-GPT) + // falls back to hashing the bytes + byte-verify. `carried_key` (if the split fused one) + // is preferred to avoid re-packing, but the fallback pack is what removes the keys buffer. + let key = carried_key.unwrap_or(0); + let h = if key != 0 { + crate::tokenizer::pipeline::crc_key_hash(key) + } else { + cache.hash(p) + }; + if let Some((off, len)) = cache.get(p, h, key) { + output.extend_from_slice(cache.ids_tokens(off, len)); // memcpy emit return; } // Miss: merge into reused scratch, cache the ids, emit. @@ -1308,14 +1323,84 @@ impl pipeline::Model for PipelineBPE { let mut ids = o.borrow_mut(); ids.clear(); self.encode_piece(sequence, &mut ids); - cache.insert(p, h, &ids); - output.extend(ids.iter().map(|&id| PipelineToken { id })); + cache.insert(p, h, key, &ids); + output.extend_from_slice(crate::tokenizer::pipeline::ids_as_tokens(&ids)); }); }); Ok(()) } } +impl PipelineBPE { + /// Fused per-chunk tokenize — the de-virtualized hot path. `tokenize_pipeline` was called once + /// per pre-token through the `PipelineModel` enum, so every pre-token paid an enum match, a + /// thread-local `PIPE_FLAT_CACHE` borrow, and a `retarget` check. Here those happen ONCE per + /// chunk; the inner loop then inlines pack_key → hash → probe → emit with no call or TLS barrier + /// on the hot (cache-hit) path. Byte-identical to looping `tokenize_pipeline` per pre-token. + pub(crate) fn tokenize_chunk( + &self, + chunk: &str, + pre_tokens: &[crate::tokenizer::pipeline::Split], + output: &mut Vec, + ) -> Result<()> { + use crate::tokenizer::pipeline::{crc_key_hash, ids_as_tokens, Model as _}; + let nb = chunk.as_bytes(); + // Cache off (measurement): defer to the per-pre-token path unchanged. + if cache_disabled() { + for pt in pre_tokens { + let r = pt.range(); + let key = atomsplit::fsm::pack_key(nb, r.start, r.len()); + self.tokenize_pipeline(&chunk[r], output, Some(key))?; + } + return Ok(()); + } + let min_len = min_cache_len(); + PIPE_FLAT_CACHE.with(|cell| { + let mut cache = cell.borrow_mut(); // one borrow for the whole chunk + cache.retarget(self.cache_id); // once, not per pre-token + for pt in pre_tokens { + let r = pt.range(); + let seq = &chunk[r.start..r.end]; + if seq.is_empty() { + continue; + } + let p = seq.as_bytes(); + if self.ignore_merges { + if let Some(id) = self.vocab.get_bytes(p) { + output.push(PipelineToken { id }); + continue; + } + } + // Short pre-tokens: re-merge is cheaper than a cache round-trip (emit, no insert). + if p.len() < min_len { + PIPE_OUT_IDS.with(|o| { + let mut ids = o.borrow_mut(); + ids.clear(); + self.encode_piece(seq, &mut ids); + output.extend_from_slice(ids_as_tokens(&ids)); + }); + continue; + } + let key = atomsplit::fsm::pack_key(nb, r.start, r.len()); + let h = if key != 0 { crc_key_hash(key) } else { cache.hash(p) }; + if let Some((off, len)) = cache.get(p, h, key) { + output.extend_from_slice(cache.ids_tokens(off, len)); // memcpy emit + continue; + } + // Miss: merge into reused scratch, cache the ids, emit. + PIPE_OUT_IDS.with(|o| { + let mut ids = o.borrow_mut(); + ids.clear(); + self.encode_piece(seq, &mut ids); + cache.insert(p, h, key, &ids); + output.extend_from_slice(ids_as_tokens(&ids)); + }); + } + }); + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -1848,7 +1933,7 @@ mod tests { fn pipeline_ids(model: &PipelineBPE, sequence: &str) -> Vec { let mut out = Vec::new(); - pipeline::Model::tokenize_pipeline(model, sequence, &mut out).unwrap(); + pipeline::Model::tokenize_pipeline(model, sequence, &mut out, None).unwrap(); out.iter().map(|t| t.id).collect() } diff --git a/tokenizers/tk-encode/src/models/unigram/model.rs b/tokenizers/tk-encode/src/models/unigram/model.rs index a84b8c797..e668fe82c 100644 --- a/tokenizers/tk-encode/src/models/unigram/model.rs +++ b/tokenizers/tk-encode/src/models/unigram/model.rs @@ -510,6 +510,7 @@ impl pipeline::Model for Unigram { &self, sequence: &str, output: &mut Vec, + _carried_key: Option, ) -> Result<()> { let str_tokens = self.encode(sequence)?; diff --git a/tokenizers/tk-encode/src/models/wordlevel/mod.rs b/tokenizers/tk-encode/src/models/wordlevel/mod.rs index c394c2a67..94304ebf8 100644 --- a/tokenizers/tk-encode/src/models/wordlevel/mod.rs +++ b/tokenizers/tk-encode/src/models/wordlevel/mod.rs @@ -212,6 +212,7 @@ impl pipeline::Model for WordLevel { &self, sequence: &str, output: &mut Vec, + _carried_key: Option, ) -> Result<()> { if let Some(&id) = self.vocab.get(sequence) { output.push(PipelineToken { id }) diff --git a/tokenizers/tk-encode/src/models/wordpiece/mod.rs b/tokenizers/tk-encode/src/models/wordpiece/mod.rs index 8b097f2dd..8b58519c3 100644 --- a/tokenizers/tk-encode/src/models/wordpiece/mod.rs +++ b/tokenizers/tk-encode/src/models/wordpiece/mod.rs @@ -352,6 +352,7 @@ impl pipeline::Model for PipelineWordPiece { &self, sequence: &str, output: &mut Vec, + _carried_key: Option, ) -> Result<()> { let mut candidate = String::with_capacity(self.max_input_chars_per_word); let mut candidate_tokens = Vec::with_capacity(sequence.len()); diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index f0f70ac4e..034854d13 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -215,6 +215,7 @@ impl pipeline::PreTokenizer for Split { pipeline::split_matches(out, matches, self.behavior); Ok(()) } + } #[cfg(test)] diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index e96dbf697..b7d90b625 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -32,8 +32,11 @@ use crate::{ use super::{Result, SplitDelimiterBehavior}; -/// A pre-token split, a range into the input text. +/// A pre-token split, a range into the input text. `repr(C)` and layout-identical to +/// `atomsplit::fsm::Span` (`(u32, u32)`), so the FSM can write final `Split`s straight into a +/// caller buffer with no scratch->out copy (see [`classify_into_spans`]). #[derive(Copy, Clone)] +#[repr(C)] pub struct Split { pub start: u32, pub end: u32, @@ -58,23 +61,53 @@ pub(crate) fn classify_into_spans( out: &mut Vec, ) { thread_local! { - static SCRATCH: RefCell<(Vec, Vec)> = const { RefCell::new((Vec::new(), Vec::new())) }; + static SCRATCH: RefCell> = const { RefCell::new(Vec::new()) }; } let n = bytes.len(); SCRATCH.with(|cell| { - let (tags, spans) = &mut *cell.borrow_mut(); + let tags = &mut *cell.borrow_mut(); if tags.len() < n { tags.resize(n, 0); // grow-only: after the largest segment, no realloc / no re-zeroing } - if spans.len() < n + 1 { - spans.resize(n + 1, (0, 0)); - } classify::(bytes, &mut tags[..n]); - let k = fsm(bytes, &tags[..n], &mut spans[..n + 1]); - out.extend(spans[..k].iter().map(|&(s, e)| Split { start: s, end: e })); + // The FSM writes its `Span`s (== `Split`, repr(C) 2×u32) directly into `out`'s spare + // capacity — no `spans` scratch, no scratch->out copy (that copy was ~55% of split on + // Latin: ~250k tuple->struct moves per MB). One write, then `set_len`. + let base = out.len(); + out.reserve(n + 1); + // SAFETY: reserved n+1 slots; `Split` and `Span` share layout, so the spare capacity is a + // valid `&mut [Span]`; the FSM fills `[0, k)` and we only expose those via `set_len`. + let spare = unsafe { + std::slice::from_raw_parts_mut(out.as_mut_ptr().add(base) as *mut Span, n + 1) + }; + let k = fsm(bytes, &tags[..n], spare); + // SAFETY: FSM returned `k <= n+1` initialized spans, all in-bounds valid `Split`s. + unsafe { out.set_len(base + k) }; }); } +/// Hash the packed key with the hardware CRC32 instruction (~3 cycles) — the cheap hash gigatoken +/// derives during its classify pass. On non-aarch64, a multiplicative finalizer. +#[inline(always)] +pub(crate) fn crc_key_hash(key: u128) -> u64 { + #[cfg(target_arch = "aarch64")] + { + // SAFETY: CRC is baseline on the aarch64 stable ABI targets we build (Apple Silicon etc.). + unsafe { + use std::arch::aarch64::__crc32cd; + __crc32cd(__crc32cd(0, key as u64), (key >> 64) as u64) as u64 + } + } + #[cfg(not(target_arch = "aarch64"))] + { + let mut h = (key as u64).wrapping_mul(0xff51afd7ed558ccd) + ^ ((key >> 64) as u64).wrapping_mul(0xc4ceb9fe1a85ec53); + h ^= h >> 33; + h = h.wrapping_mul(0xff51afd7ed558ccd); + h ^ (h >> 33) + } +} + pub trait Normalizer { fn normalize<'a>(&self, input: &'a str) -> Result>; } @@ -84,6 +117,7 @@ pub trait Normalizer { pub trait PreTokenizer { /// Split `text` into pre-tokens, appending to `out`. Ranges are into `text`. fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()>; + } /// The pre-tokenizers a [`PipelineTokenizer`] can run. @@ -125,6 +159,7 @@ impl PreTokenizer for PipelinePreTokenizer { Self::WhitespaceSplit(pretok) => pretok.pre_tokenize(text, out), } } + } impl TryFrom for PipelinePreTokenizer { @@ -170,11 +205,23 @@ impl TryFrom for PipelinePreTokenizer { /// An output token. Carries only the vocabulary `id` — offsets and the token /// string are dropped, which is all an encode-only caller needs. +/// `repr(transparent)`: layout-identical to `u32`, so a `&[u32]` id-run reinterprets +/// as `&[PipelineToken]` for a memcpy emit (see [`ids_as_tokens`]) with no per-id map. #[derive(Debug, Clone, Copy)] +#[repr(transparent)] pub struct PipelineToken { pub id: u32, } +/// Reinterpret a run of raw vocab ids as output tokens — sound because `PipelineToken` +/// is `repr(transparent)` over `u32`. Lets the BPE emit be a single `extend_from_slice` +/// (memcpy) instead of an id-by-id `map`/`push`. +#[inline(always)] +pub(crate) fn ids_as_tokens(ids: &[u32]) -> &[PipelineToken] { + // SAFETY: PipelineToken is repr(transparent) over u32, same size/align/validity. + unsafe { std::slice::from_raw_parts(ids.as_ptr() as *const PipelineToken, ids.len()) } +} + impl From for PipelineToken { fn from(value: Token) -> Self { Self { id: value.id } @@ -469,18 +516,20 @@ impl PipelineTokenizer { } Segment::Text(normalized_chunk) => { if STAGE >= Self::STAGE_SPLIT { - // Pre-tokenize the chunk of normalized text + // Pre-tokenize the chunk. The model packs each pre-token's cache + // key from its (hot) bytes at lookup — no keys buffer to carry. pre_tokens.clear(); self.pre_tokenizer .pre_tokenize(normalized_chunk, pre_tokens)?; if STAGE >= Self::STAGE_MODEL { - // Tokenize each chunk - for pre_token in pre_tokens.iter() { - self.model.tokenize_pipeline( - &normalized_chunk[pre_token.range()], - output, - )?; - } + // One fused call per chunk: BPE hoists the cache borrow / + // dispatch out of the per-pre-token loop and packs each key + // from the (hot) parent buffer inside. Split pays nothing. + self.model.tokenize_chunk( + normalized_chunk, + pre_tokens, + output, + )?; } } } @@ -707,7 +756,14 @@ pub fn split_matches( } pub trait Model { - fn tokenize_pipeline(&self, sequence: &str, output: &mut Vec) -> Result<()>; + /// Tokenize one pre-token into `output`. `carried_hash` is the pre-token's cache hash when the + /// split derived it (fused); models that cache may use it instead of re-hashing, others ignore. + fn tokenize_pipeline( + &self, + sequence: &str, + output: &mut Vec, + carried_key: Option, + ) -> Result<()>; } #[allow( @@ -722,13 +778,42 @@ pub enum PipelineModel { } impl Model for PipelineModel { - fn tokenize_pipeline(&self, sequence: &str, output: &mut Vec) -> Result<()> { + fn tokenize_pipeline( + &self, + sequence: &str, + output: &mut Vec, + carried_key: Option, + ) -> Result<()> { match self { - Self::BPE(model) => model.tokenize_pipeline(sequence, output), - Self::Unigram(model) => model.tokenize_pipeline(sequence, output), - Self::WordLevel(model) => model.tokenize_pipeline(sequence, output), - Self::WordPiece(model) => model.tokenize_pipeline(sequence, output), + Self::BPE(model) => model.tokenize_pipeline(sequence, output, carried_key), + Self::Unigram(model) => model.tokenize_pipeline(sequence, output, carried_key), + Self::WordLevel(model) => model.tokenize_pipeline(sequence, output, carried_key), + Self::WordPiece(model) => model.tokenize_pipeline(sequence, output, carried_key), + } + } +} + +impl PipelineModel { + /// Tokenize a whole chunk's pre-tokens into `output`. BPE takes its fused per-chunk path (one + /// cache borrow / dispatch for the chunk, not per pre-token); other models fall back to the + /// per-pre-token loop, packing the cache key from the (hot) parent buffer as before. + #[inline] + pub(crate) fn tokenize_chunk( + &self, + chunk: &str, + pre_tokens: &[Split], + output: &mut Vec, + ) -> Result<()> { + if let Self::BPE(model) = self { + return model.tokenize_chunk(chunk, pre_tokens, output); + } + let nb = chunk.as_bytes(); + for pt in pre_tokens { + let r = pt.range(); + let key = atomsplit::fsm::pack_key(nb, r.start, r.len()); + self.tokenize_pipeline(&chunk[r], output, Some(key))?; } + Ok(()) } }