diff --git a/.github/workflows/speed-regression.yml b/.github/workflows/speed-regression.yml index 7c9bf8b..739e077 100644 --- a/.github/workflows/speed-regression.yml +++ b/.github/workflows/speed-regression.yml @@ -62,12 +62,11 @@ jobs: fi echo "Benchmarking base commit: $BASE_SHA" git stash --include-untracked || true - git checkout "$BASE_SHA" + git checkout "$BASE_SHA" -- . || git checkout --force "$BASE_SHA" cargo run --example speed_regression --release --features hf \ > benchmark_base.json 2>benchmark_base.log || echo '{"results":[]}' > benchmark_base.json cat benchmark_base.log - git checkout - - git stash pop || true + git checkout --force - || true - name: Compare results if: github.event_name == 'push' || github.event_name == 'pull_request' diff --git a/.github/workflows/tokenizer-accuracy.yml b/.github/workflows/tokenizer-accuracy.yml index 495c8e6..3d85704 100644 --- a/.github/workflows/tokenizer-accuracy.yml +++ b/.github/workflows/tokenizer-accuracy.yml @@ -3,6 +3,8 @@ name: Tokenizer Accuracy on: push: branches: [main] + pull_request: + branches: [main] workflow_dispatch: env: @@ -41,4 +43,4 @@ jobs: fi - name: Run tokenizer accuracy tests - run: cargo test -p tokie --test accuracy --features hf -- --ignored + run: cargo test -p tokie --test accuracy --features build -- --ignored diff --git a/Cargo.lock b/Cargo.lock index 972d078..c62a33f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -37,12 +37,56 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + [[package]] name = "anstyle" version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.60.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.60.2", +] + [[package]] name = "anyhow" version = "1.0.102" @@ -198,6 +242,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" dependencies = [ "clap_builder", + "clap_derive", ] [[package]] @@ -206,8 +251,22 @@ version = "4.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" dependencies = [ + "anstream", "anstyle", "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1110bd8a634a1ab8cb04345d8d878267d57c3cf1b38d91b71af6686408bbca6a" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -216,6 +275,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "compact_str" version = "0.9.0" @@ -861,6 +926,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" version = "0.10.5" @@ -1043,6 +1114,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "onig" version = "6.5.1" @@ -1173,18 +1250,18 @@ dependencies = [ [[package]] name = "pretokie" version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cb5f02994a0b0264a8eb100161df7d304d41e3bb9d52e0c0c1c9bac101eeb9" dependencies = [ - "foldhash", "regex-automata", "unicode-general-category", ] [[package]] name = "pretokie" -version = "0.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52cb5f02994a0b0264a8eb100161df7d304d41e3bb9d52e0c0c1c9bac101eeb9" +version = "0.0.4" dependencies = [ + "foldhash", "regex-automata", "unicode-general-category", ] @@ -1716,13 +1793,23 @@ dependencies = [ "hf-hub", "icu_normalizer 1.5.0", "memchr", - "pretokie 0.0.3 (registry+https://github.com/rust-lang/crates.io-index)", + "pretokie 0.0.3", "serde", "serde_json", "smallvec", "tiktoken-rs", "tokenizers", "unicode-general-category", + "ureq", +] + +[[package]] +name = "tokie-cli" +version = "0.0.1" +dependencies = [ + "clap", + "hf-hub", + "tokie", ] [[package]] @@ -1833,6 +1920,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "version_check" version = "0.9.5" diff --git a/crates/tokie-cli/Cargo.toml b/crates/tokie-cli/Cargo.toml new file mode 100644 index 0000000..62f6e04 --- /dev/null +++ b/crates/tokie-cli/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "tokie-cli" +version = "0.0.1" +edition = "2024" +description = "CLI for building, verifying, and uploading tokie tokenizers" +license = "MIT OR Apache-2.0" +repository = "https://github.com/chonkie-inc/tokie" + +[[bin]] +name = "tokie" +path = "src/main.rs" + +[dependencies] +tokie = { path = "../tokie", features = ["hf", "build"] } +clap = { version = "4", features = ["derive"] } +hf-hub = { version = "0.4.3", default-features = false, features = ["ureq"] } diff --git a/crates/tokie-cli/src/main.rs b/crates/tokie-cli/src/main.rs new file mode 100644 index 0000000..2775758 --- /dev/null +++ b/crates/tokie-cli/src/main.rs @@ -0,0 +1,194 @@ +use std::path::PathBuf; + +use clap::{Parser, Subcommand}; +use tokie::build::VerifyResult; + +#[derive(Parser)] +#[command(name = "tokie", about = "Build, verify, and upload tokie tokenizers")] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +enum Commands { + /// Convert a HuggingFace tokenizer.json to .tkz format + Convert { + /// HuggingFace repo ID (e.g., "openai-community/gpt2", "meta-llama/Llama-3.2-1B") + repo_id: String, + + /// Output path for the .tkz file + #[arg(short, long)] + output: PathBuf, + + /// Verify the converted .tkz against the reference tokenizer + #[arg(short, long)] + verify: bool, + + /// Upload to tokiers/ org on HuggingFace Hub (implies --verify) + #[arg(short, long)] + upload: bool, + + /// HuggingFace API token (falls back to HF_TOKEN env or cached token) + #[arg(long)] + token: Option, + }, + /// Verify a .tkz tokenizer against the reference backend + Verify { + /// HuggingFace repo ID + repo_id: String, + + /// Path to the .tkz file (if omitted, downloads from tokiers/) + #[arg(long)] + tkz: Option, + }, +} + +const MAX_MISMATCHES_SHOWN: usize = 5; + +fn print_verify_result(result: &VerifyResult) { + if result.failed == 0 { + println!("OK {}/{} chunks pass", result.passed, result.total); + return; + } + + println!( + "FAILED {}/{} chunks pass ({} mismatches)", + result.passed, result.total, result.failed + ); + + let shown = result.mismatches.len().min(MAX_MISMATCHES_SHOWN); + for m in &result.mismatches[..shown] { + println!(" MISMATCH: \"{}\"", m.text); + + // Find first divergence index + let first_diff = m + .tokie_ids + .iter() + .zip(m.reference_ids.iter()) + .position(|(a, b)| a != b) + .unwrap_or(m.tokie_ids.len().min(m.reference_ids.len())); + + let start = first_diff.saturating_sub(2); + let tokie_end = (first_diff + 5).min(m.tokie_ids.len()); + let ref_end = (first_diff + 5).min(m.reference_ids.len()); + + println!( + " first diff at token {}: tokie[{}..{}]={:?} ref[{}..{}]={:?}", + first_diff, + start, + tokie_end, + &m.tokie_ids[start..tokie_end], + start, + ref_end, + &m.reference_ids[start..ref_end], + ); + } + + let remaining = result.mismatches.len() - shown; + if remaining > 0 { + println!(" ... and {remaining} more mismatches"); + } +} + +fn main() { + let cli = Cli::parse(); + + match cli.command { + Commands::Convert { + repo_id, + output, + verify, + upload, + token, + } => { + cmd_convert(&repo_id, &output, verify || upload, upload, token.as_deref()); + } + Commands::Verify { repo_id, tkz } => { + cmd_verify(&repo_id, tkz); + } + } +} + +fn cmd_convert(repo_id: &str, output: &PathBuf, verify: bool, upload: bool, token: Option<&str>) { + print!("Converting {repo_id} ... "); + + match tokie::build::convert(repo_id, output) { + Ok(result) => { + println!( + "OK vocab={} size={:.1}KB", + result.vocab_size, + result.file_size_bytes as f64 / 1024.0 + ); + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + + if verify { + print!("Verifying ... "); + match tokie::build::verify(repo_id, output) { + Ok(result) => { + let failed = result.failed > 0; + print_verify_result(&result); + if failed { + std::process::exit(1); + } + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + } + + if upload { + let tokiers_name = repo_id.rsplit('/').next().unwrap_or(repo_id); + print!("Uploading to tokiers/{tokiers_name} ... "); + match tokie::build::upload(output, tokiers_name, token) { + Ok(()) => println!("OK"), + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + } +} + +fn cmd_verify(repo_id: &str, tkz: Option) { + let tkz_path = if let Some(path) = tkz { + path + } else { + let name = repo_id.rsplit('/').next().unwrap_or(repo_id); + print!("Downloading tokiers/{name} ... "); + let api = hf_hub::api::sync::ApiBuilder::new().build().unwrap(); + let repo = hf_hub::Repo::model(format!("tokiers/{name}")); + match api.repo(repo).get("tokenizer.tkz") { + Ok(p) => { + println!("OK"); + p + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + }; + + print!("Verifying {repo_id} ... "); + match tokie::build::verify(repo_id, &tkz_path) { + Ok(result) => { + let failed = result.failed > 0; + print_verify_result(&result); + if failed { + std::process::exit(1); + } + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } +} diff --git a/crates/tokie/Cargo.toml b/crates/tokie/Cargo.toml index b5224e0..02968e4 100644 --- a/crates/tokie/Cargo.toml +++ b/crates/tokie/Cargo.toml @@ -13,6 +13,7 @@ categories = ["text-processing", "science"] [features] default = [] hf = ["hf-hub"] +build = ["hf", "dep:tokenizers", "dep:tiktoken-rs", "dep:ureq", "dep:zip", "dep:dirs-next"] [dependencies] crc32fast = "1.4" @@ -27,11 +28,14 @@ icu_normalizer = "1.5" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" unicode-general-category = "1.0" +tokenizers = { version = "0.22", features = ["http"], optional = true } +tiktoken-rs = { version = "0.6", optional = true } +ureq = { version = "2", optional = true } +zip = { version = "2", default-features = false, features = ["deflate"], optional = true } +dirs-next = { version = "2", optional = true } [dev-dependencies] criterion = "0.5" -tokenizers = { version = "0.22", features = ["http"] } -tiktoken-rs = "0.6" [[bench]] name = "tokenize" @@ -55,15 +59,15 @@ required-features = ["hf"] [[example]] name = "test_llama4" -required-features = ["hf"] +required-features = ["build"] [[example]] name = "test_enwik8" -required-features = ["hf"] +required-features = ["build"] [[example]] name = "verify_embedding_models" -required-features = ["hf"] +required-features = ["build"] [[example]] name = "test_postprocessor" @@ -71,7 +75,7 @@ required-features = ["hf"] [[example]] name = "convert_and_verify" -required-features = ["hf"] +required-features = ["build"] [[example]] name = "regenerate_all_v11" @@ -95,7 +99,7 @@ required-features = ["hf"] [[example]] name = "diagnose_failures" -required-features = ["hf"] +required-features = ["build"] [[example]] name = "speed_regression" @@ -104,3 +108,11 @@ required-features = ["hf"] [[example]] name = "regenerate_from_json" required-features = ["hf"] + +[[example]] +name = "bench_vs_hf" +required-features = ["build"] + +[[example]] +name = "check_models" +required-features = ["build"] diff --git a/crates/tokie/src/build.rs b/crates/tokie/src/build.rs new file mode 100644 index 0000000..fa492b5 --- /dev/null +++ b/crates/tokie/src/build.rs @@ -0,0 +1,272 @@ +//! Build tools for converting, verifying, and uploading .tkz tokenizers. +//! +//! This module is only available when the `build` feature is enabled: +//! ```toml +//! tokie = { version = "0.1", features = ["build"] } +//! ``` + +use std::path::{Path, PathBuf}; + +use hf_hub::Repo; + +use crate::hf::JsonLoadError; +use crate::serde::SerdeError; +use crate::Tokenizer; + +#[derive(Debug)] +pub struct ConvertResult { + pub vocab_size: usize, + pub file_size_bytes: u64, +} + +#[derive(Debug)] +pub struct Mismatch { + pub text: String, + pub tokie_ids: Vec, + pub reference_ids: Vec, +} + +#[derive(Debug)] +pub struct VerifyResult { + pub total: usize, + pub passed: usize, + pub failed: usize, + pub mismatches: Vec, +} + +#[derive(Debug)] +pub enum BuildError { + Download(String), + LoadJson(JsonLoadError), + SaveTkz(SerdeError), + Verification { result: VerifyResult }, + Upload(String), + HubInit(String), +} + +impl std::fmt::Display for BuildError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BuildError::Download(e) => write!(f, "download failed: {e}"), + BuildError::LoadJson(e) => write!(f, "failed to load tokenizer.json: {e}"), + BuildError::SaveTkz(e) => write!(f, "failed to save .tkz: {e}"), + BuildError::Verification { result } => { + write!(f, "verification failed: {}/{} texts mismatched", result.failed, result.total) + } + BuildError::Upload(e) => write!(f, "upload failed: {e}"), + BuildError::HubInit(e) => write!(f, "HF Hub init failed: {e}"), + } + } +} + +impl std::error::Error for BuildError {} + +fn tiktoken_encoding_name(repo_id: &str) -> Option<&'static str> { + match repo_id.to_ascii_lowercase().as_str() { + "xenova/gpt-4" => Some("cl100k_base"), + "xenova/gpt-4o" => Some("o200k_base"), + "xenova/text-davinci-003" => Some("p50k_base"), + _ => None, + } +} + +const ENWIK8_URL: &str = "http://mattmahoney.net/dc/enwik8.zip"; + +/// Download and cache enwik8. Returns the path to the cached file. +fn ensure_enwik8() -> Result { + let cache_dir = dirs_next::cache_dir() + .unwrap_or_else(|| PathBuf::from("/tmp")) + .join("tokie"); + let enwik8_path = cache_dir.join("enwik8"); + + if enwik8_path.exists() { + return Ok(enwik8_path); + } + + std::fs::create_dir_all(&cache_dir).ok(); + + let zip_path = cache_dir.join("enwik8.zip"); + let response = ureq::get(ENWIK8_URL) + .call() + .map_err(|e| BuildError::Download(format!("failed to download enwik8: {e}")))?; + + let mut bytes = Vec::new(); + response + .into_reader() + .read_to_end(&mut bytes) + .map_err(|e| BuildError::Download(format!("failed to read enwik8 response: {e}")))?; + std::fs::write(&zip_path, &bytes) + .map_err(|e| BuildError::Download(format!("failed to write enwik8.zip: {e}")))?; + + // Extract the zip + let file = std::fs::File::open(&zip_path) + .map_err(|e| BuildError::Download(format!("failed to open enwik8.zip: {e}")))?; + let mut archive = zip::ZipArchive::new(file) + .map_err(|e| BuildError::Download(format!("failed to read enwik8.zip: {e}")))?; + let mut entry = archive + .by_index(0) + .map_err(|e| BuildError::Download(format!("failed to extract enwik8: {e}")))?; + let mut content = Vec::new(); + std::io::Read::read_to_end(&mut entry, &mut content) + .map_err(|e| BuildError::Download(format!("failed to read enwik8 entry: {e}")))?; + std::fs::write(&enwik8_path, &content) + .map_err(|e| BuildError::Download(format!("failed to write enwik8: {e}")))?; + + // Clean up zip + std::fs::remove_file(&zip_path).ok(); + + Ok(enwik8_path) +} + + +/// Convert a HuggingFace repo's tokenizer.json to .tkz format. +pub fn convert(repo_id: &str, output: &Path) -> Result { + let api = hf_hub::api::sync::ApiBuilder::new() + .build() + .map_err(|e| BuildError::HubInit(e.to_string()))?; + + let repo = Repo::model(repo_id.to_string()); + let repo_api = api.repo(repo); + + let json_path = repo_api + .get("tokenizer.json") + .map_err(|e| BuildError::Download(e.to_string()))?; + + let tokenizer = Tokenizer::from_json(&json_path).map_err(BuildError::LoadJson)?; + + if let Some(parent) = output.parent() { + std::fs::create_dir_all(parent).ok(); + } + tokenizer.to_file(output).map_err(BuildError::SaveTkz)?; + + let file_size_bytes = std::fs::metadata(output).map(|m| m.len()).unwrap_or(0); + + Ok(ConvertResult { + vocab_size: tokenizer.vocab_size(), + file_size_bytes, + }) +} + +/// Verify a .tkz file against a reference tokenizer backend using enwik8. +/// +/// Downloads and caches enwik8 (~100MB), encodes the first 1MB as a single +/// string with both backends, and compares the full token ID vectors. +pub fn verify(repo_id: &str, tkz_path: &Path) -> Result { + let tokie_tok = Tokenizer::from_file(tkz_path).map_err(BuildError::SaveTkz)?; + + let enwik8_path = ensure_enwik8()?; + let raw = std::fs::read(&enwik8_path) + .map_err(|e| BuildError::Download(format!("failed to read enwik8: {e}")))?; + // Use first 1MB to keep verification fast + let truncated = &raw[..raw.len().min(1_000_000)]; + let text = String::from_utf8_lossy(truncated); + + let tokie_ids = tokie_tok.encode(&text, false).ids; + + let ref_ids = if let Some(encoding_name) = tiktoken_encoding_name(repo_id) { + let tiktoken = match encoding_name { + "cl100k_base" => tiktoken_rs::cl100k_base(), + "o200k_base" => tiktoken_rs::o200k_base(), + "p50k_base" => tiktoken_rs::p50k_base(), + _ => unreachable!(), + } + .expect("failed to load tiktoken encoding"); + + tiktoken + .encode_with_special_tokens(&text) + .into_iter() + .map(|id| id as u32) + .collect::>() + } else { + let mut hf_tok = tokenizers::Tokenizer::from_pretrained(repo_id, None) + .map_err(|e| BuildError::Download(format!("HF tokenizer load failed: {e}")))?; + let _ = hf_tok.with_truncation(None); + + let hf_encoding = hf_tok + .encode(text.to_string(), false) + .map_err(|e| BuildError::Download(format!("HF encode failed: {e}")))?; + hf_encoding.get_ids().to_vec() + }; + + if tokie_ids == ref_ids { + return Ok(VerifyResult { + total: tokie_ids.len(), + passed: tokie_ids.len(), + failed: 0, + mismatches: vec![], + }); + } + + // Find the first divergence point + let first_diff = tokie_ids + .iter() + .zip(ref_ids.iter()) + .position(|(a, b)| a != b) + .unwrap_or(tokie_ids.len().min(ref_ids.len())); + + let mismatch = Mismatch { + text: format!( + "enwik8 (1MB): {} tokie tokens vs {} ref tokens, first diff at token {}", + tokie_ids.len(), + ref_ids.len(), + first_diff + ), + tokie_ids: tokie_ids[first_diff.saturating_sub(2)..(first_diff + 5).min(tokie_ids.len())] + .to_vec(), + reference_ids: ref_ids + [first_diff.saturating_sub(2)..(first_diff + 5).min(ref_ids.len())] + .to_vec(), + }; + + Ok(VerifyResult { + total: tokie_ids.len().max(ref_ids.len()), + passed: first_diff, + failed: tokie_ids.len().max(ref_ids.len()) - first_diff, + mismatches: vec![mismatch], + }) +} + +/// Upload a .tkz file to the tokiers/ org on HuggingFace Hub. +/// +/// Uses the HF Hub HTTP API directly (hf-hub crate is download-only). +/// If `token` is None, falls back to HF_TOKEN env var. +pub fn upload(tkz_path: &Path, tokiers_name: &str, token: Option<&str>) -> Result<(), BuildError> { + if !tkz_path.exists() { + return Err(BuildError::Upload(format!( + "file not found: {}", + tkz_path.display() + ))); + } + + let token_str = token + .map(|t| t.to_string()) + .or_else(|| std::env::var("HF_TOKEN").ok()) + .ok_or_else(|| { + BuildError::Upload( + "no HF token found — pass --token or set HF_TOKEN env var".to_string(), + ) + })?; + + let file_content = std::fs::read(tkz_path) + .map_err(|e| BuildError::Upload(format!("failed to read {}: {e}", tkz_path.display())))?; + + let repo_id = format!("tokiers/{tokiers_name}"); + let url = format!( + "https://huggingface.co/api/models/{repo_id}/upload/main/tokenizer.tkz", + ); + + let response = ureq::put(&url) + .set("Authorization", &format!("Bearer {token_str}")) + .set("Content-Type", "application/octet-stream") + .send_bytes(&file_content) + .map_err(|e| BuildError::Upload(format!("HTTP upload to {repo_id} failed: {e}")))?; + + if response.status() >= 400 { + return Err(BuildError::Upload(format!( + "upload returned HTTP {}: check your token has write access to tokiers/ org", + response.status() + ))); + } + + Ok(()) +} diff --git a/crates/tokie/src/lib.rs b/crates/tokie/src/lib.rs index c82464e..f987900 100644 --- a/crates/tokie/src/lib.rs +++ b/crates/tokie/src/lib.rs @@ -57,6 +57,8 @@ pub mod encoder; pub mod hf; #[cfg(feature = "hf")] mod hub; +#[cfg(feature = "build")] +pub mod build; pub mod normalizer; pub mod padding; mod postprocessor; @@ -70,6 +72,8 @@ pub use decoder::{Decoder, DecoderType, VocabDecoder}; pub use hf::JsonLoadError; #[cfg(feature = "hf")] pub use hub::{FromPretrainedOptions, HubError}; +#[cfg(feature = "build")] +pub use build::{BuildError, ConvertResult, VerifyResult, Mismatch}; pub use normalizer::{bert_uncased_normalize, clean_text, fnr, metaspace_normalize, strip_accents, FnrFinder, Normalizer}; pub use padding::{Encoding, PaddingParams, PaddingStrategy, PaddingDirection, TruncationParams, TruncationStrategy, TruncationDirection}; pub use postprocessor::PostProcessor; diff --git a/crates/tokie/tests/accuracy.rs b/crates/tokie/tests/accuracy.rs index be43c0c..f3fcd84 100644 --- a/crates/tokie/tests/accuracy.rs +++ b/crates/tokie/tests/accuracy.rs @@ -4,7 +4,7 @@ //! //! Requires network access and benches/data/enwik8 (1MB used). -#![cfg(feature = "hf")] +#![cfg(feature = "build")] use std::path::Path; diff --git a/docs/superpowers/plans/2026-04-21-build-module-and-cli.md b/docs/superpowers/plans/2026-04-21-build-module-and-cli.md new file mode 100644 index 0000000..859ff0c --- /dev/null +++ b/docs/superpowers/plans/2026-04-21-build-module-and-cli.md @@ -0,0 +1,709 @@ +# Build Module + CLI Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a `build` feature to the tokie crate with convert/verify/upload functions, and a `tokie-cli` crate that exposes them as `tokie convert` and `tokie verify` CLI commands. + +**Architecture:** The `build` module lives in the core `tokie` crate behind a `build` feature flag that pulls in `tokenizers` and `tiktoken-rs` as optional deps. The `tokie-cli` crate is a thin clap wrapper that calls into `tokie::build::*`. The build module reuses existing `Tokenizer::from_json()` and `to_file()` for conversion, and compares against HF tokenizers / tiktoken-rs for verification. + +**Tech Stack:** Rust, clap (CLI), hf-hub (HF Hub API), tokenizers (HF reference), tiktoken-rs (tiktoken reference) + +--- + +### Task 1: Add `build` feature flag and move dependencies + +**Files:** +- Modify: `crates/tokie/Cargo.toml` + +- [ ] **Step 1: Update Cargo.toml — add `build` feature and move deps** + +In `crates/tokie/Cargo.toml`, add the `build` feature and move `tokenizers` and `tiktoken-rs` from `[dev-dependencies]` to optional `[dependencies]`: + +```toml +[features] +default = [] +hf = ["hf-hub"] +build = ["hf", "dep:tokenizers", "dep:tiktoken-rs"] + +[dependencies] +# ... existing deps unchanged ... +tokenizers = { version = "0.22", features = ["http"], optional = true } +tiktoken-rs = { version = "0.6", optional = true } + +[dev-dependencies] +criterion = "0.5" +# tokenizers and tiktoken-rs removed from here +``` + +Also update examples that use `tokenizers` or `tiktoken-rs` to require the `build` feature. Find all examples that import these crates: + +- `convert_and_verify` — change `required-features` to `["build"]` +- `verify_embedding_models` — change `required-features` to `["build"]` +- `bench_vs_hf` — add `required-features = ["build"]` (currently missing from Cargo.toml) +- `check_models` — add `required-features = ["build"]` if it uses tokenizers + +Leave examples that only use `tokie` and `hf` unchanged. + +- [ ] **Step 2: Verify it compiles** + +Run: `cargo check -p tokie --features build` +Expected: Compiles with no errors. + +Run: `cargo check -p tokie` +Expected: Compiles without `tokenizers` or `tiktoken-rs`. + +- [ ] **Step 3: Commit** + +```bash +git add crates/tokie/Cargo.toml +git commit -m "feat: add build feature flag, move tokenizers/tiktoken-rs to optional deps" +``` + +--- + +### Task 2: Create `build` module with types and `convert` function + +**Files:** +- Create: `crates/tokie/src/build.rs` +- Modify: `crates/tokie/src/lib.rs:58` (add module declaration) + +- [ ] **Step 1: Add module declaration to lib.rs** + +In `crates/tokie/src/lib.rs`, after line 59 (`mod hub;`), add: + +```rust +#[cfg(feature = "build")] +pub mod build; +``` + +And after line 72 (`pub use hub::{FromPretrainedOptions, HubError};`), add: + +```rust +#[cfg(feature = "build")] +pub use build::{BuildError, ConvertResult, VerifyResult, Mismatch}; +``` + +- [ ] **Step 2: Create build.rs with types and convert function** + +Create `crates/tokie/src/build.rs`: + +```rust +//! Build tools for converting, verifying, and uploading .tkz tokenizers. +//! +//! This module is only available when the `build` feature is enabled: +//! ```toml +//! tokie = { version = "0.1", features = ["build"] } +//! ``` + +use std::path::Path; + +use hf_hub::Repo; + +use crate::hf::JsonLoadError; +use crate::serde::SerdeError; +use crate::Tokenizer; + +/// Result of a successful conversion. +#[derive(Debug)] +pub struct ConvertResult { + pub vocab_size: usize, + pub file_size_bytes: u64, +} + +/// A single verification mismatch. +#[derive(Debug)] +pub struct Mismatch { + pub text: String, + pub tokie_ids: Vec, + pub reference_ids: Vec, +} + +/// Result of a verification run. +#[derive(Debug)] +pub struct VerifyResult { + pub total: usize, + pub passed: usize, + pub failed: usize, + pub mismatches: Vec, +} + +/// Errors from build operations. +#[derive(Debug)] +pub enum BuildError { + Download(String), + LoadJson(JsonLoadError), + SaveTkz(SerdeError), + Verification { result: VerifyResult }, + Upload(String), + HubInit(String), +} + +impl std::fmt::Display for BuildError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BuildError::Download(e) => write!(f, "download failed: {e}"), + BuildError::LoadJson(e) => write!(f, "failed to load tokenizer.json: {e}"), + BuildError::SaveTkz(e) => write!(f, "failed to save .tkz: {e}"), + BuildError::Verification { result } => { + write!(f, "verification failed: {}/{} texts mismatched", result.failed, result.total) + } + BuildError::Upload(e) => write!(f, "upload failed: {e}"), + BuildError::HubInit(e) => write!(f, "HF Hub init failed: {e}"), + } + } +} + +impl std::error::Error for BuildError {} + +/// Convert a HuggingFace repo's tokenizer.json to .tkz format. +pub fn convert(repo_id: &str, output: &Path) -> Result { + let api = hf_hub::api::sync::ApiBuilder::new() + .build() + .map_err(|e| BuildError::HubInit(e.to_string()))?; + + let repo = Repo::model(repo_id.to_string()); + let repo_api = api.repo(repo); + + let json_path = repo_api + .get("tokenizer.json") + .map_err(|e| BuildError::Download(e.to_string()))?; + + let tokenizer = Tokenizer::from_json(&json_path).map_err(BuildError::LoadJson)?; + + if let Some(parent) = output.parent() { + std::fs::create_dir_all(parent).ok(); + } + tokenizer.to_file(output).map_err(BuildError::SaveTkz)?; + + let file_size_bytes = std::fs::metadata(output) + .map(|m| m.len()) + .unwrap_or(0); + + Ok(ConvertResult { + vocab_size: tokenizer.vocab_size(), + file_size_bytes, + }) +} +``` + +- [ ] **Step 3: Verify it compiles** + +Run: `cargo check -p tokie --features build` +Expected: Compiles with no errors. + +- [ ] **Step 4: Commit** + +```bash +git add crates/tokie/src/build.rs crates/tokie/src/lib.rs +git commit -m "feat: add build module with convert function" +``` + +--- + +### Task 3: Add `verify` function + +**Files:** +- Modify: `crates/tokie/src/build.rs` + +- [ ] **Step 1: Add the tiktoken repo-to-encoding mapping and test texts** + +Append to `crates/tokie/src/build.rs`, before the `convert` function: + +```rust +/// Repos that should be verified against tiktoken-rs instead of HF tokenizers. +fn tiktoken_encoding_name(repo_id: &str) -> Option<&'static str> { + match repo_id.to_ascii_lowercase().as_str() { + "xenova/gpt-4" => Some("cl100k_base"), + "xenova/gpt-4o" => Some("o200k_base"), + "xenova/text-davinci-003" => Some("p50k_base"), + _ => None, + } +} + +const VERIFY_TEXTS: &[&str] = &[ + "Hello, world!", + "The quick brown fox jumps over the lazy dog.", + "Machine learning models encode text into dense vector representations.", + "Tokenization is the process of splitting text into smaller units called tokens.", + "BGE, GTE, and E5 are popular embedding models for semantic search.", + "The 18th century was a time of great change.", + "user@example.com visited https://example.org/path?query=1", + "I can't believe it's not butter! Don't you think so?", +]; +``` + +- [ ] **Step 2: Add the verify function** + +Append to `crates/tokie/src/build.rs`: + +```rust +/// Verify a .tkz file against a reference tokenizer backend. +/// +/// Auto-detects whether to use tiktoken-rs or HF tokenizers based on repo_id. +pub fn verify(repo_id: &str, tkz_path: &Path) -> Result { + let tokie_tok = Tokenizer::from_file(tkz_path) + .map_err(|e| BuildError::SaveTkz(e))?; + + let mut mismatches = Vec::new(); + let mut passed = 0; + + if let Some(encoding_name) = tiktoken_encoding_name(repo_id) { + let tiktoken = match encoding_name { + "cl100k_base" => tiktoken_rs::cl100k_base(), + "o200k_base" => tiktoken_rs::o200k_base(), + "p50k_base" => tiktoken_rs::p50k_base(), + _ => unreachable!(), + } + .expect("failed to load tiktoken encoding"); + + for text in VERIFY_TEXTS { + let tokie_ids = tokie_tok.encode(text, false).ids; + let ref_ids: Vec = tiktoken.encode_with_special_tokens(text) + .into_iter() + .map(|id| id as u32) + .collect(); + + if tokie_ids == ref_ids { + passed += 1; + } else { + mismatches.push(Mismatch { + text: text.to_string(), + tokie_ids, + reference_ids: ref_ids, + }); + } + } + } else { + let hf_tok = tokenizers::Tokenizer::from_pretrained(repo_id, None) + .map_err(|e| BuildError::Download(format!("HF tokenizer load failed: {e}")))?; + + for text in VERIFY_TEXTS { + let tokie_ids = tokie_tok.encode(text, false).ids; + let hf_encoding = hf_tok + .encode(text.to_string(), false) + .map_err(|e| BuildError::Download(format!("HF encode failed: {e}")))?; + let ref_ids: Vec = hf_encoding.get_ids().to_vec(); + + if tokie_ids == ref_ids { + passed += 1; + } else { + mismatches.push(Mismatch { + text: text.to_string(), + tokie_ids, + reference_ids: ref_ids, + }); + } + } + } + + let total = VERIFY_TEXTS.len(); + let failed = mismatches.len(); + + Ok(VerifyResult { + total, + passed, + failed, + mismatches, + }) +} +``` + +- [ ] **Step 3: Verify it compiles** + +Run: `cargo check -p tokie --features build` +Expected: Compiles with no errors. + +- [ ] **Step 4: Commit** + +```bash +git add crates/tokie/src/build.rs +git commit -m "feat: add verify function with auto-detected reference backend" +``` + +--- + +### Task 4: Add `upload` function + +**Files:** +- Modify: `crates/tokie/src/build.rs` + +- [ ] **Step 1: Add the upload function** + +Append to `crates/tokie/src/build.rs`: + +```rust +/// Upload a .tkz file to the tokiers/ org on HuggingFace Hub. +/// +/// Uses the HF Hub HTTP API directly (hf-hub crate is download-only). +/// If `token` is None, falls back to HF_TOKEN env var. +pub fn upload(tkz_path: &Path, tokiers_name: &str, token: Option<&str>) -> Result<(), BuildError> { + if !tkz_path.exists() { + return Err(BuildError::Upload(format!("file not found: {}", tkz_path.display()))); + } + + let token_str = token + .map(|t| t.to_string()) + .or_else(|| std::env::var("HF_TOKEN").ok()) + .ok_or_else(|| BuildError::Upload( + "no HF token found — pass --token or set HF_TOKEN env var".to_string(), + ))?; + + let file_content = std::fs::read(tkz_path) + .map_err(|e| BuildError::Upload(format!("failed to read {}: {e}", tkz_path.display())))?; + + let repo_id = format!("tokiers/{tokiers_name}"); + let url = format!( + "https://huggingface.co/api/models/{}/upload/main/tokenizer.tkz", + repo_id + ); + + let response = ureq::put(&url) + .header("Authorization", &format!("Bearer {token_str}")) + .content_type("application/octet-stream") + .send_bytes(&file_content) + .map_err(|e| BuildError::Upload(format!("HTTP upload to {repo_id} failed: {e}")))?; + + if response.status() >= 400 { + return Err(BuildError::Upload(format!( + "upload returned HTTP {}: check your token has write access to tokiers/ org", + response.status() + ))); + } + + Ok(()) +} +``` + +Note: `ureq` is a transitive dependency via `hf-hub`. If it's not directly accessible, add `ureq = "2"` to `[dependencies]` in `crates/tokie/Cargo.toml` gated on the `build` feature. + +- [ ] **Step 3: Verify it compiles** + +Run: `cargo check -p tokie --features build` +Expected: Compiles with no errors. + +- [ ] **Step 4: Commit** + +```bash +git add crates/tokie/src/build.rs +git commit -m "feat: add upload function for tokiers/ HF org" +``` + +--- + +### Task 5: Create `tokie-cli` crate with `convert` command + +**Files:** +- Create: `crates/tokie-cli/Cargo.toml` +- Create: `crates/tokie-cli/src/main.rs` + +- [ ] **Step 1: Create Cargo.toml** + +Create `crates/tokie-cli/Cargo.toml`: + +```toml +[package] +name = "tokie-cli" +version = "0.0.1" +edition = "2024" +description = "CLI for building, verifying, and uploading tokie tokenizers" +license = "MIT OR Apache-2.0" +repository = "https://github.com/chonkie-inc/tokie" + +[[bin]] +name = "tokie" +path = "src/main.rs" + +[dependencies] +tokie = { path = "../tokie", features = ["hf", "build"] } +clap = { version = "4", features = ["derive"] } +``` + +- [ ] **Step 2: Create main.rs with convert command** + +Create `crates/tokie-cli/src/main.rs`: + +```rust +use std::path::PathBuf; + +use clap::{Parser, Subcommand}; + +#[derive(Parser)] +#[command(name = "tokie", about = "Build, verify, and upload tokie tokenizers")] +struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +enum Commands { + /// Convert a HuggingFace tokenizer.json to .tkz format + Convert { + /// HuggingFace repo ID (e.g., "openai-community/gpt2", "meta-llama/Llama-3.2-1B") + repo_id: String, + + /// Output path for the .tkz file + #[arg(short, long)] + output: PathBuf, + + /// Verify the converted .tkz against the reference tokenizer + #[arg(short, long)] + verify: bool, + + /// Upload to tokiers/ org on HuggingFace Hub (implies --verify) + #[arg(short, long)] + upload: bool, + + /// HuggingFace API token (falls back to HF_TOKEN env or cached token) + #[arg(long)] + token: Option, + }, + /// Verify a .tkz tokenizer against the reference backend + Verify { + /// HuggingFace repo ID + repo_id: String, + + /// Path to the .tkz file (if omitted, downloads from tokiers/) + #[arg(long)] + tkz: Option, + }, +} + +fn main() { + let cli = Cli::parse(); + + match cli.command { + Commands::Convert { repo_id, output, verify, upload, token } => { + cmd_convert(&repo_id, &output, verify || upload, upload, token.as_deref()); + } + Commands::Verify { repo_id, tkz } => { + cmd_verify(&repo_id, tkz.as_deref()); + } + } +} + +fn cmd_convert(repo_id: &str, output: &PathBuf, verify: bool, upload: bool, token: Option<&str>) { + print!("Converting {repo_id} ... "); + + match tokie::build::convert(repo_id, output) { + Ok(result) => { + println!( + "OK vocab={} size={:.1}KB", + result.vocab_size, + result.file_size_bytes as f64 / 1024.0 + ); + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + + if verify { + print!("Verifying ... "); + match tokie::build::verify(repo_id, output) { + Ok(result) => { + if result.failed == 0 { + println!("OK {}/{} texts pass", result.passed, result.total); + } else { + println!("FAILED {}/{} texts pass", result.passed, result.total); + for m in &result.mismatches { + println!(" MISMATCH: \"{}\"", m.text); + println!(" tokie: {:?}", &m.tokie_ids[..m.tokie_ids.len().min(15)]); + println!(" ref: {:?}", &m.reference_ids[..m.reference_ids.len().min(15)]); + } + std::process::exit(1); + } + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + } + + if upload { + let tokiers_name = repo_id.rsplit('/').next().unwrap_or(repo_id); + print!("Uploading to tokiers/{tokiers_name} ... "); + match tokie::build::upload(output, tokiers_name, token) { + Ok(()) => println!("OK"), + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + } +} + +fn cmd_verify(repo_id: &str, tkz: Option<&PathBuf>) { + let tkz_path; + let tkz_ref: &std::path::Path; + + if let Some(path) = tkz { + tkz_ref = path; + } else { + // Download from tokiers/ via from_pretrained + let name = repo_id.rsplit('/').next().unwrap_or(repo_id); + let tmp_path = format!("/tmp/tokie-verify-{name}.tkz"); + print!("Downloading tokiers/{name} ... "); + let api = hf_hub::api::sync::ApiBuilder::new().build().unwrap(); + let repo = hf_hub::Repo::model(format!("tokiers/{name}")); + match api.repo(repo).get("tokenizer.tkz") { + Ok(p) => { + std::fs::copy(&p, &tmp_path).unwrap(); + println!("OK"); + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } + tkz_path = PathBuf::from(tmp_path); + tkz_ref = &tkz_path; + } + + print!("Verifying {repo_id} ... "); + match tokie::build::verify(repo_id, tkz_ref) { + Ok(result) => { + if result.failed == 0 { + println!("OK {}/{} texts pass", result.passed, result.total); + } else { + println!("FAILED {}/{} texts pass", result.passed, result.total); + for m in &result.mismatches { + println!(" MISMATCH: \"{}\"", m.text); + println!(" tokie: {:?}", &m.tokie_ids[..m.tokie_ids.len().min(15)]); + println!(" ref: {:?}", &m.reference_ids[..m.reference_ids.len().min(15)]); + } + std::process::exit(1); + } + } + Err(e) => { + println!("FAILED: {e}"); + std::process::exit(1); + } + } +} +``` + +- [ ] **Step 3: Verify it compiles** + +Run: `cargo build -p tokie-cli` +Expected: Compiles and produces a `tokie` binary. + +- [ ] **Step 4: Commit** + +```bash +git add crates/tokie-cli/ +git commit -m "feat: add tokie-cli crate with convert and verify commands" +``` + +--- + +### Task 6: Integration test — convert and verify gpt2 + +**Files:** +- No new files — this is a manual smoke test + +- [ ] **Step 1: Test convert** + +Run: `cargo run -p tokie-cli -- convert openai-community/gpt2 -o /tmp/gpt2-test.tkz -v` + +Expected output (approximately): +``` +Converting openai-community/gpt2 ... OK vocab=50257 size=XXX.XKB +Verifying ... OK 8/8 texts pass +``` + +- [ ] **Step 2: Test verify standalone** + +Run: `cargo run -p tokie-cli -- verify openai-community/gpt2 --tkz /tmp/gpt2-test.tkz` + +Expected output: +``` +Verifying openai-community/gpt2 ... OK 8/8 texts pass +``` + +- [ ] **Step 3: Test verify with Hub download** + +Run: `cargo run -p tokie-cli -- verify openai-community/gpt2` + +Expected output: +``` +Downloading tokiers/gpt2 ... OK +Verifying openai-community/gpt2 ... OK 8/8 texts pass +``` + +- [ ] **Step 4: Test a BERT model** + +Run: `cargo run -p tokie-cli -- convert bert-base-uncased -o /tmp/bert-test.tkz -v` + +Expected: Converts and verifies successfully. + +- [ ] **Step 5: Test a tiktoken model (cl100k)** + +Run: `cargo run -p tokie-cli -- convert Xenova/gpt-4 -o /tmp/cl100k-test.tkz -v` + +Expected: Converts and verifies against tiktoken-rs. + +- [ ] **Step 6: Test error case — invalid repo** + +Run: `cargo run -p tokie-cli -- convert nonexistent/model -o /tmp/bad.tkz` + +Expected: Prints a clear error and exits with code 1. + +- [ ] **Step 7: Commit any fixes from testing** + +If any bugs were found and fixed during testing, commit them: + +```bash +git add -A +git commit -m "fix: address issues found during integration testing" +``` + +--- + +### Task 7: Update examples to use `build` feature + +**Files:** +- Modify: `crates/tokie/Cargo.toml` (example entries) + +- [ ] **Step 1: Update example required-features** + +In `crates/tokie/Cargo.toml`, update all examples that use `tokenizers` or `tiktoken-rs` to require the `build` feature: + +For these examples, change `required-features = ["hf"]` to `required-features = ["build"]`: +- `convert_and_verify` +- `verify_embedding_models` +- `tiktoken_compat` (if it uses tiktoken-rs directly) +- `regenerate_from_json` +- `regenerate_all_v11` + +For these examples that only need `hf`, leave unchanged: +- `from_pretrained` +- `test_gpt2` +- `test_llama4` +- `test_enwik8` +- `basic_usage` +- `padding_truncation` +- `batch_processing` +- `cross_encoder` + +Check each example file's imports to determine which group it belongs to. Run: + +```bash +grep -l 'use tokenizers\|use tiktoken_rs' crates/tokie/examples/*.rs +``` + +- [ ] **Step 2: Verify all examples compile** + +Run: `cargo build -p tokie --features build --examples` +Expected: All examples compile. + +Run: `cargo build -p tokie --features hf --examples` +Expected: Only non-build examples compile (build-gated ones are skipped). + +- [ ] **Step 3: Commit** + +```bash +git add crates/tokie/Cargo.toml +git commit -m "chore: update examples to use build feature flag" +``` diff --git a/docs/superpowers/specs/2026-04-21-build-module-and-cli-design.md b/docs/superpowers/specs/2026-04-21-build-module-and-cli-design.md new file mode 100644 index 0000000..4b24ab2 --- /dev/null +++ b/docs/superpowers/specs/2026-04-21-build-module-and-cli-design.md @@ -0,0 +1,111 @@ +# tokie build module + CLI design + +## Summary + +Add a `build` feature to the core `tokie` crate with functions for converting `tokenizer.json` to `.tkz`, verifying correctness against reference backends, and uploading to the `tokiers/` HF org. Add a `tokie-cli` crate that wires these functions to a CLI binary named `tokie`. + +## Motivation + +The build pipeline is currently scattered across 18 Rust examples and ad-hoc Python scripts, with the model list duplicated in three places (`regenerate_from_json.rs`, `upload_tkz_to_hf.py`, `hub.rs::tokiers_repo_name`). There is no single command to convert, verify, and upload a tokenizer. Moving the logic into library functions makes it reusable and testable, and the CLI gives a single entry point. + +## Design + +### `tokie::build` module (feature-gated) + +Gated behind `feature = "build"`, which implies `hf` and pulls in `tokenizers` (HuggingFace Rust crate) and `tiktoken-rs` as optional dependencies. + +Three public functions: + +#### `convert` + +```rust +pub fn convert(repo_id: &str, output: &Path) -> Result +``` + +Downloads `tokenizer.json` from the HF repo, loads it via `Tokenizer::from_json()`, saves as `.tkz` via `to_file()`. Returns `ConvertResult { vocab_size: u32, file_size_bytes: u64 }`. + +#### `verify` + +```rust +pub fn verify(repo_id: &str, tkz_path: &Path) -> Result +``` + +Loads the `.tkz` from disk and the reference tokenizer from HF. Auto-detects the reference backend: +- tiktoken-rs for repos that map to tiktoken encodings: `Xenova/gpt-4` (cl100k_base), `Xenova/gpt-4o` (o200k_base), `Xenova/text-davinci-003` (p50k_base). These repos don't have a `tokenizer.json` that HF tokenizers can load, so tiktoken-rs is the only viable reference. +- HuggingFace `tokenizers` crate for everything else (including `openai-community/gpt2`, which has a `tokenizer.json`) + +Encodes a standard set of test texts with both backends, compares token IDs. Returns `VerifyResult { total: usize, passed: usize, failed: usize, mismatches: Vec }` where `Mismatch` contains the text, tokie IDs, and reference IDs. + +#### `upload` + +```rust +pub fn upload(tkz_path: &Path, tokiers_name: &str, token: Option<&str>) -> Result<(), BuildError> +``` + +Uploads the `.tkz` file to `tokiers/{tokiers_name}/tokenizer.tkz` on HF Hub. Uses the provided token, or falls back to the HF token from environment/cache. Fails with a clear error if the user lacks write permissions to the `tokiers` org. + +#### Error type + +```rust +pub enum BuildError { + Download(String), + LoadJson(JsonLoadError), + SaveTkz(SerdeError), + Verification { result: VerifyResult }, + Upload(String), + MissingToken, +} +``` + +### `tokie-cli` crate + +A thin binary crate at `crates/tokie-cli/`. Depends on `tokie = { path = "../tokie", features = ["hf", "build"] }` and `clap`. + +Binary name is `tokie` (set via `[[bin]]` in Cargo.toml), so users install with `cargo install tokie-cli` and get the `tokie` command. + +#### CLI surface + +``` +tokie convert -o [-v] [-u] +tokie verify [--tkz ] +``` + +**`tokie convert -o `** +- Calls `build::convert(repo_id, output)` +- `-v` (verify): also calls `build::verify()` after conversion. Prints pass/fail summary. +- `-u` (upload): also calls `build::upload()`. Derives tokiers name from repo_id by stripping the org prefix (e.g., `meta-llama/Llama-3.2-1B` -> `Llama-3.2-1B`). `-u` implies `-v` — unverified tokenizers are never uploaded. +- Prints vocab size, file size, and verification results to stdout. + +**`tokie verify `** +- If `--tkz ` is provided, verifies that local file against the reference backend. +- If no `--tkz`, downloads the `.tkz` from `tokiers/{name}` via Hub and verifies it. +- Prints per-text pass/fail and a summary. + +### Dependency changes to `tokie` crate + +```toml +[features] +build = ["hf", "dep:tokenizers", "dep:tiktoken-rs"] + +[dependencies] +tokenizers = { version = "0.22", features = ["http"], optional = true } +tiktoken-rs = { version = "0.6", optional = true } +``` + +These move from `[dev-dependencies]` to optional `[dependencies]` gated on `build`. Existing examples that use them get `required-features = ["build"]`. + +### Files to change + +- `crates/tokie/Cargo.toml` — add `build` feature, move deps +- `crates/tokie/src/lib.rs` — add `#[cfg(feature = "build")] pub mod build;` +- `crates/tokie/src/build.rs` — new module with `convert`, `verify`, `upload`, types +- `crates/tokie-cli/Cargo.toml` — new crate +- `crates/tokie-cli/src/main.rs` — clap CLI wiring +- Existing examples — update `required-features` or remove if fully replaced + +### What this does NOT change + +- `from_pretrained` resolution logic in `hub.rs` (already updated separately to remove `tokenizer.json` fallback) +- The `tokiers_repo_name()` mapping in `hub.rs` (still used by inference path) +- The `upload_tkz_to_hf.py` script (can be deprecated later) +- The `regenerate_from_json.rs` batch example (can be replaced later with `tokie convert` in a loop)