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.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/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)"); +} 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/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 7fff60fc8..196d714d2 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("Invalid merge rule #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")); + } }