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//! * [`estimate`] -- a dependency-free, deterministic *token count
45//!   estimator* ([`estimate::estimate_tokens`]) for when there is no
46//!   loaded vocab to run the real BPE against. It exists to make LLM cost
47//!   visible (read-vs-outline) rather than to reproduce exact ids.
48//!
49//! # What this crate does *not* do
50//!
51//! No model inference, no chat templating, no attention-mask bookkeeping.
52//! Those belong to `kopitiam-runtime`, which will code against the
53//! [`Tokenizer`] trait rather than `BpeTokenizer` directly, the same way
54//! it will code against traits from `kopitiam-tensor` and `kopitiam-loader`
55//! rather than their concrete types.
56
57pub mod bpe;
58pub mod byte_map;
59pub mod estimate;
60pub mod loader;
61pub mod merges;
62pub mod pretokenize;
63pub mod specials;
64pub mod vocab;
65
66pub use bpe::BpeTokenizer;
67pub use estimate::{LineEstimate, estimate_tokens, estimate_tokens_by_line};
68pub use loader::from_tokenizer_json;
69
70use kopitiam_core::Result;
71
72/// What every tokenizer implementation in the Kopitiam Runtime must
73/// provide. `kopitiam-runtime` codes against this trait, not
74/// [`BpeTokenizer`] directly, so a future tokenizer scheme (e.g.
75/// SentencePiece for a model family that needs it) can be dropped in
76/// without touching the runtime's call sites.
77pub trait Tokenizer {
78    /// Encodes `text` into token ids. Byte-level implementations never
79    /// fail here -- every possible `&str` (any valid UTF-8, including
80    /// emoji, CJK, and arbitrary punctuation runs) is representable as
81    /// some sequence of byte-derived tokens -- so this returning `Result`
82    /// is about leaving room for non-byte-level implementations, not
83    /// because `BpeTokenizer::encode` can actually fail today.
84    fn encode(&self, text: &str) -> Result<Vec<u32>>;
85
86    /// Decodes token ids back into text. Fails gracefully (rather than
87    /// panicking) on an id outside the vocab; never fails just because the
88    /// concatenated bytes happen not to be valid UTF-8 on their own (see
89    /// [`bpe::BpeTokenizer::decode`]'s docs for why that case is handled
90    /// losslessly-in-practice-but-not-by-contract via lossy UTF-8
91    /// conversion instead of an error).
92    fn decode(&self, ids: &[u32]) -> Result<String>;
93
94    /// Total number of ids in the vocabulary, including any registered
95    /// special tokens.
96    fn vocab_size(&self) -> usize;
97}