forge/tokenizer/
char.rs

1//! Character-level tokenizer, matching nanoGPT's `shakespeare_char` prepare
2//! step: the vocabulary is the sorted set of distinct characters in the
3//! corpus, and a token id is that character's index in the sorted order.
4//!
5//! For Tiny Shakespeare this yields 65 tokens, which shrinks the (weight-tied)
6//! embedding table from 9.65M parameters at the GPT-2 BPE vocab to 24,960 —
7//! the difference between a 45.9 MB checkpoint that is 84% untrained random
8//! rows and a 43.1 MB one where every parameter does work.
9
10use std::collections::{BTreeSet, HashMap};
11use std::path::Path;
12
13use crate::error::{ForgeError, Result};
14use crate::tokenizer::Tokenizer;
15
16pub struct CharTokenizer {
17    /// id -> char, indexed directly.
18    itos: Vec<char>,
19    stoi: HashMap<char, u32>,
20}
21
22impl CharTokenizer {
23    /// Build the vocab from the sorted unique characters of `text`
24    /// (nanoGPT's rule). Ids are indices into that sorted order.
25    pub fn from_corpus(text: &str) -> Self {
26        Self::from_chars(text.chars().collect::<BTreeSet<char>>())
27    }
28
29    fn from_chars(sorted: BTreeSet<char>) -> Self {
30        let itos: Vec<char> = sorted.into_iter().collect();
31        let stoi = itos
32            .iter()
33            .enumerate()
34            .map(|(i, &c)| (c, i as u32))
35            .collect();
36        CharTokenizer { itos, stoi }
37    }
38
39    /// The vocabulary in id order.
40    pub fn chars(&self) -> &[char] {
41        &self.itos
42    }
43
44    /// Characters of `text` that are outside the vocabulary, deduplicated and
45    /// in first-appearance order. Callers that want to warn a user before
46    /// [`CharTokenizer::encode`] rejects the input use this.
47    pub fn unknown_chars(&self, text: &str) -> Vec<char> {
48        let mut seen = Vec::new();
49        for c in text.chars() {
50            if !self.stoi.contains_key(&c) && !seen.contains(&c) {
51                seen.push(c);
52            }
53        }
54        seen
55    }
56
57    /// Encode, silently dropping characters outside the vocabulary. Pair with
58    /// [`CharTokenizer::unknown_chars`] to tell the user what was dropped —
59    /// the strict [`CharTokenizer::encode`] is the default for good reason.
60    pub fn encode_lossy(&self, text: &str) -> Vec<u32> {
61        text.chars()
62            .filter_map(|c| self.stoi.get(&c).copied())
63            .collect()
64    }
65
66    /// `{"<char>": id}` — the same shape as GPT-2's `vocab.json`, so the same
67    /// loader path and the same filename work for both tokenizers.
68    pub fn to_json(&self) -> String {
69        let map: HashMap<String, u32> = self
70            .stoi
71            .iter()
72            .map(|(&c, &i)| (c.to_string(), i))
73            .collect();
74        // Infallible: keys are strings and values are u32.
75        serde_json::to_string(&map).expect("char vocab serializes")
76    }
77
78    pub fn from_json(s: &str) -> Result<Self> {
79        let map: HashMap<String, u32> = serde_json::from_str(s)?;
80        let mut itos = vec!['\0'; map.len()];
81        let mut filled = vec![false; map.len()];
82        for (k, id) in &map {
83            let mut cs = k.chars();
84            let (Some(c), None) = (cs.next(), cs.next()) else {
85                return Err(ForgeError::Tokenizer(format!(
86                    "char vocab entry {k:?} is not a single character"
87                )));
88            };
89            let idx = *id as usize;
90            if idx >= itos.len() {
91                return Err(ForgeError::Tokenizer(format!(
92                    "char vocab id {id} out of range for a {}-entry vocab",
93                    itos.len()
94                )));
95            }
96            if filled[idx] {
97                return Err(ForgeError::Tokenizer(format!(
98                    "char vocab id {id} assigned twice"
99                )));
100            }
101            itos[idx] = c;
102            filled[idx] = true;
103        }
104        let stoi = itos
105            .iter()
106            .enumerate()
107            .map(|(i, &c)| (c, i as u32))
108            .collect();
109        Ok(CharTokenizer { itos, stoi })
110    }
111
112    /// Write `vocab.json` next to a checkpoint. Re-deriving the vocab at
113    /// inference time from a different corpus would silently shift every id,
114    /// so it ships with the weights.
115    pub fn save_json(&self, path: impl AsRef<Path>) -> Result<()> {
116        std::fs::write(path, self.to_json())?;
117        Ok(())
118    }
119
120    pub fn from_json_file(path: impl AsRef<Path>) -> Result<Self> {
121        Self::from_json(&std::fs::read_to_string(path)?)
122    }
123}
124
125impl Tokenizer for CharTokenizer {
126    /// Unknown characters are an error rather than a silent drop: a prompt
127    /// containing them would otherwise generate a plausible continuation of
128    /// something the user did not type. See [`CharTokenizer::encode_lossy`].
129    fn encode(&self, text: &str) -> Result<Vec<u32>> {
130        text.chars()
131            .map(|c| {
132                self.stoi.get(&c).copied().ok_or_else(|| {
133                    ForgeError::Tokenizer(format!(
134                        "character {c:?} is outside the {}-token vocabulary",
135                        self.itos.len()
136                    ))
137                })
138            })
139            .collect()
140    }
141
142    fn decode(&self, ids: &[u32]) -> String {
143        ids.iter()
144            .filter_map(|&i| self.itos.get(i as usize))
145            .collect()
146    }
147
148    /// A char vocab has no partial-UTF-8 problem — every token is a whole
149    /// character — but the streaming path calls this for both tokenizers.
150    fn decode_bytes(&self, ids: &[u32]) -> Vec<u8> {
151        self.decode(ids).into_bytes()
152    }
153
154    fn vocab_size(&self) -> usize {
155        self.itos.len()
156    }
157}
forge/tokenizer/char.rs

forge/tokenizer/
char.rs
