Skip to main content

kopitiam_tokenizer/
lib.rs

1//! Kopitiam Runtime: byte-level BPE tokenizer.
2//!
3//! This crate turns text into the token ids a GPT-2/GPT-3/GPT-4- or
4//! Qwen-family model was trained on, and back. It is a from-scratch,
5//! original Rust implementation -- not a wrapper around the `tokenizers`
6//! crate -- because the whole point of the Kopitiam Runtime (see
7//! `docs/ai-decisions/AID-0001`) is to own this layer rather than depend
8//! on it, in keeping with this workspace's Pure Rust Core commitment.
9//!
10//! # Byte-level BPE, in one paragraph
11//!
12//! Classic word-level BPE needs an `<UNK>` token for anything outside its
13//! training vocabulary -- a stray emoji, a typo, a byte sequence that
14//! isn't valid text at all. Byte-level BPE sidesteps this entirely by
15//! making the *256 possible byte values* the base alphabet instead of
16//! characters or words: every possible input, valid UTF-8 or not, is some
17//! sequence of bytes, and every byte already has a token id. There is
18//! therefore no `<UNK>` in this crate's design and no failure mode for
19//! encoding -- only [`Tokenizer::decode`] can fail, and only on a
20//! genuinely unknown id (see its docs). This is the scheme GPT-2 through
21//! GPT-4 and the Qwen family all use.
22//!
23//! # How the pieces fit together
24//!
25//! * [`byte_map`] -- the reversible byte <-> "printable Unicode character"
26//!   mapping that lets a byte-level vocab be written as JSON text at all.
27//!   Used only when loading/saving `tokenizer.json`; the runtime path
28//!   never touches it.
29//! * [`vocab`] -- [`vocab::Vocab`], the token (raw bytes) <-> id table.
30//! * [`merges`] -- [`merges::MergeTable`], the ordered "which adjacent pair
31//!   merges first, into what" rules learned during BPE training.
32//! * [`pretokenize`] -- splits text into words/numbers/punctuation/
33//!   whitespace-run chunks *before* BPE runs, matching the GPT-2 pattern
34//!   without `regex`'s unsupported lookahead (see that module's docs --
35//!   this is the single easiest place to introduce a silent correctness
36//!   bug in a tokenizer).
37//! * [`specials`] -- atomic special-token matching (`<|endoftext|>` and
38//!   friends), so they are pulled out before pre-tokenization ever sees
39//!   them.
40//! * [`bpe`] -- [`bpe::BpeTokenizer`], the concrete [`Tokenizer`] that
41//!   wires the above into `encode`/`decode`.
42//! * [`loader`] -- parses the HuggingFace `tokenizer.json` shape into a
43//!   [`bpe::BpeTokenizer`].
44//!
45//! # What this crate does *not* do
46//!
47//! No model inference, no chat templating, no attention-mask bookkeeping.
48//! Those belong to `kopitiam-runtime`, which will code against the
49//! [`Tokenizer`] trait rather than `BpeTokenizer` directly, the same way
50//! it will code against traits from `kopitiam-tensor` and `kopitiam-loader`
51//! rather than their concrete types.
52
53pub mod bpe;
54pub mod byte_map;
55pub mod loader;
56pub mod merges;
57pub mod pretokenize;
58pub mod specials;
59pub mod vocab;
60
61pub use bpe::BpeTokenizer;
62pub use loader::from_tokenizer_json;
63
64use kopitiam_core::Result;
65
66/// What every tokenizer implementation in the Kopitiam Runtime must
67/// provide. `kopitiam-runtime` codes against this trait, not
68/// [`BpeTokenizer`] directly, so a future tokenizer scheme (e.g.
69/// SentencePiece for a model family that needs it) can be dropped in
70/// without touching the runtime's call sites.
71pub trait Tokenizer {
72    /// Encodes `text` into token ids. Byte-level implementations never
73    /// fail here -- every possible `&str` (any valid UTF-8, including
74    /// emoji, CJK, and arbitrary punctuation runs) is representable as
75    /// some sequence of byte-derived tokens -- so this returning `Result`
76    /// is about leaving room for non-byte-level implementations, not
77    /// because `BpeTokenizer::encode` can actually fail today.
78    fn encode(&self, text: &str) -> Result<Vec<u32>>;
79
80    /// Decodes token ids back into text. Fails gracefully (rather than
81    /// panicking) on an id outside the vocab; never fails just because the
82    /// concatenated bytes happen not to be valid UTF-8 on their own (see
83    /// [`bpe::BpeTokenizer::decode`]'s docs for why that case is handled
84    /// losslessly-in-practice-but-not-by-contract via lossy UTF-8
85    /// conversion instead of an error).
86    fn decode(&self, ids: &[u32]) -> Result<String>;
87
88    /// Total number of ids in the vocabulary, including any registered
89    /// special tokens.
90    fn vocab_size(&self) -> usize;
91}