From 51254f220656ca926fe1c1eeb86e3947f715c7a8 Mon Sep 17 00:00:00 2001 From: xrl1 Date: Sun, 12 Jul 2026 14:24:58 +0300 Subject: [PATCH 1/3] bench: add load_mem benchmark for tokenizer load memory dhat-based bench reporting peak live bytes and total churn of Tokenizer::from_file, complementing the latency-only criterion benches. --- .gitignore | 2 ++ tokenizers/Cargo.toml | 5 +++ tokenizers/benches/load_mem.rs | 60 ++++++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+) create mode 100644 tokenizers/benches/load_mem.rs diff --git a/.gitignore b/.gitignore index c76bbbb9e..e73e09d55 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,5 @@ pip-wheel-metadata *.code-workspace uv.lock + +dhat-heap.json diff --git a/tokenizers/Cargo.toml b/tokenizers/Cargo.toml index 716a1afa7..7203d8bbf 100644 --- a/tokenizers/Cargo.toml +++ b/tokenizers/Cargo.toml @@ -77,6 +77,10 @@ harness = false name = "ci_benchmark" harness = false +[[bench]] +name = "load_mem" +harness = false + [dependencies] tk-encode = { path = "tk-encode", version = "0.23.2-dev.0", default-features = false } tk-train = { path = "tk-train", version = "0.23.2-dev.0", default-features = false, optional = true } @@ -97,6 +101,7 @@ parity-aware-bpe = ["train", "tk-train/parity-aware-bpe"] criterion = "0.6" serde_json = "1.0" ahash = "0.8.11" +dhat = "0.3" tempfile = "3.10" assert_approx_eq = "1.1" tracing = "0.1" diff --git a/tokenizers/benches/load_mem.rs b/tokenizers/benches/load_mem.rs new file mode 100644 index 000000000..6b476cbf2 --- /dev/null +++ b/tokenizers/benches/load_mem.rs @@ -0,0 +1,60 @@ +//! Heap-allocation benchmark: profiles the heap (peak live bytes and total +//! churn) via `dhat` while `Tokenizer::from_file` runs, to measure load memory. +//! Set `LOAD_MEM_TOKENIZER` / `LOAD_MEM_ITERS` to override the file and loop count. + +use std::hint::black_box; + +use tokenizers::Tokenizer; + +#[global_allocator] +static ALLOC: dhat::Alloc = dhat::Alloc; + +fn mib(bytes: u64) -> f64 { + bytes as f64 / (1024.0 * 1024.0) +} + +fn main() { + let tokenizer_path = std::env::var("LOAD_MEM_TOKENIZER") + .unwrap_or_else(|_| "data/llama-3-tokenizer.json".to_string()); + let iters: usize = std::env::var("LOAD_MEM_ITERS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(1); + + if !std::path::Path::new(&tokenizer_path).exists() { + eprintln!("load_mem: tokenizer file not found: {tokenizer_path}"); + eprintln!(); + eprintln!("Pass a path to any `tokenizer.json` via LOAD_MEM_TOKENIZER:"); + eprintln!(" LOAD_MEM_TOKENIZER= cargo bench --bench load_mem"); + eprintln!(); + eprintln!("Or fetch the default llama-3 tokenizer (a large 128k-vocab BPE):"); + eprintln!(" make data # downloads data/llama-3-tokenizer.json"); + std::process::exit(2); + } + + let profiler = dhat::Profiler::new_heap(); + + for _ in 0..iters { + let tok = Tokenizer::from_file(&tokenizer_path) + .unwrap_or_else(|e| panic!("failed to load {}: {}", tokenizer_path, e)); + black_box(&tok); + drop(tok); + } + + let stats = dhat::HeapStats::get(); + drop(profiler); + + println!(); + println!("load_mem: {tokenizer_path} (iters={iters})"); + println!( + " peak : {:>8.2} MiB in {:>9} blocks (max live, t-gmax)", + mib(stats.max_bytes as u64), + stats.max_blocks + ); + println!( + " total : {:>8.2} MiB in {:>9} blocks (churn over region)", + mib(stats.total_bytes), + stats.total_blocks + ); + println!(" wrote dhat-heap.json (open in https://nnethercote.github.io/dh_view/dh_view.html)"); +} From 7213d35e179e13b3a0680c24d3fedbdf6268c80c Mon Sep 17 00:00:00 2001 From: xrl1 Date: Sun, 12 Jul 2026 14:27:27 +0300 Subject: [PATCH 2/3] perf(models): drop serde_json::Value round-trip when loading a model ModelWrapper::deserialize buffered the whole model into an owned Value and reparsed it, allocating the vocab twice. Read it as a borrowed RawValue, peek the `type` tag, and parse once into the concrete model. Halves load-time allocation for every model type, and preserves on-disk field order (a Value's map sorts keys), which the BPE change relies on. Enables the serde_json `raw_value` feature. --- tokenizers/Cargo.lock | 460 ++++++++++++++++--------- tokenizers/tk-encode/Cargo.toml | 2 +- tokenizers/tk-encode/src/models/mod.rs | 133 ++++--- 3 files changed, 384 insertions(+), 211 deletions(-) diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index 4e5a39e84..27d10193e 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "addr2line" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" +dependencies = [ + "gimli", +] + [[package]] name = "adler2" version = "2.0.1" @@ -84,6 +93,21 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", + "windows-link", +] + [[package]] name = "base64" version = "0.13.1" @@ -117,9 +141,9 @@ dependencies = [ "proc-macro2", "quote", "regex", - "rustc-hash", + "rustc-hash 2.1.3", "shlex 1.3.0", - "syn", + "syn 2.0.119", ] [[package]] @@ -154,9 +178,9 @@ checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" [[package]] name = "bitflags" -version = "2.13.0" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "bitmap_gen" @@ -201,9 +225,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "cast" @@ -222,9 +246,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.65" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" dependencies = [ "find-msvc-tools", "jobserver", @@ -249,9 +273,9 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" [[package]] name = "chacha20" @@ -313,18 +337,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.1" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.6.0" +version = "4.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" dependencies = [ "anstyle", "clap_lex", @@ -375,9 +399,9 @@ dependencies = [ [[package]] name = "console" -version = "0.16.3" +version = "0.16.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" dependencies = [ "encode_unicode", "libc", @@ -464,9 +488,9 @@ dependencies = [ [[package]] name = "crossbeam-deque" -version = "0.8.6" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" dependencies = [ "crossbeam-epoch", "crossbeam-utils", @@ -474,18 +498,18 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] [[package]] name = "crossbeam-utils" -version = "0.8.21" +version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" [[package]] name = "crunchy" @@ -501,9 +525,9 @@ checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" [[package]] name = "daachorse" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99251f238b74cd219a86fe6ea9328308ebb223fcbb5b8eb5aa400b847a41dded" +checksum = "5614204febbc33cc07a2806aa6440b904ac012b68eecc37f4493ea4a76455a3d" [[package]] name = "darling" @@ -526,7 +550,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.119", ] [[package]] @@ -537,7 +561,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ "darling_core", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -567,7 +591,7 @@ dependencies = [ "darling", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -577,7 +601,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" dependencies = [ "derive_builder_core", - "syn", + "syn 2.0.119", +] + +[[package]] +name = "dhat" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98cd11d84628e233de0ce467de10b8633f4ddaecafadefc86e13b84b8739b827" +dependencies = [ + "backtrace", + "lazy_static", + "mintex", + "parking_lot", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "thousands", ] [[package]] @@ -609,7 +649,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -673,9 +713,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" [[package]] name = "find-msvc-tools" @@ -722,53 +762,53 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "futures-channel" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" dependencies = [ "futures-core", ] [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" [[package]] name = "futures-io" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" [[package]] name = "futures-macro" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "futures-sink" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" [[package]] name = "futures-task" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" [[package]] name = "futures-util" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ "futures-core", "futures-io", @@ -823,16 +863,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 6.0.0", "rand_core 0.10.1", + "wasm-bindgen", ] +[[package]] +name = "gimli" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" + [[package]] name = "glob" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" [[package]] name = "half" @@ -873,7 +921,7 @@ dependencies = [ "indicatif 0.17.11", "libc", "log", - "rand 0.9.4", + "rand 0.9.5", "reqwest", "serde", "serde_json", @@ -894,9 +942,9 @@ dependencies = [ [[package]] name = "http-body" -version = "1.0.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" dependencies = [ "bytes", "http", @@ -904,9 +952,9 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" dependencies = [ "bytes", "futures-core", @@ -923,9 +971,9 @@ checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" [[package]] name = "hyper" -version = "1.10.1" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +checksum = "d22053281f852e11534f5198498373cbb59295120a20771d90f7ed1897490a72" dependencies = [ "atomic-waker", "bytes", @@ -954,7 +1002,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.8", + "webpki-roots 1.0.9", ] [[package]] @@ -1104,11 +1152,11 @@ dependencies = [ [[package]] name = "indicatif" -version = "0.18.5" +version = "0.18.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "993f007684f2e9727160da8b960ec161264703bfd1af084fd2e34d040c9a0dd4" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" dependencies = [ - "console 0.16.3", + "console 0.16.4", "portable-atomic", "unicode-width", "unit-prefix", @@ -1203,9 +1251,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.186" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libloading" @@ -1219,9 +1267,9 @@ dependencies = [ [[package]] name = "libredox" -version = "0.1.17" +version = "0.1.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f02ab6bace2054fb888a3c16f990117b579d14a3088e472d63c6011fa185c9d3" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" dependencies = [ "libc", ] @@ -1238,6 +1286,15 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -1266,7 +1323,7 @@ dependencies = [ "quote", "regex-syntax", "rustc_version", - "syn", + "syn 2.0.119", ] [[package]] @@ -1319,14 +1376,14 @@ checksum = "73acd151c6ce84a41d8d6fb0958d9a3d5a18d649ad5a85ad5b719439af8ad257" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "memchr" -version = "2.8.2" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "minimal-lexical" @@ -1344,11 +1401,17 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "mintex" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c505b3e17ed6b70a7ed2e67fbb2c560ee327353556120d6e72f5232b6880d536" + [[package]] name = "mio" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi", @@ -1374,7 +1437,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1417,6 +1480,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "object" +version = "0.37.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +dependencies = [ + "memchr", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1458,6 +1530,29 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "partition" version = "0.1.2" @@ -1540,9 +1635,9 @@ dependencies = [ [[package]] name = "portable-atomic" -version = "1.13.1" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" [[package]] name = "potential_utf" @@ -1575,23 +1670,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.119", ] [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] [[package]] name = "ptr_hash" -version = "2.0.1" +version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a847c2cc746ab2aeba36aad3e75fc417b47539603298c12d8373e388890aad3c" +checksum = "9f184d2c69ac0853853275df42e7160a7dc4f3248d93434002c28de27ed3f6d0" dependencies = [ "bitvec", "colored", @@ -1621,7 +1716,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.3", "rustls", "socket2", "thiserror", @@ -1632,16 +1727,17 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.15" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.2", + "rand_pcg", "ring", - "rustc-hash", + "rustc-hash 2.1.3", "rustls", "rustls-pki-types", "slab", @@ -1653,23 +1749,23 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ "cfg_aliases", "libc", "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "quote" -version = "1.0.46" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -1694,9 +1790,9 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.9.4" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.5", @@ -1748,6 +1844,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + [[package]] name = "rayon" version = "1.12.0" @@ -1794,6 +1899,15 @@ dependencies = [ "voracious_radix_sort", ] +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_users" version = "0.5.2" @@ -1807,9 +1921,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.4" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -1819,9 +1933,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.14" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" dependencies = [ "aho-corasick", "memchr", @@ -1872,7 +1986,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots 1.0.8", + "webpki-roots 1.0.9", ] [[package]] @@ -1889,11 +2003,23 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rustc-demangle" +version = "0.1.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b74b56ffa8bb2830709a538c2cbcae9aa062db0d2a42563bfb09bdaae44020eb" + [[package]] name = "rustc-hash" -version = "2.1.2" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc_version" @@ -1919,9 +2045,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.41" +version = "0.23.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "log", "once_cell", @@ -1934,9 +2060,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" dependencies = [ "web-time", "zeroize", @@ -1955,9 +2081,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "ryu" @@ -1974,6 +2100,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "semver" version = "1.0.28" @@ -1982,9 +2114,9 @@ checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -1992,29 +2124,29 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] [[package]] name = "serde_json" -version = "1.0.150" +version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ "itoa", "memchr", @@ -2058,9 +2190,9 @@ checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" [[package]] name = "simd-adler32" -version = "0.3.9" +version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" [[package]] name = "slab" @@ -2076,9 +2208,9 @@ checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" [[package]] name = "socket2" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", "windows-sys 0.61.2", @@ -2133,9 +2265,20 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.118" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" dependencies = [ "proc-macro2", "quote", @@ -2159,7 +2302,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2183,29 +2326,35 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] +[[package]] +name = "thousands" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820" + [[package]] name = "thread_local" -version = "1.1.9" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" dependencies = [ "cfg-if", ] @@ -2252,9 +2401,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" dependencies = [ "tinyvec_macros", ] @@ -2274,13 +2423,13 @@ dependencies = [ "atomsplit", "compact_str", "criterion 0.6.0", - "daachorse 3.0.2", + "daachorse 3.0.3", "dary_heap", "derive_builder", "fancy-regex 0.17.0", "getrandom 0.3.4", "hf-hub", - "indicatif 0.18.5", + "indicatif 0.18.6", "itertools 0.14.0", "log", "logos", @@ -2291,7 +2440,7 @@ dependencies = [ "paste", "pcre2", "ptr_hash", - "rand 0.9.4", + "rand 0.9.5", "rayon", "rayon-cond", "regex", @@ -2320,7 +2469,7 @@ dependencies = [ "dary_heap", "derive_builder", "esaxx-rs", - "indicatif 0.18.5", + "indicatif 0.18.6", "itertools 0.14.0", "log", "rayon", @@ -2344,14 +2493,14 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", - "indicatif 0.18.5", + "indicatif 0.18.6", "itertools 0.14.0", "log", "macro_rules_attribute", "monostate", "onig", "paste", - "rand 0.9.4", + "rand 0.9.5", "rayon", "rayon-cond", "regex", @@ -2372,6 +2521,7 @@ dependencies = [ "ahash", "assert_approx_eq", "criterion 0.6.0", + "dhat", "serde_json", "tempfile", "tk-encode", @@ -2382,9 +2532,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.3" +version = "1.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" dependencies = [ "bytes", "libc", @@ -2406,9 +2556,9 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.18" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" dependencies = [ "bytes", "futures-core", @@ -2481,7 +2631,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2719,7 +2869,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "wasm-bindgen-shared", ] @@ -2771,14 +2921,14 @@ version = "0.26.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" dependencies = [ - "webpki-roots 1.0.8", + "webpki-roots 1.0.9", ] [[package]] name = "webpki-roots" -version = "1.0.8" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" dependencies = [ "rustls-pki-types", ] @@ -3008,9 +3158,9 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" [[package]] name = "yada" @@ -3037,28 +3187,28 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.52" +version = "0.8.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.52" +version = "0.8.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -3078,7 +3228,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] @@ -3118,11 +3268,11 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index 1b1597e30..d20c329e0 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -33,7 +33,7 @@ regex = "1.10" rayon = "1.10" rayon-cond = "0.4" serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0" +serde_json = { version = "1.0", features = ["raw_value"] } unicode-normalization-alignments = "0.1" unicode_categories = "0.1" unicode-segmentation = "1.11" diff --git a/tokenizers/tk-encode/src/models/mod.rs b/tokenizers/tk-encode/src/models/mod.rs index 7fff60fc8..bbeab334e 100644 --- a/tokenizers/tk-encode/src/models/mod.rs +++ b/tokenizers/tk-encode/src/models/mod.rs @@ -9,7 +9,8 @@ use ahash::AHashMap; use std::collections::HashMap; use std::path::{Path, PathBuf}; -use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error}; +use serde_json::value::RawValue; use crate::models::bpe::BPE; use crate::models::unigram::Unigram; @@ -75,65 +76,44 @@ impl<'de> Deserialize<'de> for ModelWrapper { where D: Deserializer<'de>, { - #[derive(Deserialize)] - pub struct Tagged { - #[serde(rename = "type")] - variant: EnumType, - #[serde(flatten)] - rest: serde_json::Value, - } - #[derive(Deserialize)] - pub enum EnumType { - BPE, - WordPiece, - WordLevel, - Unigram, - } + let raw: &'de RawValue = Deserialize::deserialize(deserializer)?; + let model_json = raw.get(); #[derive(Deserialize)] - #[serde(untagged)] - pub enum ModelHelper { - Tagged(Tagged), - Legacy(serde_json::Value), - } - - #[derive(Deserialize)] - #[serde(untagged)] - pub enum ModelUntagged { - BPE(BPE), - // WordPiece must stay before WordLevel here for deserialization (for retrocompatibility - // with the versions not including the "type"), since WordLevel is a subset of WordPiece - WordPiece(WordPiece), - WordLevel(WordLevel), - Unigram(Unigram), + struct Tag { + #[serde(rename = "type")] + variant: Option, } + let tag: Tag = serde_json::from_str(model_json).map_err(D::Error::custom)?; - let helper = ModelHelper::deserialize(deserializer)?; - Ok(match helper { - ModelHelper::Tagged(model) => match model.variant { - EnumType::BPE => ModelWrapper::BPE( - serde_json::from_value(model.rest).map_err(serde::de::Error::custom)?, - ), - EnumType::WordPiece => ModelWrapper::WordPiece( - serde_json::from_value(model.rest).map_err(serde::de::Error::custom)?, - ), - EnumType::WordLevel => ModelWrapper::WordLevel( - serde_json::from_value(model.rest).map_err(serde::de::Error::custom)?, - ), - EnumType::Unigram => ModelWrapper::Unigram( - serde_json::from_value(model.rest).map_err(serde::de::Error::custom)?, - ), - }, - ModelHelper::Legacy(value) => { - let untagged = serde_json::from_value(value).map_err(serde::de::Error::custom)?; - match untagged { - ModelUntagged::BPE(bpe) => ModelWrapper::BPE(bpe), - ModelUntagged::WordPiece(bpe) => ModelWrapper::WordPiece(bpe), - ModelUntagged::WordLevel(bpe) => ModelWrapper::WordLevel(bpe), - ModelUntagged::Unigram(bpe) => ModelWrapper::Unigram(bpe), + match tag.variant.as_deref() { + Some("BPE") => Ok(ModelWrapper::BPE( + serde_json::from_str(model_json).map_err(D::Error::custom)?, + )), + Some("WordPiece") => Ok(ModelWrapper::WordPiece( + serde_json::from_str(model_json).map_err(D::Error::custom)?, + )), + Some("WordLevel") => Ok(ModelWrapper::WordLevel( + serde_json::from_str(model_json).map_err(D::Error::custom)?, + )), + Some("Unigram") => Ok(ModelWrapper::Unigram( + serde_json::from_str(model_json).map_err(D::Error::custom)?, + )), + Some(other) => Err(D::Error::custom(format!("Unknown model type `{other}`"))), + None => { + if let Ok(m) = serde_json::from_str::(model_json) { + Ok(ModelWrapper::BPE(m)) + } else if let Ok(m) = serde_json::from_str::(model_json) { + Ok(ModelWrapper::WordPiece(m)) + } else if let Ok(m) = serde_json::from_str::(model_json) { + Ok(ModelWrapper::WordLevel(m)) + } else if let Ok(m) = serde_json::from_str::(model_json) { + Ok(ModelWrapper::Unigram(m)) + } else { + Err(D::Error::custom("Model is not a known variant")) } } - }) + } } } @@ -271,8 +251,51 @@ mod tests { let reconstructed: std::result::Result = serde_json::from_str(invalid); match reconstructed { - Err(err) => assert_eq!(err.to_string(), "Merges text file invalid at line 1"), + Err(err) => assert!( + err.to_string().starts_with("Merges text file invalid at line 1"), + "unexpected error: {}", + err + ), _ => panic!("Expected an error here"), } } + + // The `type` tag must dispatch to the right variant for every model, and an + // unknown tag must error (RawValue-based deserialization). + #[test] + fn model_wrapper_dispatches_on_type() { + use crate::models::unigram::Unigram; + use crate::models::wordlevel::WordLevel; + use crate::models::wordpiece::WordPiece; + + // Round-trip each model's default through `ModelWrapper` and check the + // resolved variant. + let bpe = serde_json::to_string(&BPE::default()).unwrap(); + assert!(matches!( + serde_json::from_str(&bpe).unwrap(), + ModelWrapper::BPE(_) + )); + + let wordpiece = serde_json::to_string(&WordPiece::default()).unwrap(); + assert!(matches!( + serde_json::from_str(&wordpiece).unwrap(), + ModelWrapper::WordPiece(_) + )); + + let wordlevel = serde_json::to_string(&WordLevel::default()).unwrap(); + assert!(matches!( + serde_json::from_str(&wordlevel).unwrap(), + ModelWrapper::WordLevel(_) + )); + + let unigram = serde_json::to_string(&Unigram::default()).unwrap(); + assert!(matches!( + serde_json::from_str(&unigram).unwrap(), + ModelWrapper::Unigram(_) + )); + + let unknown = r#"{"type":"NotAModel","vocab":{"a":0}}"#; + let err = serde_json::from_str::(unknown).unwrap_err(); + assert!(err.to_string().starts_with("Unknown model type")); + } } From fa36a76b68f4b5f4ba594293d1253675c4278185 Mon Sep 17 00:00:00 2001 From: xrl1 Date: Sun, 12 Jul 2026 14:28:41 +0300 Subject: [PATCH 3/3] perf(bpe): resolve merges to ids while deserializing, not after BPE loading buffered every merge into an owned Vec<(String, String)> and re-walked it in build(). BPEVisitor now resolves each merge to ids as it is parsed, borrowing tokens via Cow so only escaped tokens allocate. Merge resolution is shared with build() so both paths stay identical; when merges precede vocab it falls back to the old buffer-then-resolve. Reword Error::BadMerges to "Invalid merge rule #N" (the position is a JSON array index, not a file line). With the previous commit, loading a 128k-vocab BPE drops from 172 to 44 MiB peak. Output is byte-identical (round-trip and encode parity verified). --- tokenizers/tk-encode/src/models/bpe/mod.rs | 5 +- tokenizers/tk-encode/src/models/bpe/model.rs | 123 +++++++----- .../tk-encode/src/models/bpe/serialization.rs | 183 ++++++++++++++++-- tokenizers/tk-encode/src/models/mod.rs | 2 +- 4 files changed, 244 insertions(+), 69 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/mod.rs b/tokenizers/tk-encode/src/models/bpe/mod.rs index 6e1cb2da9..6c2d990f8 100644 --- a/tokenizers/tk-encode/src/models/bpe/mod.rs +++ b/tokenizers/tk-encode/src/models/bpe/mod.rs @@ -19,9 +19,8 @@ pub enum Error { /// When the vocab.json file is in the wrong format #[error("Bad vocabulary json file")] BadVocabulary, - /// When the merges.txt file is in the wrong format. This error holds the line - /// number of the line that caused the error. - #[error("Merges text file invalid at line {0}")] + /// A merge rule is in the wrong format. Holds its 1-based position. + #[error("Invalid merge rule #{0}")] BadMerges(usize), /// If a token found in merges, is not in the vocab #[error("Token `{0}` out of vocabulary")] diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 4c1fc07f5..79390dc54 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -95,10 +95,15 @@ thread_local! { } pub type Merges = Vec<(String, String)>; +pub(super) enum MergeInput { + Raw(Merges), + Resolved(MergeMap), +} + struct Config { files: Option<(String, String)>, vocab: Vocab, - merges: Merges, + merges: MergeInput, cache_capacity: usize, dropout: Option, unk_token: Option, @@ -120,7 +125,7 @@ impl Default for BpeBuilder { config: Config { files: None, vocab: AHashMap::new(), - merges: vec![], + merges: MergeInput::Raw(vec![]), cache_capacity: DEFAULT_CACHE_CAPACITY, dropout: None, unk_token: None, @@ -150,11 +155,18 @@ impl BpeBuilder { /// Set the vocab (token -> ID) and merges mappings. #[must_use] pub fn vocab_and_merges>>( - mut self, + self, vocab: V, merges: Merges, ) -> Self { - self.config.vocab = vocab.into(); + self.vocab_and_merge_input(vocab.into(), MergeInput::Raw(merges)) + } + + /// Set the vocab and merges, where merges may be raw token strings (resolved + /// to ids in `build`) or already resolved by the deserializer. + #[must_use] + pub(super) fn vocab_and_merge_input(mut self, vocab: Vocab, merges: MergeInput) -> Self { + self.config.vocab = vocab; self.config.merges = merges; self } @@ -227,7 +239,7 @@ impl BpeBuilder { if let Some((vocab, merges)) = self.config.files { let (v, m) = BPE::read_file(&vocab, &merges)?; self.config.vocab = v; - self.config.merges = m; + self.config.merges = MergeInput::Raw(m); } let mut max_len = 0; @@ -247,31 +259,19 @@ impl BpeBuilder { } else { 0 }; - let mut buffer: Vec = vec![0; max_len]; - let merge_map: MergeMap = self - .config - .merges - .into_iter() - .enumerate() - .map(|(i, (a, b))| -> Result<(Pair, (u32, u32))> { - let a_id = vocab - .get(&a) - .ok_or_else(|| Error::MergeTokenOutOfVocabulary(a.to_owned()))?; - let b_id = vocab - .get(&b) - .ok_or_else(|| Error::MergeTokenOutOfVocabulary(b.to_owned()))?; - buffer[0..a.len()].copy_from_slice(a.as_bytes()); - let b_len = b.len() - prefix_len; - let merge_len = a.len() + b_len; - buffer[a.len()..merge_len].copy_from_slice(&b.as_bytes()[prefix_len..]); - // SAFETY: buffer contains a concatenation of two valid UTF-8 strings, so it is itself valid UTF-8, even considering prefix_len - let new_token = unsafe { from_utf8_unchecked(&buffer[..merge_len]) }; - let new_id = vocab - .get(new_token) - .ok_or_else(|| Error::MergeTokenOutOfVocabulary(new_token.to_owned()))?; - Ok(((*a_id, *b_id), (i as u32, *new_id))) - }) - .collect::>()?; + let merge_map: MergeMap = match self.config.merges { + MergeInput::Resolved(resolved) => resolved, + MergeInput::Raw(merges) => { + let mut buffer: Vec = vec![0; max_len]; + merges + .into_iter() + .enumerate() + .map(|(i, (a, b))| { + resolve_merge(&vocab, &mut buffer, prefix_len, i as u32, &a, &b) + }) + .collect::>()? + } + }; // merges.insert(pair, (rank as u32, *new_id)); @@ -370,24 +370,56 @@ impl Clone for BPE { } } +/// Resolves a merge rule `(a, b)` into `((a_id, b_id), (rank, new_id))`. +pub(super) fn resolve_merge( + vocab: &Vocab, + buffer: &mut [u8], + prefix_len: usize, + rank: u32, + a: &str, + b: &str, +) -> Result<(Pair, (u32, u32))> { + let a_id = vocab + .get(a) + .ok_or_else(|| Error::MergeTokenOutOfVocabulary(a.to_owned()))?; + let b_id = vocab + .get(b) + .ok_or_else(|| Error::MergeTokenOutOfVocabulary(b.to_owned()))?; + buffer[0..a.len()].copy_from_slice(a.as_bytes()); + let b_len = b.len() - prefix_len; + let merge_len = a.len() + b_len; + buffer[a.len()..merge_len].copy_from_slice(&b.as_bytes()[prefix_len..]); + // SAFETY: buffer contains a concatenation of two valid UTF-8 strings, so it is itself valid UTF-8, even considering prefix_len + let new_token = unsafe { from_utf8_unchecked(&buffer[..merge_len]) }; + let new_id = vocab + .get(new_token) + .ok_or_else(|| Error::MergeTokenOutOfVocabulary(new_token.to_owned()))?; + Ok(((*a_id, *b_id), (rank, *new_id))) +} + +/// Parses one legacy `"{a} {b}"` merge line. Returns `Ok(None)` for a +/// `#version` header, which is skipped without consuming a rank. `rank` is the +/// 1-based position of the rule, used in the error message. +pub(super) fn parse_legacy_merge(line: &str, rank: usize) -> Result> { + if line.starts_with("#version") { + return Ok(None); + } + let parts = line.split(' ').collect::>(); + if parts.len() != 2 { + return Err(Error::BadMerges(rank).into()); + } + Ok(Some((parts[0].to_string(), parts[1].to_string()))) +} + /// Converts the merges strings (for example from `merges.txt` file) with the format /// "{pair_a} {pair_b}" into the format expected by the BPE struct -pub(crate) fn convert_merges_to_hashmap>( - iter: I, - _vocab: &Vocab, -) -> Result { +pub(crate) fn convert_merges_to_hashmap>(iter: I) -> Result { let mut merges = vec![]; - - let lines = iter.filter(|l| !l.starts_with("#version")); - for (rank, line) in lines.enumerate() { - let parts = line.split(' ').collect::>(); - if parts.len() != 2 { - return Err(Error::BadMerges(rank + 1).into()); + for line in iter { + if let Some(pair) = parse_legacy_merge(&line, merges.len() + 1)? { + merges.push(pair); } - - merges.push((parts[0].to_string(), parts[1].to_string())); } - Ok(merges) } @@ -435,9 +467,8 @@ impl BPE { // Read merges file let merge_file = File::open(merges)?; let merge_file = BufReader::new(merge_file); - let merges = ResultShunt::process(merge_file.lines(), |iter| { - convert_merges_to_hashmap(iter, &vocab) - })??; + let merges = + ResultShunt::process(merge_file.lines(), |iter| convert_merges_to_hashmap(iter))??; Ok((vocab, merges)) } diff --git a/tokenizers/tk-encode/src/models/bpe/serialization.rs b/tokenizers/tk-encode/src/models/bpe/serialization.rs index eb49922a4..fee4c0780 100644 --- a/tokenizers/tk-encode/src/models/bpe/serialization.rs +++ b/tokenizers/tk-encode/src/models/bpe/serialization.rs @@ -1,10 +1,14 @@ -use super::{super::OrderedVocabIter, BPE, BpeBuilder, Pair, convert_merges_to_hashmap}; +use super::{ + super::OrderedVocabIter, BPE, BpeBuilder, MergeInput, MergeMap, Merges, Pair, + convert_merges_to_hashmap, parse_legacy_merge, resolve_merge, +}; use ahash::AHashMap; use serde::{ Deserialize, Deserializer, Serialize, Serializer, de::{Error, MapAccess, Visitor}, ser::SerializeStruct, }; +use std::borrow::Cow; impl Serialize for BPE { fn serialize(&self, serializer: S) -> Result @@ -88,14 +92,10 @@ impl<'de> Visitor<'de> for BPEVisitor { { let mut builder = BpeBuilder::new(); let mut vocab: Option> = None; + // Needed to resolve merges on the fly; `None` until the field is seen. + let mut prefix_len: Option = None; + let mut merges: Option = None; - #[derive(Debug, Deserialize)] - #[serde(untagged)] - enum MergeType { - Tuple(Vec<(String, String)>), - Legacy(Vec), - } - let mut merges: Option = None; while let Some(key) = map.next_key::()? { match key.as_ref() { "dropout" => { @@ -109,7 +109,9 @@ impl<'de> Visitor<'de> for BPEVisitor { } } "continuing_subword_prefix" => { - if let Some(prefix) = map.next_value()? { + let prefix: Option = map.next_value()?; + prefix_len = Some(prefix.as_ref().map_or(0, String::len)); + if let Some(prefix) = prefix { builder = builder.continuing_subword_prefix(prefix); } } @@ -134,7 +136,21 @@ impl<'de> Visitor<'de> for BPEVisitor { } } "vocab" => vocab = Some(map.next_value()?), - "merges" => merges = Some(map.next_value()?), + "merges" => { + merges = Some(match (&vocab, prefix_len) { + // Fast path: resolve merges to ids as they are parsed. + (Some(vocab), Some(prefix_len)) => { + let max_len = vocab.keys().map(|k| k.len()).max().unwrap_or(0); + MergeInput::Resolved(map.next_value_seed(MergesResolver { + vocab, + prefix_len, + max_len, + })?) + } + // vocab/prefix not seen yet: buffer raw, resolve in `build`. + _ => MergeInput::Raw(map.next_value::()?.into_pairs()?), + }); + } "type" => match map.next_value()? { "BPE" => {} u => { @@ -147,18 +163,92 @@ impl<'de> Visitor<'de> for BPEVisitor { _ => {} } } - if let (Some(vocab), Some(merges)) = (vocab, merges) { - let merges = match merges { - MergeType::Tuple(merges) => merges, - MergeType::Legacy(merges) => { - convert_merges_to_hashmap(merges.into_iter(), &vocab).map_err(Error::custom)? + match (vocab, merges) { + (Some(vocab), Some(merges)) => builder + .vocab_and_merge_input(vocab, merges) + .build() + .map_err(Error::custom), + _ => Err(Error::custom("Missing vocab/merges")), + } + } +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum RawMergeType { + Tuple(Vec<(String, String)>), + Legacy(Vec), +} + +impl RawMergeType { + fn into_pairs(self) -> Result { + match self { + RawMergeType::Tuple(pairs) => Ok(pairs), + RawMergeType::Legacy(lines) => { + convert_merges_to_hashmap(lines.into_iter()).map_err(E::custom) + } + } + } +} + +/// Deserializes the `merges` array straight into a resolved `MergeMap`, so the +/// merge tokens are never collected into an owned `Vec<(String, String)>`. +struct MergesResolver<'v> { + vocab: &'v AHashMap, + prefix_len: usize, + max_len: usize, +} + +impl<'de, 'v> serde::de::DeserializeSeed<'de> for MergesResolver<'v> { + type Value = MergeMap; + + fn deserialize(self, deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_seq(self) + } +} + +impl<'de, 'v> Visitor<'de> for MergesResolver<'v> { + type Value = MergeMap; + + fn expecting(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(fmt, "a sequence of BPE merge rules") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: serde::de::SeqAccess<'de>, + { + #[derive(Deserialize)] + #[serde(untagged)] + enum MergeElem<'a> { + Pair(#[serde(borrow)] (Cow<'a, str>, Cow<'a, str>)), + Legacy(#[serde(borrow)] Cow<'a, str>), + } + + let mut buffer: Vec = vec![0; self.max_len]; + let mut merge_map = + MergeMap::with_capacity_and_hasher(seq.size_hint().unwrap_or(0), Default::default()); + let mut rank: u32 = 0; + while let Some(elem) = seq.next_element::()? { + let (a, b): (Cow, Cow) = match elem { + MergeElem::Pair((a, b)) => (a, b), + MergeElem::Legacy(line) => { + match parse_legacy_merge(&line, rank as usize + 1).map_err(Error::custom)? { + Some((a, b)) => (Cow::Owned(a), Cow::Owned(b)), + None => continue, + } } }; - builder = builder.vocab_and_merges(vocab, merges); - Ok(builder.build().map_err(Error::custom)?) - } else { - Err(Error::custom("Missing vocab/merges")) + let (pair, value) = + resolve_merge(self.vocab, &mut buffer, self.prefix_len, rank, &a, &b) + .map_err(Error::custom)?; + merge_map.insert(pair, value); + rank += 1; } + Ok(merge_map) } } @@ -242,4 +332,59 @@ mod test { let bpe_string = r#"{"type":"BPE","dropout":null,"unk_token":"","continuing_subword_prefix":null,"end_of_word_suffix":null,"fuse_unk":false,"byte_fallback":false,"vocab":{"":0,"a":1,"b":2},"merges":[]}"#; assert_eq!(serde_json::from_str::(bpe_string).unwrap(), bpe); } + + // Deserialization must work independent of JSON field order. + #[test] + fn test_deserialize_field_order_independent() { + use itertools::Itertools; + + let fields = [ + r#""type":"BPE""#, + r#""dropout":null"#, + r#""unk_token":"[UNK]""#, + r###""continuing_subword_prefix":"##""###, + r#""end_of_word_suffix":null"#, + r#""fuse_unk":false"#, + r#""byte_fallback":false"#, + r#""ignore_merges":true"#, + r###""vocab":{"[UNK]":0,"a":1,"##b":2,"##c":3,"ab":4,"abc":5}"###, + r###""merges":[["a","##b"],["ab","##c"]]"###, + ]; + + let expected: BPE = + serde_json::from_str(&format!("{{{}}}", fields.iter().join(","))).unwrap(); + + for i in 0..fields.len() { + let mut rotated = fields.iter().cycle().skip(i).take(fields.len()); + let json = format!("{{{}}}", rotated.join(",")); + let bpe: BPE = serde_json::from_str(&json) + .unwrap_or_else(|e| panic!("failed to deserialize {}: {}", json, e)); + assert_eq!(bpe, expected, "field order {} changed the model", json); + } + } + + // Legacy "a b" merges must parse the same whether merges come before or + // after vocab (fallback vs fast path). + #[test] + fn test_deserialize_legacy_merges_both_orders() { + let vocab = r#""vocab":{"a":0,"b":1,"ab":2}"#; + let merges = r#""merges":["a b"]"#; + let vocab_first: BPE = + serde_json::from_str(&format!(r#"{{"type":"BPE",{vocab},{merges}}}"#)).unwrap(); + let merges_first: BPE = + serde_json::from_str(&format!(r#"{{"type":"BPE",{merges},{vocab}}}"#)).unwrap(); + assert_eq!(vocab_first, merges_first); + } + + // A merge referencing a token outside the vocab is rejected. + #[test] + fn test_deserialize_merge_out_of_vocab() { + let json = r#"{"type":"BPE","vocab":{"a":0,"b":1,"ab":2},"merges":[["a","zzz"]]}"#; + let err = serde_json::from_str::(json).unwrap_err(); + assert!( + err.to_string().starts_with("Token `zzz` out of vocabulary"), + "unexpected error: {}", + err + ); + } } diff --git a/tokenizers/tk-encode/src/models/mod.rs b/tokenizers/tk-encode/src/models/mod.rs index bbeab334e..196d714d2 100644 --- a/tokenizers/tk-encode/src/models/mod.rs +++ b/tokenizers/tk-encode/src/models/mod.rs @@ -252,7 +252,7 @@ mod tests { serde_json::from_str(invalid); match reconstructed { Err(err) => assert!( - err.to_string().starts_with("Merges text file invalid at line 1"), + err.to_string().starts_with("Invalid merge rule #1"), "unexpected error: {}", err ),