gigatoken 0.1.0 → 0.2.0

data/src/lib.rs

@@ -32,7 +32,9 @@ use crate::bindings::matcher::{SpecialTokenFound, SubstringMatcher};
 #[cfg(feature = "python")]
 use crate::bindings::padding;
 #[cfg(feature = "python")]
-use crate::bindings::pretokenize::{PretokenizerIter, pretokenized_counts, pretokenizer};
+use crate::bindings::pretokenize::{
+    PretokenizerIter, pretokenized_counts, pretokenizer, pretokenizer_scheme,
+};
 #[cfg(feature = "python")]
 use crate::bindings::sources::{
     BytesSource, FileSource, JsonlFileSource, ParquetFileSource, TextFileSource,
@@ -51,6 +53,8 @@ use pyo3::pybacked::{PyBackedBytes, PyBackedStr};
 #[cfg(feature = "python")]
 use pyo3::types::{PyBytes, PyDict, PyList};
 #[cfg(feature = "python")]
+use std::collections::HashMap;
+#[cfg(feature = "python")]
 use std::path::PathBuf;
 
 #[cfg(feature = "python")]
@@ -83,30 +87,24 @@ impl BPETokenizer {
 #[cfg(feature = "python")]
 #[pymethods]
 impl BPETokenizer {
+    /// Load from a .tiktoken rank file with the named pretokenizer scheme
+    /// and, optionally, a {content: id} mapping of special tokens. The file
+    /// carries neither, so both are the caller's to supply — see
+    /// `gigatoken.Tokenizer.from_tiktoken`, which knows them for the
+    /// encodings OpenAI publishes.
     #[staticmethod]
-    fn from_tiktoken(path: PathBuf) -> PyResult<Self> {
-        Ok(Self {
-            tokenizer: load_tokenizer::tiktoken::load_tiktoken(&path)?,
-            workers: WorkerPool::new(),
-        })
-    }
-    /// Load from a tiktoken rank file plus a tokenizer_config.json carrying
-    /// the special tokens — the layout of repos that ship no tokenizer.json
-    /// (e.g. the moonshotai Kimi line) — with the named pretokenizer scheme.
-    #[staticmethod]
-    fn from_tiktoken_model(
-        model_path: PathBuf,
-        config_path: PathBuf,
+    #[pyo3(signature = (path, pretokenizer, special_tokens = None))]
+    fn from_tiktoken(
+        path: PathBuf,
         pretokenizer: &str,
+        special_tokens: Option<HashMap<String, u32>>,
     ) -> PyResult<Self> {
-        let scheme = pretokenize::PretokenizerType::from_name(pretokenizer).ok_or_else(|| {
-            pyo3::exceptions::PyValueError::new_err(format!(
-                "unknown pretokenizer scheme {pretokenizer:?}; expected one of \
-                 gpt2, gpt4, qwen2, qwen35, olmo3, deepseek_v3, o200k, nemotron, kimi"
-            ))
-        })?;
+        let scheme = pretokenizer_scheme(pretokenizer)?;
+        let special_tokens = special_tokens.unwrap_or_default().into_iter().collect();
         Ok(Self {
-            tokenizer: load_tokenizer::tiktoken::load_tiktoken_model(&model_path, &config_path, scheme)?,
+            tokenizer: bindings::cache::apply_max_cache_bytes(
+                load_tokenizer::tiktoken::load_tiktoken(&path, scheme, special_tokens)?,
+            ),
             workers: WorkerPool::new(),
         })
     }
@@ -114,7 +112,9 @@ impl BPETokenizer {
     #[staticmethod]
     fn from_hf(path: PathBuf) -> PyResult<Self> {
         Ok(Self {
-            tokenizer: load_tokenizer::hf::load_hf_bpe(&path)?,
+            tokenizer: bindings::cache::apply_max_cache_bytes(load_tokenizer::hf::load_hf_bpe(
+                &path,
+            )?),
             workers: WorkerPool::new(),
         })
     }
@@ -276,6 +276,14 @@ impl BPETokenizer {
         Ok(self.tokenizer.decode(ids.as_slice()?).collect())
     }
 
+    /// Cached pretoken entries on the single-document `encode` path's
+    /// tokenizer (batch workers keep their own caches): grows as text is
+    /// encoded, drops back toward vocab-seed level when a budgeted cache
+    /// wipes (see gigatoken.set_max_cache_bytes).
+    fn cache_entries(&self) -> usize {
+        self.tokenizer.cache_entries()
+    }
+
     fn __repr__(&self) -> PyResult<String> {
         Ok(format!("{:?}", self.tokenizer))
     }
@@ -344,6 +352,15 @@ impl SentencePieceTokenizer {
             }
         })
     }
+
+    /// Wrap a loaded model, applying the global cache budget to it and to
+    /// the single-document encode state.
+    fn with_model(model: bpe::SentencePieceBPE) -> Self {
+        let tokenizer = bindings::cache::apply_max_cache_bytes_sp(model);
+        let state =
+            bpe::sentencepiece::EncodeState::with_budget(tokenizer.max_cache_bytes());
+        Self { tokenizer, state }
+    }
 }
 
 #[cfg(feature = "python")]
@@ -351,10 +368,7 @@ impl SentencePieceTokenizer {
 impl SentencePieceTokenizer {
     #[staticmethod]
     fn from_hf(path: PathBuf) -> PyResult<Self> {
-        Ok(Self {
-            tokenizer: load_tokenizer::hf::load_hf_sentencepiece(&path)?,
-            state: bpe::sentencepiece::EncodeState::new(),
-        })
+        Ok(Self::with_model(load_tokenizer::hf::load_hf_sentencepiece(&path)?))
     }
 
     /// Encode a batch of documents in parallel, releasing the GIL. Accepts
@@ -502,6 +516,12 @@ impl SentencePieceTokenizer {
         Ok(self.tokenizer.decode(ids.as_slice()?))
     }
 
+    /// Cached unit entries on the single-document `encode` path's state
+    /// (batch encoders are per-call); see `BPETokenizer.cache_entries`.
+    fn cache_entries(&self) -> usize {
+        self.state.cache_size()
+    }
+
     fn __repr__(&self) -> PyResult<String> {
         Ok(format!("{:?}", self.tokenizer))
     }
@@ -532,19 +552,14 @@ fn load_hf_json(py: Python<'_>, data: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
         load_tokenizer::hf::HfTokenizer::Bpe(tokenizer) => Ok(Py::new(
             py,
             BPETokenizer {
-                tokenizer,
+                tokenizer: bindings::cache::apply_max_cache_bytes(tokenizer),
                 workers: WorkerPool::new(),
             },
         )?
         .into_any()),
-        load_tokenizer::hf::HfTokenizer::SentencePiece(tokenizer) => Ok(Py::new(
-            py,
-            SentencePieceTokenizer {
-                tokenizer,
-                state: bpe::sentencepiece::EncodeState::new(),
-            },
-        )?
-        .into_any()),
+        load_tokenizer::hf::HfTokenizer::SentencePiece(tokenizer) => {
+            Ok(Py::new(py, SentencePieceTokenizer::with_model(tokenizer))?.into_any())
+        }
     }
 }
 
@@ -574,5 +589,7 @@ fn gigatoken_rs<'py>(py: Python, m: &Bound<'py, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(bindings::hub::hub_file, m)?)?;
     m.add_function(wrap_pyfunction!(bindings::hub::looks_like_repo_id, m)?)?;
     m.add_function(wrap_pyfunction!(bindings::hub::get_hf_token, m)?)?;
+    m.add_function(wrap_pyfunction!(bindings::cache::set_max_cache_bytes, m)?)?;
+    m.add_function(wrap_pyfunction!(bindings::cache::get_max_cache_bytes, m)?)?;
     Ok(())
 }

data/src/load_tokenizer/hf.rs

@@ -413,6 +413,7 @@ fn build_sentencepiece(tj: &TokenizerJson) -> Result<SentencePieceBPE> {
         split_safe: Vec::new(),
         cross_pieces: Vec::new(),
         cross_prev: [0; 4],
+        max_cache_bytes: Some(crate::bpe::Tokenizer::DEFAULT_MAX_CACHE_BYTES),
     };
     model.norm_added_tokens = norm_added_tokens
         .into_iter()

data/src/load_tokenizer/tiktoken.rs

@@ -26,62 +26,30 @@ fn load_tiktoken_ranks(file_path: impl AsRef<Path>) -> Result<Vec<Vec<u8>>> {
         .collect()
 }
 
-pub fn load_tiktoken(file_path: impl AsRef<Path>) -> Result<Tokenizer> {
-    let rank_vocab = load_tiktoken_ranks(file_path)?;
-    let n_ranks = rank_vocab.len() as u32;
-    let mut tokenizer = Tokenizer::from_ranks(rank_vocab)?;
-    // Tiktoken vocab files carry no special tokens; GPT-2-family vocabs
-    // (gpt2/r50k) place <|endoftext|> at the id right after the mergeable
-    // ranks. Register it so tiktoken- and tokenizer.json-loaded tokenizers
-    // encode and decode identically.
-    tokenizer.add_special_token(b"<|endoftext|>".to_vec(), n_ranks.into());
-    Ok(tokenizer)
-}
-
-/// The fields of a HuggingFace tokenizer_config.json a tiktoken-rank repo
-/// needs: the special tokens (there is no tokenizer.json to carry them).
-#[derive(serde::Deserialize)]
-struct TokenizerConfigJson {
-    #[serde(default)]
-    added_tokens_decoder: std::collections::BTreeMap<String, AddedTokenJson>,
-}
-
-#[derive(serde::Deserialize)]
-struct AddedTokenJson {
-    content: String,
-}
-
-/// Load a tokenizer from a tiktoken rank file plus a HuggingFace
-/// `tokenizer_config.json` whose `added_tokens_decoder` carries the special
-/// tokens — the layout of repos that ship no tokenizer.json (e.g. the
-/// moonshotai Kimi/Moonlight line's `tiktoken.model`). The pretokenizer
-/// scheme is passed by the caller: such repos define their split regex only
-/// in remote code, so no shipped file can name it.
-pub fn load_tiktoken_model(
-    model_path: impl AsRef<Path>,
-    config_path: impl AsRef<Path>,
+/// Load a tokenizer from a tiktoken rank file with the pretokenizer scheme
+/// and special tokens (`(content, id)`) supplied by the caller. A .tiktoken
+/// file carries neither — its split regex and specials live in the code
+/// that defines the encoding — and a wrong scheme silently changes every
+/// encode, so neither may be defaulted here.
+pub fn load_tiktoken(
+    file_path: impl AsRef<Path>,
     pretokenizer: PretokenizerType,
+    special_tokens: Vec<(String, u32)>,
 ) -> Result<Tokenizer> {
-    let rank_vocab = load_tiktoken_ranks(model_path)?;
+    let rank_vocab = load_tiktoken_ranks(file_path)?;
     let n_ranks = rank_vocab.len() as u32;
     let mut tokenizer = Tokenizer::from_ranks(rank_vocab)?;
     tokenizer.set_pretokenizer_type(pretokenizer);
-    let config_path = config_path.as_ref();
-    let config: TokenizerConfigJson = sonic_rs::from_slice(
-        &std::fs::read(config_path)
-            .with_context(|| format!("Failed to read {}", config_path.display()))?,
-    )
-    .map_err(|e| eyre::eyre!("Failed to parse {}: {e}", config_path.display()))?;
-    for (id, token) in &config.added_tokens_decoder {
-        let id = id.parse::<u32>().with_context(|| {
-            format!("added_tokens_decoder id {id:?} in {}", config_path.display())
-        })?;
+    for (content, id) in &special_tokens {
         ensure!(
-            id >= n_ranks,
-            "added token {:?} (id {id}) overlaps the {n_ranks} mergeable ranks",
-            token.content
+            *id >= n_ranks,
+            "special token {content:?} (id {id}) overlaps the {n_ranks} mergeable ranks"
         );
-        tokenizer.add_special_token(token.content.clone().into_bytes(), id.into());
     }
+    tokenizer.add_special_tokens(
+        special_tokens
+            .into_iter()
+            .map(|(content, id)| (content.into_bytes(), id.into())),
+    );
     Ok(tokenizer)
 }

data/src/main.rs

@@ -73,8 +73,12 @@ pub fn main() {
             .unwrap()
             .join("data/tokenizers/r50k_base.tiktoken")
     });
-    let mut r50k =
-        load_tokenizer::tiktoken::load_tiktoken(&r50k_path).expect("Failed to load r50k tokenizer");
+    let mut r50k = load_tokenizer::tiktoken::load_tiktoken(
+        &r50k_path,
+        pretokenize::PretokenizerType::GPT2,
+        vec![("<|endoftext|>".to_string(), 50256)],
+    )
+    .expect("Failed to load r50k tokenizer");
     eprintln!("Loaded r50k: {:?}", r50k);
 
     // Encode with r50k (GPT-2 style: pretokenize + memoized BPE)

data/src/pretokenize/options.rs

@@ -66,10 +66,25 @@ impl PretokenizerType {
         }
     }
 
+    /// The canonical name of each variant, in variant order — the
+    /// identifiers `from_name` accepts (plus the aliases listed there).
+    /// Error messages naming the valid schemes derive from this list.
+    pub const NAMES: [&'static str; 9] = [
+        "gpt2",
+        "gpt4",
+        "qwen2",
+        "qwen35",
+        "olmo3",
+        "deepseek_v3",
+        "o200k",
+        "nemotron",
+        "kimi",
+    ];
+
     /// The scheme named by a lowercase identifier, as used by loaders whose
     /// source format carries no split regex (e.g. tiktoken-rank repos, whose
     /// regex lives in remote code) and the Python bindings. Accepts each
-    /// variant's canonical name plus the common tiktoken aliases.
+    /// canonical name from [`Self::NAMES`] plus the common tiktoken aliases.
     pub fn from_name(name: &str) -> Option<Self> {
         Some(match name {
             "gpt2" | "r50k" => PretokenizerType::GPT2,

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gigatoken
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - Eric Jacobs
@@ -68,6 +68,7 @@ files:
 - exe/gigatoken
 - ext/gigatoken/Cargo.toml
 - ext/gigatoken/extconf.rb
+- ext/gigatoken/src/cache.rs
 - ext/gigatoken/src/error.rs
 - ext/gigatoken/src/gvl.rs
 - ext/gigatoken/src/lib.rs
@@ -79,6 +80,11 @@ files:
 - lib/gigatoken/cli/bench.rb
 - lib/gigatoken/cli/support.rb
 - lib/gigatoken/cli/validate.rb
+- lib/gigatoken/encodings.rb
+- lib/gigatoken/encodings/PROVENANCE.md
+- lib/gigatoken/encodings/cl100k_base.tiktoken
+- lib/gigatoken/encodings/o200k_base.tiktoken
+- lib/gigatoken/encodings/r50k_base.tiktoken
 - lib/gigatoken/hub.rb
 - lib/gigatoken/packed_result.rb
 - lib/gigatoken/tokenizer.rb
@@ -86,6 +92,7 @@ files:
 - rust-toolchain.toml
 - src/batch.rs
 - src/bindings/bridge.rs
+- src/bindings/cache.rs
 - src/bindings/hub.rs
 - src/bindings/matcher.rs
 - src/bindings/mod.rs
@@ -139,6 +146,7 @@ licenses:
 - MIT
 metadata:
   bug_tracker_uri: https://github.com/jetpks/gigatoken-rb/issues
+  changelog_uri: https://github.com/jetpks/gigatoken-rb/blob/main/CHANGELOG.md
   homepage_uri: https://github.com/jetpks/gigatoken-rb
   source_code_uri: https://github.com/jetpks/gigatoken-rb
 rdoc_options: []