diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 675cfc18f..2bc0dd74c 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -12,7 +12,7 @@ use crate::pipeline::{self, PipelineToken}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; use crate::utils::word_cache::{Lookup, WordCache}; -use crate::vocab::bucket_vocab_store::BucketVocabStore; +use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash}; const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; @@ -248,16 +248,16 @@ impl PipelineBPE { #[inline(always)] fn fold_id(&self, sequence: &str) -> Option { let bytes = sequence.as_bytes(); - self.fold_id_hashed(bytes, self.vocab.hash_word(bytes)) + let (key, hash) = key_and_hash(bytes); + self.fold_id_keyed(bytes, key, hash) } - /// [`Self::fold_id`] for a caller that already hashed the word with - /// [`BucketVocabStore::hash_word`]. + /// [`Self::fold_id`] for a caller that already ran [`key_and_hash`] on the word. #[inline(always)] - fn fold_id_hashed(&self, bytes: &[u8], hash: u64) -> Option { + fn fold_id_keyed(&self, bytes: &[u8], key: u64, hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries // carry it was settled at load -- see `from_bpe`. - let (id, foldable) = self.vocab.get_bytes_foldable_hashed(bytes, hash)?; + let (id, foldable) = self.vocab.get_bytes_foldable_keyed(bytes, key, hash)?; foldable.then_some(id) } @@ -326,8 +326,8 @@ impl pipeline::Model for PipelineBPE { // way -- the vocabulary compares the entry's bytes, the cache compares its key -- so this // shares the hash and nothing else. let bytes = sequence.as_bytes(); - let hash = self.vocab.hash_word(bytes); - if let Some(id) = self.fold_id_hashed(bytes, hash) { + let (key, hash) = key_and_hash(bytes); + if let Some(id) = self.fold_id_keyed(bytes, key, hash) { output.push(PipelineToken { id }); return Ok(()); } diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index fd21898d9..08e077874 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -42,7 +42,9 @@ const LEN_TAG: [u64; INLINE_KEY_BYTES + 1] = { let mut t = [0u64; INLINE_KEY_BYTES + 1]; let mut len = 0; while len <= INLINE_KEY_BYTES { - t[len] = (len as u64) << 56; + // `len + 1`, not `len`, so the top byte is never zero and no real key can be 0. A phantom + // slot's key is 0, which is what lets one compare reject it without a second load. + t[len] = (len as u64 + 1) << 56; len += 1; } t @@ -73,9 +75,20 @@ fn mix(z: u64) -> u64 { /// word under one value and look it up under another. #[inline] pub fn word_hash(word: &[u8]) -> u64 { + key_and_hash(word).1 +} + +/// The key a slot is verified by, and the `u64` the MPHF is indexed by. +/// +/// Up to [`INLINE_KEY_BYTES`] bytes the key *is* the word, so comparing it to a slot's key is +/// proof of identity and the byte slab never has to be read. Longer words key by aHash, which is +/// not proof, so those still confirm against the slab. +#[inline] +pub fn key_and_hash(word: &[u8]) -> (u64, u64) { let len = word.len(); if len > INLINE_KEY_BYTES { - return KEY_HASHER.hash_one(word); + let hash = KEY_HASHER.hash_one(word); + return (hash, hash); } // Reading past the word is not allowed, so read a head and a tail that overlap and stitch them: // still register-only, no `memcpy`. @@ -91,7 +104,8 @@ pub fn word_hash(word: &[u8]) -> u64 { } else { 0 }; - mix((raw & KEY_MASK[len]) | LEN_TAG[len]) + let key = (raw & KEY_MASK[len]) | LEN_TAG[len]; + (key, mix(key)) } /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be @@ -105,14 +119,27 @@ const FOLD_BIT: u32 = 1 << 31; /// The id half. 2^31 ids is far past any vocabulary. const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; -#[derive(Clone, Copy, Debug)] +/// Everything a probe reads, and nothing it does not. +/// +/// The key used to live in the byte slab, so verifying a slot meant a third dependent load after +/// the MPHF pilot and this entry. Holding it here makes the probe two loads. `(start, len)` moved +/// to [`Span`]: only the reverse lookup and enumeration want them, and keeping them here made +/// every probe drag six dead bytes through cache. +#[derive(Clone, Copy, Debug, Default)] struct Entry { - start: u32, - len: u16, + /// 0 for a phantom slot, which no real key can be -- see [`LEN_TAG`]. + key: u64, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } +/// `slot -> (offset into `bytes`, length)`. Off the probe path on purpose. +#[derive(Clone, Copy, Debug, Default)] +struct Span { + start: u32, + len: u16, +} + /// The BucketVocabStore optimizes for space and speed. We don't use a HashMap to prevent duplicating the /// keys. Instead, we just use an `id_to_slot` and `entries` table. When you query bytes, you hash /// on the fly and get an `index` into the `entries` table. When you query an `id`, you fetch in @@ -136,8 +163,11 @@ pub struct BucketVocabStore { mphf: Mphf, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, - /// `entries[slot]` -> (offset into `bytes`, length, id). Ordered by MPHF slot. + /// `entries[slot]` -> (key, id). Ordered by MPHF slot. entries: Box<[Entry]>, + /// `spans[slot]` -> where the token's bytes live. Parallel to `entries`, read only by the + /// reverse lookup and by enumeration. + spans: Box<[Span]>, /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, /// Number of real tokens. Cached at build so `len()` is O(1): `entries` is sized to the @@ -218,14 +248,8 @@ impl BucketVocabStore { let total: usize = tokens.iter().map(|(s, _)| s.len()).sum(); let max_id = tokens.iter().map(|(_, id)| *id).max().unwrap(); let mut bytes = Vec::with_capacity(total); - let mut entries = vec![ - Entry { - start: 0, - len: 0, - id: 0 - }; - n_slots - ]; + let mut entries = vec![Entry::default(); n_slots]; + let mut spans = vec![Span::default(); n_slots]; let mut id_to_slot = vec![u32::MAX; max_id as usize + 1]; for (s, id) in &tokens { assert!( @@ -236,11 +260,12 @@ impl BucketVocabStore { *id <= VOCAB_ID_MASK, "token id {id} needs bit 31, which holds FOLD_BIT" ); - let slot = mphf.index(&word_hash(s.as_slice())); - entries[slot] = Entry { + let (key, hash) = key_and_hash(s.as_slice()); + let slot = mphf.index(&hash); + entries[slot] = Entry { key, id: *id }; + spans[slot] = Span { start: bytes.len() as u32, len: s.len() as u16, - id: *id, }; id_to_slot[*id as usize] = slot as u32; bytes.extend_from_slice(s); @@ -250,6 +275,7 @@ impl BucketVocabStore { mphf, bytes: bytes.into_boxed_slice(), entries: entries.into_boxed_slice(), + spans: spans.into_boxed_slice(), id_to_slot: id_to_slot.into_boxed_slice(), n, } @@ -262,6 +288,7 @@ impl BucketVocabStore { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), bytes: Box::new([]), entries: Box::new([]), + spans: Box::new([]), id_to_slot: Box::new([]), n: 0, } @@ -276,24 +303,38 @@ impl BucketVocabStore { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&word_hash(q)); - + let (key, hash) = key_and_hash(q); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - // Byte equality: confirms `q` really is the token at this slot (perfect hashing only - // guarantees a valid slot for in-vocab keys; this rejects collisions and Out Of Vocab queries). - if len == q.len() && self.bytes[start..start + len] == *q { - Some(e.id & VOCAB_ID_MASK) - } else { - None + // Perfect hashing only promises a valid slot for in-vocab keys, so the slot still has to be + // verified; this rejects collisions, phantom slots and out-of-vocabulary queries. + if e.key != key || !self.confirm(slot, q) { + return None; } + Some(e.id & VOCAB_ID_MASK) + } + + /// Whether the token at `slot` really is `q`. + /// + /// A word of [`INLINE_KEY_BYTES`] bytes or fewer has already proved it: its key *is* its bytes + /// and its length, so the caller's key compare was exact and this is free. Only a longer word, + /// whose key is a hash, reads the byte slab -- the load the probe used to pay unconditionally. + #[inline(always)] + fn confirm(&self, slot: usize, q: &[u8]) -> bool { + if q.len() <= INLINE_KEY_BYTES { + return true; + } + let s = self.spans[slot]; + let start = s.start as usize; + self.bytes.get(start..start + s.len as usize) == Some(q) } /// The id for `q`, together with whether that entry may be folded. One probe and one entry /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { - self.get_bytes_foldable_hashed(q, self.hash_word(q)) + let (key, hash) = key_and_hash(q); + self.get_bytes_foldable_keyed(q, key, hash) } /// The hash this store keys `q` by. Exposed so a caller that also has to hash the same word for @@ -303,25 +344,26 @@ impl BucketVocabStore { word_hash(q) } - /// [`Self::get_bytes_foldable`] for a caller that already hashed the word with - /// [`Self::hash_word`]. + /// [`Self::get_bytes_foldable`] for a caller that already ran [`key_and_hash`] on the word. /// /// Verification is unchanged: the MPHF hands back a slot for *any* query, so the entry's bytes /// are still compared to `q` in full. Only the hashing is shared, never the check. #[inline] - pub fn get_bytes_foldable_hashed(&self, q: &[u8], hash: u64) -> Option<(u32, bool)> { + pub fn get_bytes_foldable_keyed(&self, q: &[u8], key: u64, hash: u64) -> Option<(u32, bool)> { if self.entries.is_empty() { return None; } - debug_assert_eq!(hash, self.hash_word(q), "hash does not belong to this word"); + debug_assert_eq!( + (key, hash), + key_and_hash(q), + "key/hash pair does not belong to this word" + ); let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - if len == q.len() && self.bytes[start..start + len] == *q { - Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) - } else { - None + if e.key != key || !self.confirm(slot, q) { + return None; } + Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) } /// Records that this token folds to itself. Called once per entry at load, after the proof. @@ -345,9 +387,9 @@ impl BucketVocabStore { if slot == u32::MAX { return None; // id is within range but absent from the vocab } - let e = self.entries[slot as usize]; - let start = e.start as usize; - self.bytes.get(start..start + e.len as usize) + let s = self.spans[slot as usize]; + let start = s.start as usize; + self.bytes.get(start..start + s.len as usize) } #[inline] @@ -375,11 +417,12 @@ impl BucketVocabStore { } pub fn content(&self) -> Vec<(String, u32)> { - self.entries + self.spans .iter() - .filter(|e| e.len > 0) + .zip(self.entries.iter()) + .filter(|(s, _)| s.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(_, m)| m.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -391,11 +434,12 @@ impl BucketVocabStore { /// convenient when we want to re-build a vocab pub fn byte_content(&self) -> Vec<(Vec, u32)> { - self.entries + self.spans .iter() - .filter(|e| e.len > 0) + .zip(self.entries.iter()) + .filter(|(s, _)| s.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(_, m)| m.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() } @@ -431,6 +475,66 @@ mod tests { } } + /// A slot is verified by comparing its key. For a word of [`INLINE_KEY_BYTES`] bytes or fewer + /// the key *is* the bytes and the length, so that compare is proof and the byte slab is never + /// read; a longer word's key is a hash, so it still confirms against the slab. Both lengths + /// have to reject an out-of-vocabulary word, including ones that differ only in length. + #[test] + fn out_of_vocabulary_words_are_rejected() { + let words: [&[u8]; 7] = [ + b"a", + b"ab", + b"the", + b" the", + b"abcdefg", // exactly INLINE_KEY_BYTES: key is proof + b"abcdefgh", // one over: key is a hash, slab confirms + b"a much longer token than the inline key can hold", + ]; + let toks: Vec<(Vec, u32)> = words + .iter() + .enumerate() + .map(|(i, w)| (w.to_vec(), i as u32)) + .collect(); + let vocab = BucketVocabStore::build(toks.clone()); + + for (bytes, id) in &toks { + assert_eq!(vocab.get_bytes(bytes), Some(*id), "{bytes:?} should be found"); + } + for miss in [ + &b""[..], + b"b", + b"ba", + b"abc", + b"abcdef", // prefix of a present token + b"abcdefi", // same length as a present short token + b"abcdefghi",// same length class as a present long token + b"a much longer token than the inline key can hold!", + ] { + assert_eq!(vocab.get_bytes(miss), None, "{miss:?} is not in the vocab"); + } + } + + /// `Entry::default()` leaves key 0 in the padding slots a non-minimal MPHF returns, and the + /// probe rejects those with the same single compare it uses for everything else. That only + /// works while no real word can key to 0 -- which is why [`LEN_TAG`] biases the length by one. + #[test] + fn no_real_key_is_zero() { + for w in [ + &b""[..], + b"a", + b"\0", + b"\0\0\0\0\0\0\0", // seven zero bytes: only the length tag keeps this off 0 + b"1234567", + b"12345678", + ] { + assert_ne!( + key_and_hash(w).0, + 0, + "key of {w:?} collides with the phantom-slot sentinel" + ); + } + } + #[test] fn single_token() { let vocab = BucketVocabStore::build(vec![(b"Hel".to_vec(), 0)]);