diff --git a/bindings/python/Cargo.lock b/bindings/python/Cargo.lock index acbc8bf8d..791ea4921 100644 --- a/bindings/python/Cargo.lock +++ b/bindings/python/Cargo.lock @@ -1597,6 +1597,7 @@ dependencies = [ "getrandom 0.3.4", "indicatif", "itertools 0.14.0", + "libc", "log", "macro_rules_attribute", "memchr", diff --git a/bindings/python/py_src/tokenizers/__init__.py b/bindings/python/py_src/tokenizers/__init__.py index 07e1e85be..3756296c9 100644 --- a/bindings/python/py_src/tokenizers/__init__.py +++ b/bindings/python/py_src/tokenizers/__init__.py @@ -88,8 +88,10 @@ class SplitDelimiterBehavior(Enum): from .tokenizers import ( AddedToken, + EncodeHandle, Encoding, NormalizedString, + PipelineTokenizer, PreTokenizedString, Regex, Token, diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index b670e95c2..9d24dae40 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -28,6 +28,7 @@ mod processors; mod token; mod tokenizer; mod trainers; +mod pipeline; mod utils; use pyo3::prelude::*; @@ -60,6 +61,10 @@ pub mod tokenizers { #[pymodule_export] pub use super::tokenizer::PyTokenizer; #[pymodule_export] + pub use super::pipeline::PyPipelineTokenizer; + #[pymodule_export] + pub use super::pipeline::PyEncodeHandle; + #[pymodule_export] pub use super::utils::PyNormalizedString; #[pymodule_export] pub use super::utils::PyPreTokenizedString; diff --git a/bindings/python/src/pipeline.rs b/bindings/python/src/pipeline.rs new file mode 100644 index 000000000..67488cdfd --- /dev/null +++ b/bindings/python/src/pipeline.rs @@ -0,0 +1,171 @@ +//! Pipeline python bindings + +use std::sync::Mutex; + +use pyo3::exceptions; +use pyo3::prelude::*; +use pyo3::types::PyString; + +use tk::pipeline::{EncodeHandle, Inputs, IntoInputs, PipelineToken, PipelineTokenizer}; + +use crate::error::ToPyResult; +use crate::tokenizer::PyTokenizer; + +fn ids(tokens: Vec) -> Vec { + tokens.iter().map(|t| t.id).collect() +} + +struct PyStrView { + ptr: *const u8, + len: usize, +} + +/// Keep alive mechanism for zero-copy access to utf8 python strings +struct PyStrBatch { + /// the keep alive ref + _owners: Vec>, + /// (ptr, len) + views: Vec, +} + +// SAFETY: the raw pointers reference CPython's cached UTF-8 buffers, which are +// immutable and outlive the owning `Py`s in `_owners`; the batch +// never mutates or frees them, and pool workers only read through `get`. +unsafe impl Send for PyStrBatch {} +unsafe impl Sync for PyStrBatch {} + +impl PyStrBatch { + fn new(strings: Vec>) -> PyResult { + let mut owners = Vec::with_capacity(strings.len()); + let mut views = Vec::with_capacity(strings.len()); + for s in strings { + let view = s.to_str()?; + views.push(PyStrView { + ptr: view.as_ptr(), + len: view.len(), + }); + owners.push(s.unbind()); + } + Ok(Self { + _owners: owners, + views, + }) + } +} + +impl Inputs for PyStrBatch { + fn len(&self) -> usize { + self.views.len() + } + + fn get(&self, i: usize) -> &str { + let (ptr, len) = (self.views[i].ptr, self.views[i].len); + // SAFETY: ptr is kept alive by PyStrBatch::_owners, and the underlying buffer is immutable UTF-8 + unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, len)) } + } +} + +impl IntoInputs for PyStrBatch { + type Inputs = PyStrBatch; + fn into_inputs(self) -> PyStrBatch { + self + } +} + +#[pyclass(module = "tokenizers", name = "PipelineTokenizer")] +pub struct PyPipelineTokenizer { + pipeline: PipelineTokenizer, +} + +#[pymethods] +impl PyPipelineTokenizer { + #[staticmethod] + fn from_tokenizer(tokenizer: &PyTokenizer) -> PyResult { + let json = { + let guard = tokenizer.read_inner()?; + serde_json::to_string(&*guard) + .map_err(|e| exceptions::PyException::new_err(format!("{e}")))? + }; + let tok: tk::Tokenizer = serde_json::from_str(&json) + .map_err(|e| exceptions::PyException::new_err(format!("{e}")))?; + let pipeline = PyResult::from(ToPyResult(PipelineTokenizer::try_from(&tok)))?; + Ok(Self { pipeline }) + } + + #[staticmethod] + fn from_file(path: &str) -> PyResult { + let tok = PyResult::from(ToPyResult(tk::Tokenizer::from_file(path)))?; + let pipeline = PyResult::from(ToPyResult(PipelineTokenizer::try_from(&tok)))?; + Ok(Self { pipeline }) + } + + /// Encode a batch of `str`, returning an [`PyEncodeHandle`] as soon as possible: + /// pool workers encode in the background while you hold the job, `wait()` + /// for everything (returned in input order), or iterate `(index, ids)` as + /// each input completes. Zero-copy: the job reads the Python strings' UTF-8 + /// buffers in place. + fn encode_batch( + &self, + py: Python<'_>, + input: Vec>, + ) -> PyResult { + let batch = PyStrBatch::new(input)?; + // Below the cost gate the job is computed inline — release the GIL. + let job = py.detach(|| self.pipeline.encode(batch)); + Ok(PyEncodeHandle { + job: Mutex::new(Some(job)), + }) + } +} + +/// An in-flight (or completed) batch encode. `wait()` blocks for all inputs and +/// returns their ids lists in input order; iterating yields `(index, ids)` for +/// each input as it completes (completion order, not input order — use `index` +/// to place it). The consuming thread *assists* the pool while it waits (all +/// with the GIL released). Dropping the job cancels unclaimed work. +#[pyclass(module = "tokenizers", name = "EncodeHandle")] +pub struct PyEncodeHandle { + job: Mutex>, +} + +impl PyEncodeHandle { + fn take(&self) -> PyResult { + self.job + .lock() + .unwrap_or_else(|e| e.into_inner()) + .take() + .ok_or_else(|| exceptions::PyException::new_err("EncodeHandle already consumed")) + } +} + +#[pymethods] +impl PyEncodeHandle { + /// Block until every input is encoded, returning ids lists in input order + fn wait(&self, py: Python<'_>) -> PyResult>> { + let job = self.take()?; + let res = py.detach(|| job.wait_for_completion()); + let out = PyResult::from(ToPyResult(res))?; + Ok(out.into_iter().map(ids).collect()) + } + + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + /// Yield `(input index, ids)` for each input as it finishes — completion + /// order, not input order. `index` is the position in the batch passed to + /// `encode_batch` (always `0` for a single input). + fn __next__(&self, py: Python<'_>) -> PyResult)>> { + let next = py.detach(|| { + self.job + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_mut() + .and_then(|j| j.next()) + }); + match next { + None => Ok(None), + Some((seq, res)) => Ok(Some((seq, ids(PyResult::from(ToPyResult(res))?)))), + } + } +} diff --git a/tokenizers/Cargo.lock b/tokenizers/Cargo.lock index 4e5a39e84..e2ef61613 100644 --- a/tokenizers/Cargo.lock +++ b/tokenizers/Cargo.lock @@ -119,7 +119,7 @@ dependencies = [ "regex", "rustc-hash", "shlex 1.3.0", - "syn", + "syn 2.0.119", ] [[package]] @@ -154,9 +154,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 +201,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 +222,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 +249,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 +313,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.1" +version = "4.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +checksum = "dd059f9da4f5c36b3787f65d38ccaab1cc315f07b01f89abc8359ee6a8205011" 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 +375,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 +464,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 +474,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 +501,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 +526,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.119", ] [[package]] @@ -537,7 +537,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ "darling_core", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -567,7 +567,7 @@ dependencies = [ "darling", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -577,7 +577,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" dependencies = [ "derive_builder_core", - "syn", + "syn 2.0.119", ] [[package]] @@ -609,7 +609,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -673,9 +673,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 +722,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,9 +823,11 @@ 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]] @@ -873,7 +875,7 @@ dependencies = [ "indicatif 0.17.11", "libc", "log", - "rand 0.9.4", + "rand 0.9.5", "reqwest", "serde", "serde_json", @@ -894,9 +896,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 +906,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", @@ -954,7 +956,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.8", + "webpki-roots 1.0.9", ] [[package]] @@ -1104,11 +1106,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", @@ -1219,9 +1221,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", ] @@ -1266,7 +1268,7 @@ dependencies = [ "quote", "regex-syntax", "rustc_version", - "syn", + "syn 2.0.119", ] [[package]] @@ -1319,14 +1321,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" @@ -1346,9 +1348,9 @@ dependencies = [ [[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 +1376,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1540,9 +1542,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 +1577,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", @@ -1632,14 +1634,15 @@ 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", "rustls", @@ -1653,23 +1656,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 +1697,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 +1751,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" @@ -1807,9 +1819,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 +1831,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 +1884,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots 1.0.8", + "webpki-roots 1.0.9", ] [[package]] @@ -1891,9 +1903,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc_version" @@ -1919,9 +1931,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 +1946,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 +1967,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" @@ -1982,9 +1994,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 +2004,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.2", ] [[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 +2070,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 +2088,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 +2145,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 = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +checksum = "a207d6d6a2b7fc470b80443726053f18a2481b7e1eee970597051596567987a3" dependencies = [ "proc-macro2", "quote", @@ -2159,7 +2182,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2183,29 +2206,29 @@ 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.2", ] [[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 +2275,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,14 +2297,15 @@ 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", + "libc", "log", "logos", "macro_rules_attribute", @@ -2291,7 +2315,7 @@ dependencies = [ "paste", "pcre2", "ptr_hash", - "rand 0.9.4", + "rand 0.9.5", "rayon", "rayon-cond", "regex", @@ -2320,7 +2344,7 @@ dependencies = [ "dary_heap", "derive_builder", "esaxx-rs", - "indicatif 0.18.5", + "indicatif 0.18.6", "itertools 0.14.0", "log", "rayon", @@ -2344,14 +2368,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", @@ -2382,9 +2406,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.3" +version = "1.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee" dependencies = [ "bytes", "libc", @@ -2481,7 +2505,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2719,7 +2743,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "wasm-bindgen-shared", ] @@ -2771,14 +2795,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 +3032,9 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "985eec839aaf2a1270af8f4ebcf63cf9401cfd90f0902f97c28d9f104ffbde72" [[package]] name = "yada" @@ -3037,28 +3061,28 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.52" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.52" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -3078,7 +3102,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] @@ -3118,11 +3142,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 a946c1e1b..b67ebfccc 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -71,6 +71,10 @@ onig = { version = "6.5.1", optional = true } pcre2 = { version = "0.2", optional = true } logos = { version = "0.15", optional = true } # compile-time DFA lexer reference (pure Rust) +[target.'cfg(unix)'.dependencies] +# pthread_atfork registration for the fork-safe encode pool (poison-and-rebuild). +libc = "0.2" + [features] # `fancy-regex` is the OPTIONAL system-regex backend, needed ONLY for a `Split` pre-tokenizer with an # arbitrary (non-GPT) regex or the `Replace` normalizer — the atomsplit-native pre-tokenizers (GPT-2, diff --git a/tokenizers/tk-encode/benches/pipeline_benchmark.rs b/tokenizers/tk-encode/benches/pipeline_benchmark.rs index 190d0807a..fbb61bb56 100644 --- a/tokenizers/tk-encode/benches/pipeline_benchmark.rs +++ b/tokenizers/tk-encode/benches/pipeline_benchmark.rs @@ -16,8 +16,6 @@ use criterion::{BenchmarkId, Criterion, Throughput}; use tk_encode::pipeline::PipelineTokenizer; use tk_encode::Tokenizer; -// (name, tokenizer.json): bert-wiki exercises the Bert/WordPiece path; dsv4 exercises the deepseek -// pre-tokenizer unroll (`fsm_deepseek`) + BPE end-to-end. const TOKENIZERS: &[(&str, &str)] = &[ ("bert", "../data/bert-wiki.json"), ("dsv4", "../data/deepseek-v4-flash-base-tokenizer.json"), @@ -74,15 +72,14 @@ fn bench_pipeline(c: &mut Criterion) { let mut group = c.benchmark_group(format!("{tok_name}-{corpus}")); for (target_bytes, label) in CHUNK_SIZES { let chunks = make_chunks(&lines, *target_bytes); + let refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); let total_bytes: u64 = chunks.iter().map(|s| s.len() as u64).sum(); group.throughput(Throughput::Bytes(total_bytes)); group.bench_function(BenchmarkId::from_parameter(label), |b| { b.iter(|| { - let mut n = 0usize; - for chunk in &chunks { - n += pipeline.encode(chunk, false).unwrap().len(); - } - black_box(n) + pipeline.encode(&refs[..]).for_each(|r| { + black_box(r.1.unwrap()); + }); }) }); } diff --git a/tokenizers/tk-encode/examples/binsize_pipeline.rs b/tokenizers/tk-encode/examples/binsize_pipeline.rs index 09b953735..1d8a0e564 100644 --- a/tokenizers/tk-encode/examples/binsize_pipeline.rs +++ b/tokenizers/tk-encode/examples/binsize_pipeline.rs @@ -18,5 +18,8 @@ fn main() { .expect("usage: binsize_pipeline "); let tok = Tokenizer::from_file(path).unwrap(); let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); - println!("{}", pipeline.encode(text.as_str(), false).unwrap().len()); + println!( + "{}", + pipeline.encode(text.as_str()).map(|r| r.1.unwrap().len()).sum::() + ); } diff --git a/tokenizers/tk-encode/examples/fixture_bench.rs b/tokenizers/tk-encode/examples/fixture_bench.rs index 0dcde13c2..d3f0ac65c 100644 --- a/tokenizers/tk-encode/examples/fixture_bench.rs +++ b/tokenizers/tk-encode/examples/fixture_bench.rs @@ -34,7 +34,7 @@ use logos::Logos; use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use rayon::ThreadPoolBuilder; use serde_json::{json, Value}; -use tk_encode::pipeline::{Model, PipelineTokenizer}; +use tk_encode::pipeline::PipelineTokenizer; use tk_encode::{AddedToken, ModelWrapper, Tokenizer}; use tokenizers_release::{AddedToken as BaselineAddedToken, Tokenizer as BaselineTokenizer}; @@ -134,19 +134,14 @@ fn thread_counts() -> Vec { /// Median MB/s of encoding `chunks` across `n` threads in a *private* rayon pool (so the sweep can't /// perturb — or be perturbed by — the global pool). One `encode` call per chunk; the sum is `black_box`'d /// so the work can't be elided. -fn par_mbps( - encode: impl Fn(&str) -> usize + Sync, - chunks: &[String], - bytes: usize, - n: usize, -) -> f64 { +fn par_mbps(encode: impl Fn(&str) + Sync, chunks: &[String], bytes: usize, n: usize) -> f64 { let pool = ThreadPoolBuilder::new().num_threads(n).build().unwrap(); - let run = || pool.install(|| chunks.par_iter().map(|c| encode(c.as_str())).sum::()); - black_box(run()); // warm the pool + lazy structures + let run = || pool.install(|| chunks.par_iter().for_each(|c| encode(c.as_str()))); + run(); // warm the pool + lazy structures let mut samples = Vec::with_capacity(REPS); for _ in 0..REPS { let t = Instant::now(); - black_box(run()); + run(); samples.push(t.elapsed().as_secs_f64()); } bytes as f64 / median_secs(samples) / 1e6 @@ -164,9 +159,22 @@ fn bench_threads( let counts = thread_counts(); let (mut pipe, mut base) = (Vec::new(), Vec::new()); for &n in &counts { - let b = baseline.map(|b| par_mbps(|s| b.encode(s, false).unwrap().len(), chunks, bytes, n)); + let b = baseline.map(|b| { + par_mbps( + |s| { + black_box(b.encode(s, false).unwrap()); + }, + chunks, + bytes, + n, + ) + }); let p = par_mbps( - |s| pipeline.encode(s, false).unwrap().len(), + |s| { + pipeline.encode(s).for_each(|r| { + black_box(r.1.unwrap()); + }); + }, chunks, bytes, n, @@ -181,44 +189,106 @@ fn bench_threads( json!({ "counts": counts, "pipeline_mbps": pipe, "baseline_mbps": base }) } -fn time_pass(encode: &dyn Fn(&str) -> usize, chunks: &[String]) -> f64 { +fn time_pass(encode: &dyn Fn(&str), chunks: &[String]) -> f64 { let start = Instant::now(); - let mut n = 0usize; for chunk in chunks { - n += encode(chunk); + encode(chunk); } - black_box(n); start.elapsed().as_secs_f64() } -/// Median wall-time (seconds) of one pass over `chunks` through the shared encode core -/// `PipelineTokenizer::encode_generic::`. `STAGE` is a const generic, so each -/// level is a branchless specialization with the later stages compiled out — timing -/// successive levels and subtracting gives each stage's marginal cost (the ablation -/// ladder), no profiler and no per-segment instrumentation. +/// Cumulative pipeline stage levels for the ablation ladder, in execution +/// order: each level runs every stage up to and including itself; `Model` is a +/// full encode, `Frame` is the special-token scan + iteration only. +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum Stage { + Frame, + Normalize, + Split, + Model, +} + +/// Median wall-time (seconds) of one pass over `chunks`, running the pipeline's +/// stages up to and including `stage`. The partial pipeline is **recomposed +/// here in the bench** from the pipeline's public components +/// (`get_added_vocabulary` / `get_normalizer` / `get_pre_tokenizer` / +/// `get_model`), mirroring the walker in the library's encode kernel — the +/// library itself carries no staging or timing scaffolding. Timing successive +/// levels and subtracting gives each stage's marginal cost (the ablation +/// ladder). A runtime branch per segment is fine in a bench; the subtraction +/// cancels its cost. /// -/// Both caller-owned buffers are reused across chunks and `black_box`'d each iteration: -/// `output` anchors the special-scan/normalize/model work, `pre_tokens` anchors the -/// split stage, so under fat LTO no dead partial stage gets optimized away. The -/// `black_box` lives here, in the bench — never in the library. -fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String]) -> f64 { - let mut out = Vec::new(); +/// The caller-owned buffers are reused across chunks and `black_box`'d each +/// iteration: `out` anchors the special-scan/normalize/model work, `pre_tokens` +/// anchors the split stage, so under fat LTO no dead partial stage gets +/// optimized away. +fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String], stage: Stage) -> f64 { + use std::borrow::Cow; + use tk_encode::pipeline::{ + Model as _, Normalizer as _, PipelineToken, PreTokenizer as _, Segment, + SpecialSegmentIterator, + }; + + let added_vocabulary = pipeline.get_added_vocabulary(); + let normalizer = pipeline.get_normalizer(); + let pre_tokenizer = pipeline.get_pre_tokenizer(); + let model = pipeline.get_model(); + + let mut out: Vec = Vec::new(); let mut pre_tokens = Vec::new(); - let mut scratch = pipeline.get_model().init_scratch(); - let mut run = || { + let mut scratch = model.init_scratch(); + let mut run = |out: &mut Vec, pre_tokens: &mut Vec<_>| { for chunk in chunks { out.clear(); - let _ = - pipeline.encode_generic::(chunk, &mut pre_tokens, &mut scratch, &mut out); + for segment in SpecialSegmentIterator::new(chunk, added_vocabulary, false) { + match segment { + Segment::SpecialToken(id) => out.push(PipelineToken { id }), + Segment::Text(text) => { + let normalized: Cow = if stage >= Stage::Normalize { + match normalizer { + Some(n) => n.normalize(text).unwrap(), + None => Cow::Borrowed(text), + } + } else { + Cow::Borrowed(text) + }; + for seg in SpecialSegmentIterator::new(&normalized, added_vocabulary, true) + { + match seg { + Segment::SpecialToken(id) => out.push(PipelineToken { id }), + Segment::Text(normalized_chunk) => { + if stage >= Stage::Split { + pre_tokens.clear(); + pre_tokenizer + .pre_tokenize(normalized_chunk, pre_tokens) + .unwrap(); + if stage >= Stage::Model { + for pre_token in pre_tokens.iter() { + model + .tokenize_pipeline( + &normalized_chunk[pre_token.range()], + &mut scratch, + out, + ) + .unwrap(); + } + } + } + } + } + } + } + } + } black_box(&out); black_box(&pre_tokens); } }; - run(); // warm-up + run(&mut out, &mut pre_tokens); // warm-up let mut samples = Vec::with_capacity(REPS); for _ in 0..REPS { let start = Instant::now(); - run(); + run(&mut out, &mut pre_tokens); samples.push(start.elapsed().as_secs_f64()); } median_secs(samples) @@ -350,7 +420,9 @@ fn memory_child(which: &str, model: &Path) { drop(tok); let after_load = rss_now().unwrap_or(0); for c in &chunks { - n += pipeline.encode(c, false).unwrap().len(); + pipeline.encode(c).for_each(|r| { + black_box(r.1.unwrap()); + }); } (after_load, rss_now().unwrap_or(0)) } @@ -672,10 +744,18 @@ fn bench_model( files: &[(String, PathBuf)], model_json: &Path, ) -> Vec { - let pipe_enc = |s: &str| pipeline.encode(s, false).unwrap().len(); + let pipe_enc = |s: &str| { + pipeline.encode(s).for_each(|r| { + black_box(r.1.unwrap()); + }); + }; // per-model: the reference regex(es) onig will time on each fixture (empty for non-regex pretoks). let regexes = pretok_regexes(model_json); - let base_enc = baseline.map(|b| move |s: &str| b.encode(s, false).unwrap().len()); + let base_enc = baseline.map(|b| { + move |s: &str| { + black_box(b.encode(s, false).unwrap()); + } + }); let mut rows = Vec::new(); for (group, path) in files { @@ -686,9 +766,8 @@ fn bench_model( let pipe_ids = |c: &String| -> Vec { pipeline - .encode(c, false) - .unwrap() - .iter() + .encode(c) + .flat_map(|r| r.1.unwrap()) .map(|t| t.id) .collect() }; @@ -727,14 +806,13 @@ fn bench_model( fmt(base_mbps) ); - // Staged decomposition of the pipeline's own encode via the ablation ladder: + // Staged decomposition of the pipeline's encode via the ablation ladder: // time each cumulative stage level, then subtract to isolate each stage's cost - // (e.g. model = t_model - t_split). Levels are the named STAGE_* consts on - // PipelineTokenizer rather than bare 0/1/2/3. - let t_frame = stage_secs::<{ PipelineTokenizer::STAGE_FRAME }>(pipeline, &chunks); - let t_norm = stage_secs::<{ PipelineTokenizer::STAGE_NORMALIZE }>(pipeline, &chunks); - let t_split = stage_secs::<{ PipelineTokenizer::STAGE_SPLIT }>(pipeline, &chunks); - let t_model = stage_secs::<{ PipelineTokenizer::STAGE_MODEL }>(pipeline, &chunks); + // (e.g. model = t_model - t_split). + let t_frame = stage_secs(pipeline, &chunks, Stage::Frame); + let t_norm = stage_secs(pipeline, &chunks, Stage::Normalize); + let t_split = stage_secs(pipeline, &chunks, Stage::Split); + let t_model = stage_secs(pipeline, &chunks, Stage::Model); // Two distinct "split" costs are separated here: `added_split` is the // added/special-token scan (the SpecialSegmentIterator over the AddedVocabulary, // captured by the FRAME level), `pre_tokenize` is the pre-tokenizer split. All @@ -880,7 +958,7 @@ fn main() { // and has no range-based impl. Probe once and downgrade to "unsupported" // (with the reason) instead of panicking partway through the bench. let pipeline = match PipelineTokenizer::try_from(&tok) { - Ok(p) => match p.encode(PROBE, false) { + Ok(p) => match p.encode(PROBE).map(|(_, r)| r).collect::, _>>() { Ok(_) => p, Err(e) => { eprintln!(" pipeline builds but can't encode yet ({shape}): {e}"); diff --git a/tokenizers/tk-encode/src/normalizers/mod.rs b/tokenizers/tk-encode/src/normalizers/mod.rs index 9d4cd31a8..1763101c3 100644 --- a/tokenizers/tk-encode/src/normalizers/mod.rs +++ b/tokenizers/tk-encode/src/normalizers/mod.rs @@ -238,6 +238,31 @@ impl pipeline::Normalizer for NormalizerWrapper { } } +impl NormalizerWrapper { + /// checks if the following is possible for the given normalizer: + /// ```rs + /// let (a, b) = input.split_at(n); + /// assert!(normalizer.normalize(a) + normalizer.normalize(b), normalizer.normalize(input)); + /// ``` + pub(crate) fn preserves_stride_boundaries(&self) -> bool { + match self { + NormalizerWrapper::NFC(_) + | NormalizerWrapper::NFD(_) + | NormalizerWrapper::NFKC(_) + | NormalizerWrapper::NFKD(_) + | NormalizerWrapper::Lowercase(_) + | NormalizerWrapper::BertNormalizer(_) + | NormalizerWrapper::StripAccents(_) => true, + NormalizerWrapper::Sequence(seq) => { + seq.as_ref().iter().all(Self::preserves_stride_boundaries) + } + // Strip, Prepend, Replace, Precompiled, Nmt, ByteLevel: may cross + // chunk boundaries. + _ => false, + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs index 85fb87b91..d652285cd 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs @@ -62,7 +62,7 @@ impl PipelineSequence { /// Isolated, non-inverted `Split`s carrying deepseek's `[\p{N}{1,3}, CJK, big]` regexes (the trailing /// byte-map `ByteLevel` converts to `PipelinePreTokenizer::None`). Routes the whole split to one /// `fsm_deepseek` pass. - fn is_deepseek(&self) -> bool { + pub(crate) fn is_deepseek(&self) -> bool { use crate::pre_tokenizers::split::SplitPattern; use crate::tokenizer::SplitDelimiterBehavior::Isolated; let regex = |i: usize| match self.pre_tokenizers.get(i) { @@ -79,6 +79,11 @@ impl PipelineSequence { (Some(a), Some(b), Some(c)) if crate::utils::is_deepseek(a, b, c) ) } + + /// The sequence members, in application order. + pub(crate) fn members(&self) -> &[PipelinePreTokenizer] { + &self.pre_tokenizers + } } impl TryFrom for PipelineSequence { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index e7b33b134..903bb61fa 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -103,6 +103,12 @@ impl Split { let regex = match compiled { Ok(re) => Some(re), Err(_) if fsm.is_some() => None, + // deepseek's member patterns aren't standalone FSMs but never split on + // their own — a recognized deepseek `Sequence` drives them via + // `fsm_deepseek`, so they need no backend to construct. + Err(_) if matches!(&pattern, SplitPattern::Regex(r) if crate::utils::is_deepseek_member(r)) => { + None + } Err(e) => return Err(e), }; @@ -121,6 +127,16 @@ impl Split { /// Isolated)`, the form the native FSM fast path requires (the inverted match /// set is the gaps, and these patterns leave no gaps). Rewrite to it so /// cl100k/o200k route to `fsm_cl100k`/`fsm_o200k` instead of the SysRegex fallback. + /// Whether the pipeline drives this `Split` with a native `atomsplit` FSM + /// (a recognized GPT regex in its canonical `Isolated`, non-inverted + /// usage). Every recognized family (gpt2 / cl100k / o200k) tokenizes such + /// that no token contains a non-whitespace→space transition — the invariant + /// the parallel runtime relies on to cut raw text at those boundaries. + pub(crate) fn has_pipeline_fsm(&self) -> bool { + use crate::tokenizer::SplitDelimiterBehavior::Isolated; + self.fsm.is_some() && !self.invert && self.behavior == Isolated + } + pub(crate) fn canonicalized_for_pipeline(self) -> Result { use crate::tokenizer::SplitDelimiterBehavior::{Isolated, Removed}; if self.fsm.is_some() && self.invert && self.behavior == Removed { diff --git a/tokenizers/tk-encode/src/tokenizer/mod.rs b/tokenizers/tk-encode/src/tokenizer/mod.rs index 607abc237..f4eff8021 100644 --- a/tokenizers/tk-encode/src/tokenizer/mod.rs +++ b/tokenizers/tk-encode/src/tokenizer/mod.rs @@ -27,6 +27,7 @@ mod encoding; pub mod normalizer; pub mod pattern; pub mod pipeline; +pub(crate) mod pool; pub mod pre_tokenizer; mod serialization; diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 9a8ee0abc..509c34609 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1,8 +1,49 @@ +//! Parallel encode runtime for [`PipelineTokenizer`]. +//! +//! [`PipelineTokenizer::encode`] returns an owned [`EncodeHandle`] immediately; +//! workers encode in the background. Results surface through the handle's two +//! faces: the blocking [`Iterator`], yielding `(input index, result)` in +//! completion order (and *assisting* — the caller drains pending work instead of +//! idling), and [`EncodeHandle::wait_for_completion`], which blocks and returns +//! results back in input order. +//! +//! **Job model.** Work is split into `Item`s behind one atomic claim cursor in +//! `JobCore`; rayon pool tasks and the consuming thread both drain it, so +//! there is no scheduler — just a cursor and caller-assist. Per-worker +//! `EncodeState` scratch persists (thread-local, reset-not-realloc) for cache +//! warmth. Panics are caught per unit and delivered as that input's `Err`. +//! +//! **Rung ladder.** Every encode is decomposed the same way, outermost stage +//! first, cheapest split first: +//! +//! * **Rung 0 (always):** special tokens are the outermost pipeline stage, so +//! cutting at each match is unconditionally safe. The builder peels them +//! (matched specials resolve immediately, no worker) and hands each +//! special-free text segment down to the sub-rung. This costs one matcher +//! scan (~0.1% of encode) and is what makes the byte-level boundary able to +//! cut at number transitions — a strided segment can never contain a special. +//! * **Sub-rung (`ParallelPlan`, chosen per config):** applied to each segment. +//! `SplitRaw` strides it into ~`PARALLEL_MIN_TOTAL_BYTES` windows workers +//! boundary-snap locally via a `StrideBoundary` (a pure predicate over the +//! `atomsplit` tag stream the config's grammar provably cannot cross); +//! `SplitNormalized` serial-normalizes the segment then strides the buffer; +//! `SplitAtModel` serial-normalizes + pre-tokenizes and parallelizes the model +//! over span-groups; `WholeInput` leaves the segment to one worker (batch-level +//! parallelism only). The two serial-prefix rungs run only when the batch alone +//! can't feed the pool. Never wrong, only less parallel. +//! +//! **Fork safety** lives in `pool`: a `pthread_atfork` child abandons the +//! stale pool without touching it and lazily rebuilds. + use std::cell::RefCell; +use std::collections::VecDeque; use std::convert::TryInto; +use std::ops::Range; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{mpsc, Arc}; use std::{borrow::Cow, convert::TryFrom}; -use atomsplit::classify::classify; +use atomsplit::classify::{classify, in_mask, mask, Atom}; use crate::models::bpe::{BpeScratch, PipelineBPE}; use crate::models::unigram::{Unigram, UnigramScratch}; @@ -25,10 +66,10 @@ use crate::{ unicode_scripts::UnicodeScripts, whitespace::{Whitespace, WhitespaceSplit}, }, - ModelWrapper, PostProcessorWrapper, PreTokenizerWrapper, Token, Tokenizer, + ModelWrapper, PostProcessorWrapper, PreTokenizerWrapper, Tokenizer, }; -use super::{Result, SplitDelimiterBehavior}; +use super::{pool, Result, SplitDelimiterBehavior}; pub use atomsplit::fsm::Span; @@ -63,7 +104,7 @@ pub trait Normalizer { /// Range-based pre-tokenization: yields spans into the input rather than owned /// substrings, so the pipeline can pre-tokenize without allocating. pub trait PreTokenizer { - /// Split `text` into pre-tokens, appending to `out`. Ranges are into `text`. + /// Span `text` into pre-tokens, appending to `out`. Ranges are into `text`. fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()>; } @@ -108,6 +149,35 @@ impl PreTokenizer for PipelinePreTokenizer { } } +impl PipelinePreTokenizer { + /// The `StrideBoundary` this pre-tokenizer guarantees, if any: a cut + /// finder its splitting rules provably cannot tokenize across. This is the + /// extension point for raw-text parallel chunking — a pre-tokenizer that + /// can prove a boundary for its grammar adds an arm here and large inputs + /// split under it automatically; anything without one falls back to the + /// model-stage split (still correct, just a serial pre-tokenize). + pub(crate) fn stride_boundary(&self) -> Option { + match self { + // Whitespace-delimiting: whitespace is a dropped delimiter, so any + // whitespace character is a safe cut (see `boundary_at_whitespace`). + Self::Whitespace(_) | Self::WhitespaceSplit(_) | Self::Bert(_) => { + Some(boundary_at_whitespace) + } + // atomsplit-FSM byte-level families: cut at a space after non-ws, or + // a newline after a word character (see `boundary_fsm`). + Self::Split(s) if s.has_pipeline_fsm() => Some(boundary_fsm), + Self::Sequence(seq) if seq.is_deepseek() => Some(boundary_fsm), + Self::Sequence(seq) => match seq.members() { + [only] => only.stride_boundary(), + // The llama/GPT-4 shape: [Split(FSM), ByteLevel(no regex) -> None]. + [split, Self::None] => split.stride_boundary(), + _ => None, + }, + _ => None, + } + } +} + impl TryFrom for PipelinePreTokenizer { type Error = crate::Error; @@ -156,12 +226,6 @@ pub struct PipelineToken { pub id: u32, } -impl From for PipelineToken { - fn from(value: Token) -> Self { - Self { id: value.id } - } -} - /// Finds special/added tokens in a text segment so the pipeline can carve them /// out before running the model. pub trait PipelinePatternMatcher { @@ -179,7 +243,7 @@ pub trait PipelinePatternMatcher { /// A piece of the input produced by [`SpecialSegmentIterator`]. pub enum Segment<'a> { - /// Ordinary text still to be (optonally normalized), pre-tokenized and run through the model. + /// Ordinary text still to be (optionally) normalized, pre-tokenized and run through the model. Text(&'a str), /// A matched special token, identified by its vocabulary id. SpecialToken(u32), @@ -211,13 +275,9 @@ pub struct SpecialSegmentIterator<'a, 'b, PatternMatcher: PipelinePatternMatcher impl<'a, 'b, PatternMatcher: PipelinePatternMatcher> SpecialSegmentIterator<'a, 'b, PatternMatcher> { - /// Create a new iterator over [`Segment`] of the [`input`]. + /// Create a new iterator over the [`Segment`]s of `input`. /// This iterator will yield [`Segment`] in order. - pub(crate) fn new( - input: &'a str, - pattern_matcher: &'b PatternMatcher, - normalized: bool, - ) -> Self { + pub fn new(input: &'a str, pattern_matcher: &'b PatternMatcher, normalized: bool) -> Self { Self { input, pattern_matcher, @@ -269,12 +329,843 @@ impl<'a, 'b, PatternMatcher: PipelinePatternMatcher> Iterator /// Experimental encode-only pipeline built from a [`Tokenizer`]. Runs the same /// stages (special-token split → normalize → pre-tokenize → model) over borrowed /// ranges to avoid the reference path's allocations. +/// +/// This is a cheap-clone **handle** over an `Arc`'d frozen core (the v2 +/// `Arc` shape): cloning shares the vocabulary/model, and an owned +/// [`EncodeHandle`] keeps the tokenizer alive past the `encode` call by holding a +/// clone. +#[derive(Clone)] pub struct PipelineTokenizer { + inner: Arc, +} + +/// The frozen pipeline components, shared by every clone of the handle. +struct PipelineInner { added_vocabulary: BucketAddedVocabulary, normalizer: Option, pre_tokenizer: PipelinePreTokenizer, model: PipelineModel, _post_processor: Option, + /// Whether some `normalized`-form added token can't survive a stride cut — + /// either its content carries an internal [`StrideBoundary`] cut, or it is an + /// affix (`lstrip`/`rstrip`) token whose absorbed adjacent whitespace a cut + /// could split. Its presence disables `SplitRaw`/`SplitNormalized`. Only + /// `normalized` tokens count — raw-matched ones are peeled by rung 0 before + /// any striding. Precomputed once at build (the `plan()` gate reads it per + /// `encode`, so it must not re-scan the vocab). + normalized_added_token_blocks_stride: bool, +} + +// comptime verification that PipelinePreTokenizer is Send + Sync +const _: fn() = || { + fn assert_send_sync() {} + assert_send_sync::(); +}; + +/// Reusable per-encode scratch, kept alive across sequences processed by the +/// same worker thread (`reset`, never realloc'd). Today it holds only the +/// pre-token span buffer; as the fast (SIMD) pre-tokenizer lands, the byte +/// `tags` buffer and chunk spans will live here too, so the parallel runtime +/// keeps a single reset-not-realloc scratch per worker. +#[derive(Default)] +pub(crate) struct EncodeState { + pre_tokens: Vec, + /// Classify tag scratch for the stride boundary scans (grow-and-reuse). + tags: Vec, + /// Model-specific heap buffers (merge heap, word symbols, candidate string, + /// …), lazily built for the model kind this thread last encoded with and + /// re-initialized on a kind switch — thread-local state is shared across + /// tokenizers on a thread. Deliberately *not* cleared by [`reset`](Self::reset): + /// persistence across sequences and encode calls is the point. + model_scratch: Option, +} + +impl EncodeState { + pub(crate) fn new() -> Self { + Self::default() + } + + /// Clear the scratch for reuse on the next sequence, keeping capacity. + pub(crate) fn reset(&mut self) { + self.pre_tokens.clear(); + } +} + +/// Total input bytes below which `encode` runs inline on the calling thread (no pool): +/// under this, the pool's dispatch/wakeup cost outweighs the parallelism gain. The same +/// threshold sets the stride size for raw-cut chunking and the span-group target for +/// `SplitAtModel`, so it is the single "unit of parallel work" knob. Bench-tunable — +/// splitting the three roles into separate knobs is a tracked tuning item. +const PARALLEL_MIN_TOTAL_BYTES: usize = 8 * 1024; + +/// The byte range that stride `index` of `text` owns after boundary snapping, +/// or `None` when the stride's window contains no safe boundary (its bytes +/// belong to the previous stride's chunk, which extends across it — that +/// stride emits an empty result to keep per-input chunk counts static). +/// +/// This is the **fused executor**'s scheduling primitive: there is no central +/// cut list and no serial pre-scan — each worker claims a stride index and +/// derives its chunk locally. Strides are fixed +/// `PARALLEL_MIN_TOTAL_BYTES`-sized windows; a chunk starts at the first safe +/// boundary inside its own window (stride 0 starts at 0) and ends at the first +/// safe boundary found in the following windows — the exact scan those +/// windows' owners also perform, so adjacent workers agree on every cut with +/// no communication. Total scanning is O(input): each window is scanned at +/// most twice (once by its owner, once by an extending predecessor). Text with +/// no boundaries at all degrades to stride 0 owning the whole input. +fn stride_range( + text: &str, + index: usize, + boundary: StrideBoundary, + tags: &mut Vec, +) -> Option> { + const STRIDE: usize = PARALLEL_MIN_TOTAL_BYTES; + let len = text.len(); + let start = if index == 0 { + 0 + } else { + boundary_in_window( + text, + index * STRIDE, + ((index + 1) * STRIDE).min(len), + boundary, + tags, + )? + }; + let mut window = index + 1; + let end = loop { + let lo = window * STRIDE; + if lo >= len { + break len; + } + if let Some(b) = + boundary_in_window(text, lo, ((window + 1) * STRIDE).min(len), boundary, tags) + { + break b; + } + window += 1; + }; + (start < end).then_some(start..end) +} + +/// A safe-cut predicate for one pre-tokenizer family, expressed over the +/// `atomsplit` **tag stream** (the same SIMD `classify` substrate the FSM +/// pre-tokenizers run on — boundary logic is not reimplemented byte-matching). +/// `window_tags[0]` is one byte of left context; the fn returns the first +/// window-relative position `i >= 1` that is a safe cut: no pre-token — under +/// that family's splitting rules — can span it, and tokenization of each side +/// is independent of the other. +/// +/// The boundary definition **belongs to the pre-tokenizer** +/// ([`PipelinePreTokenizer::stride_boundary`]): each configuration exposes a +/// predicate its own token grammar provably cannot cross, and the scheduler +/// ([`stride_range`]) is agnostic to what the boundary is. Must be a pure +/// function of the tags — workers rely on recomputing identical results. +type StrideBoundary = fn(window_tags: &[u8]) -> Option; + +/// Tag classes that are a complete non-whitespace character — the previous-char +/// requirement for [`boundary_fsm`]'s space cut. Excludes whitespace and the +/// bytes that don't name a class on their own (`Cont`/`MultiByte`/`Sentinel`) +/// plus control∪unassigned. Read via [`prev_char_tag`], so a multibyte +/// character is seen at its resolved lead, never a continuation byte. +const NON_WS_PREV: u16 = Atom::Letter.bit() + | Atom::NumWord.bit() + | Atom::NumOther.bit() + | Atom::Mark.bit() + | Atom::Connector.bit() + | Atom::Punct.bit() + | Atom::Apostrophe.bit() + | Atom::SymOther.bit() + | Atom::NumericOther.bit(); + +/// Tag classes whose token can never carry a *trailing* newline — the +/// previous-char requirement for [`boundary_fsm`]'s newline cut. Letters +/// (`\p{L}+`, which is also every CJK character), numbers (`\p{N}{1,3}`), and +/// marks all live in rules that exclude `\r\n`, so the word token ends before +/// the newline. Punctuation/symbol classes are deliberately absent: the +/// cl100k/o200k/deepseek punct rule ` ?[…]+[\r\n]*` absorbs following newlines +/// into the punct token, so a cut after punct-then-newline is not +/// side-independent. +const NEWLINE_PREV: u16 = + Atom::Letter.bit() | Atom::NumWord.bit() | Atom::NumOther.bit() | Atom::Mark.bit(); + +/// Whitespace classes other than newline (space + other unicode whitespace) — +/// the `cur` side of [`boundary_fsm`]'s whitespace cut. Newline is handled on +/// its own ([`NEWLINE_PREV`]) because it alone can be absorbed by a trailing +/// `[\r\n]*` in the punct rule. +const WS_NON_NEWLINE: u16 = Atom::Space.bit() | Atom::WsOther.bit(); + +/// Tag classes a number run can start *after* — the previous-char requirement +/// for [`safe_fsm_cut`]'s letter/connector/punct→number cut. The byte-level +/// families tokenize digit runs as `\p{N}{1,3}`, which never joins a preceding +/// non-number, and every preceding token rule (letters, punct) excludes `\p{N}`, +/// so the number run starts a fresh token regardless of side. (Number→number is +/// excluded — a long digit run splits into `{1,3}` groups the cut would break.) +const NUM_START_PREV: u16 = Atom::Letter.bit() + | Atom::Mark.bit() + | Atom::Connector.bit() + | Atom::Punct.bit() + | Atom::Apostrophe.bit() + | Atom::SymOther.bit() + | Atom::NumericOther.bit(); + +/// The class tag of the character ending just before window position `i`: walk +/// back over `Cont` continuation bytes to the character's lead. `i >= 1`, and +/// `window_tags[0]` is always a lead (the classify window starts on a char +/// boundary), so the walk stays in bounds. Lets a boundary predicate judge the +/// preceding character by its real class even when it is multibyte (e.g. CJK, +/// whose lead is `Letter` and whose trailing bytes are `Cont`). +fn prev_char_tag(window_tags: &[u8], i: usize) -> u8 { + let mut j = i - 1; + while j > 0 && in_mask(window_tags[j], Atom::Cont.bit()) { + j -= 1; + } + window_tags[j] +} + +/// `StrideBoundary` for whitespace-delimiting pre-tokenizers (`Whitespace`, +/// `WhitespaceSplit`, `Bert`): *any* whitespace character is a safe cut. These +/// splitters drop whitespace as a delimiter and never emit a token that spans +/// it, so a cut placed at a whitespace character leaves the tokens on either +/// side unchanged — even one inside a whitespace run, since each side then +/// drops its own run. Cutting *at* the character (it joins the right chunk) +/// keeps the cut on a char boundary: a whitespace lead is classed `WS`, its +/// continuation bytes `Cont`. Denser than a newline-only rule, which matters +/// for single-line inputs (minified JSON, long CSV rows) that would otherwise +/// see no intra-input split. +fn boundary_at_whitespace(window_tags: &[u8]) -> Option { + (1..window_tags.len()).find(|&i| in_mask(window_tags[i], mask::WS)) +} + +/// Whether a cut at a `prev`→`cur` character transition is side-independent for +/// the atomsplit-FSM byte-level families (gpt2 / cl100k / o200k / deepseek) — +/// the cut character `cur` joins the right chunk. Two safe transitions, each +/// checked branch by branch against `atomsplit::regexes` for all four families: +/// +/// * **→ space / other whitespace** after a non-whitespace character — the ws +/// joins the right chunk, where a leading space binds to the following token +/// (` ?\p{L}+`); no rule matches a non-whitespace character directly followed +/// by a (non-newline) space *inside* one token. +/// * **→ newline** after a word character ([`NEWLINE_PREV`]) — a letter/number/ +/// mark token never carries a trailing newline, so the newline opens a fresh +/// whitespace match (`\s*[\r\n]+`). Gives space-sparse scripts (CJK) +/// intra-document parallelism. +/// +/// * **→ number** after a letter/connector/punct/symbol ([`NUM_START_PREV`]) — +/// a `\p{N}{1,3}` run opens a fresh token; no preceding rule absorbs digits. +/// * **number → letter** — the digit run ended (`\p{N}{1,3}`), so a following +/// letter opens a ` ?\p{L}+` token. Together these two give whitespace-free +/// byte-level content (base64, minified JSON, long IDs) intra-input cuts. +/// +/// The number transitions *can* fall inside a special (`<…token_0|>`), so they +/// are only sound because **rung 0 always peels specials before any striding** — +/// a strided segment is special-free, so no cut can bisect one. The gate then +/// only needs to guard normalized-form added tokens (see `normalized_added_token_blocks_stride`). +fn safe_fsm_cut(prev: u8, cur: u8) -> bool { + (in_mask(cur, WS_NON_NEWLINE) && in_mask(prev, NON_WS_PREV)) + || (in_mask(cur, mask::NEWLINE) && in_mask(prev, NEWLINE_PREV)) + || (in_mask(cur, mask::NUMBER) && in_mask(prev, NUM_START_PREV)) + || (in_mask(cur, mask::LETTER) && in_mask(prev, mask::NUMBER)) +} + +/// `StrideBoundary` for the atomsplit-FSM byte-level families: the first +/// [`safe_fsm_cut`] transition in the window. The previous character is read +/// through [`prev_char_tag`], so every branch also fires after a multibyte +/// character (e.g. CJK). +fn boundary_fsm(window_tags: &[u8]) -> Option { + (1..window_tags.len()).find(|&i| safe_fsm_cut(prev_char_tag(window_tags, i), window_tags[i])) +} + +/// Whether `token`'s own bytes contain a cut of `boundary` strictly inside it — +/// i.e. a stride could split this token and each half would mis-frame it. Used +/// by `normalized_added_token_blocks_stride` to disqualify striding for a `normalized` +/// added token whose content carries an internal cut (e.g. a `_→0` number +/// transition). Raw-matched tokens don't need this — rung 0 peels them before +/// any striding — so only normalized-form tokens reach this check. +fn has_internal_boundary(token: &str, boundary: StrideBoundary) -> bool { + if token.len() < 2 { + return false; + } + let mut tags = vec![0u8; token.len()]; + classify(token.as_bytes(), &mut tags); + boundary(&tags).is_some() +} + +/// Find the first safe cut in `text[lo..hi)` (absolute byte index; `lo >= 1`): +/// classify in blocks — SIMD [`classify`] into the caller's `tags` scratch, +/// with one byte of left context so the predicate can see the previous tag — +/// and apply `boundary` per block, early-exiting on the first hit. Blocked so +/// the common case (a boundary within the first block) classifies ~1 KB, not +/// the whole window. +/// +/// Block edges are snapped to char boundaries (`classify` must see complete +/// characters). The snap is a pure function of the text, so adjacent workers +/// still agree; the extra bytes a snap adds are continuation bytes, which tag +/// as `Cont` and can never satisfy a boundary predicate. +fn boundary_in_window( + text: &str, + lo: usize, + hi: usize, + boundary: StrideBoundary, + tags: &mut Vec, +) -> Option { + const BLOCK: usize = 1024; + debug_assert!(1 <= lo && lo <= hi && hi <= text.len()); + let mut block_lo = lo; + while block_lo < hi { + let mut ctx = block_lo - 1; + while !text.is_char_boundary(ctx) { + ctx -= 1; + } + let mut block_hi = (block_lo + BLOCK).min(hi); + while block_hi < text.len() && !text.is_char_boundary(block_hi) { + block_hi += 1; + } + let window = &text.as_bytes()[ctx..block_hi]; + tags.clear(); + tags.resize(window.len(), 0); + classify(window, tags); + if let Some(rel) = boundary(tags) { + return Some(ctx + rel); + } + block_lo = block_hi; + } + None +} + +/// The **sub-rung** applied to each special-free text segment (rung 0 always +/// peels the specials first), chosen per tokenizer configuration: cut the raw +/// text when the components allow it; otherwise normalize serially and cut the +/// normalized text; otherwise pre-tokenize serially and parallelize the model; +/// otherwise rely on batch-level parallelism only. Never wrong, only less +/// parallel. +#[derive(Clone, Copy, Debug)] +enum ParallelPlan { + /// Rung 1: cut raw text at safe boundaries; each chunk runs the full + /// pipeline (`encode_one_with`). Requires a boundary-preserving normalizer. + SplitRaw(StrideBoundary), + /// Rung 2: frame + normalize serially on the caller, then cut the normalized + /// text at the pre-tokenizer's boundary and parallelize pre-tokenize + model + /// (`encode_normalized`) per chunk. The serial normalize takes the normalizer + /// out of the correctness equation, so this needs only a splittable + /// pre-tokenizer — the fallback when a non-boundary-preserving normalizer + /// rules out `SplitRaw` but the pre-tokenizer still cuts. + SplitNormalized(StrideBoundary), + /// Rung 3: frame + normalize + pre-tokenize serially on the caller; + /// parallelize the model over ~`PARALLEL_MIN_TOTAL_BYTES`-sized groups of + /// pre-token spans. Valid for every config (each pre-token is modeled + /// independently). + SplitAtModel, + /// No intra-input parallelism. + WholeInput, +} + +/// Serial-prefix output of the `SplitAtModel` sub-rung for **one special-free +/// text segment** (rung 0 already peeled raw specials): normalization, +/// normalized-special framing and pre-tokenization are done; what remains is +/// model work, packaged as span-groups. Everything is index/range-based (no +/// borrows), so both the scoped and the owned paths can consume it. +struct ModelPrefix { + /// Ordered pipeline units for this segment. + units: Vec, + /// The owned normalized segment ([`PrefixText::Norm`] points here) — `Some` + /// only when the normalizer rewrote this chunk. One chunk is normalized once, + /// so there is at most one buffer (all `Norm` units point at it). + norm_buf: Option, + /// Flat pool of pre-token spans; groups reference sub-ranges of it. + spans: Vec, +} + +enum PrefixUnit { + /// A matched special/added token — already resolved, no worker needed. + Special(u32), + /// A group of consecutive pre-token spans over one text segment. + Group { + text: PrefixText, + spans: Range, + }, +} + +/// Where a span-group's text lives. Spans are relative to the resolved text. +enum PrefixText { + /// A sub-range of the raw input (normalizer absent or returned `Borrowed`). + Raw(Range), + /// A sub-range of `norm_bufs[buf]` (the normalizer rewrote this segment). + Norm { buf: usize, range: Range }, +} + +/// Byte offset of the subslice `sub` within its base string `base`. +/// Both must be views into the same allocation (`sub` derived from `base`). +fn offset_within(base: &str, sub: &str) -> usize { + let off = sub.as_ptr() as usize - base.as_ptr() as usize; + debug_assert!(off + sub.len() <= base.len()); + off +} + +thread_local! { + /// Per-thread reusable encode scratch: pool drainers, inline callers and the + /// caller-assist path all keep their buffers warm across encode calls + /// (reset, never realloc'd) — pool threads persist for the process, so this + /// is the cache-warmth contract of `EncodeState`. + static SCRATCH: RefCell = RefCell::new(EncodeState::new()); +} + +/// One unit's result travelling from a worker back to the job: the tokens (or +/// error) of the `idx`-th chunk of input `seq`. +struct UnitResult { + seq: usize, + idx: usize, + tokens: Result>, +} + +enum HandleInner { + /// Fully computed; yields `(input index, result)` in input order. + Ready(std::iter::Enumerate>>>), + /// Being filled by pool workers over a channel; each input surfaces (tagged + /// with its index) as soon as all of its chunks are in — completion order, + /// not input order. + Streaming { + core: Arc, + state: StreamState, + }, +} + +/// Reassembly state for a streaming job. An input is split into parallel chunks +/// that arrive interleaved with other inputs'; this buffers them per input and +/// surfaces each input (with its index) the moment its last chunk lands — no +/// input-order cursor, so a completed input never waits behind an earlier-index +/// one that is still encoding. +struct StreamState { + rx: mpsc::Receiver, + /// `slots[seq][chunk_idx]` — a chunk's tokens once received; the row length + /// is the number of chunks expected for that input. + slots: Vec>>>, + /// chunks received so far per input. + filled: Vec, + /// first error seen per input (if any). + err: Vec>, + /// Inputs whose chunks are all in, awaiting emit (also seeded with any + /// zero-chunk inputs, which are complete from the start). + ready: VecDeque, + /// Inputs not yet returned by [`next_completed`](Self::next_completed); when + /// zero, the job is fully drained. + remaining: usize, + /// Caller-assist bookkeeping: set once the job's cursor is exhausted so + /// the blocking faces stop trying to claim units. + assist_done: bool, +} + +impl StreamState { + fn new(rx: mpsc::Receiver, counts: Vec) -> Self { + let slots: Vec>>> = counts + .iter() + .map(|&n| (0..n).map(|_| None).collect()) + .collect(); + let ready = counts + .iter() + .enumerate() + .filter_map(|(i, &n)| (n == 0).then_some(i)) + .collect(); + let filled = vec![0; counts.len()]; + let err = counts.iter().map(|_| None).collect(); + let remaining = counts.len(); + Self { + rx, + slots, + filled, + err, + ready, + remaining, + assist_done: false, + } + } + + fn record(&mut self, UnitResult { seq, idx, tokens }: UnitResult) { + match tokens { + Ok(t) => self.slots[seq][idx] = Some(t), + Err(e) => { + if self.err[seq].is_none() { + self.err[seq] = Some(e); + } + self.slots[seq][idx] = Some(Vec::new()); + } + } + self.filled[seq] += 1; + if self.filled[seq] == self.slots[seq].len() { + self.ready.push_back(seq); + } + } + + /// Concatenate input `seq`'s chunks in order, or its first error. + fn take_result(&mut self, seq: usize) -> Result> { + if let Some(e) = self.err[seq].take() { + return Err(e); + } + Ok(self.slots[seq] + .iter_mut() + .flat_map(|slot| slot.take().unwrap_or_default()) + .collect()) + } + + /// Channel closed with inputs still short of their chunks: a worker died + /// without sending (a bug or a killed thread). Surface an error for one such + /// input; repeated calls drain the rest, so a consumer never hangs. + fn disconnected(&mut self) -> (usize, Result>) { + let seq = (0..self.slots.len()) + .find(|&s| self.filled[s] < self.slots[s].len()) + .expect("channel closed with work outstanding but no unfinished input"); + self.filled[seq] = self.slots[seq].len(); + self.remaining -= 1; + ( + seq, + Err("encode: worker exited before all chunks were produced".into()), + ) + } + + /// Yield the next input to *complete*, tagged with its index — not input + /// order. While waiting it absorbs whatever the channel already holds, then + /// encodes pending units itself (caller-assist) rather than idling. `None` + /// once every input has been emitted. + fn next_completed(&mut self, core: &JobCore) -> Option<(usize, Result>)> { + loop { + if let Some(seq) = self.ready.pop_front() { + self.remaining -= 1; + return Some((seq, self.take_result(seq))); + } + if self.remaining == 0 { + return None; + } + // Record anything already delivered without blocking. + match self.rx.try_recv() { + Ok(msg) => { + self.record(msg); + continue; + } + Err(mpsc::TryRecvError::Empty) => {} + Err(mpsc::TryRecvError::Disconnected) => return Some(self.disconnected()), + } + // Nothing ready: claim a unit of work ourselves rather than idling. + if !self.assist_done { + if SCRATCH.with(|st| core.run_one(&mut st.borrow_mut())) { + continue; + } + // Cursor exhausted: nothing left to claim. + self.assist_done = true; + continue; + } + match self.rx.recv() { + Ok(msg) => self.record(msg), + Err(_) => return Some(self.disconnected()), + } + } + } +} + +/// Input storage an [`EncodeHandle`] keeps alive for its whole lifetime — the +/// contract that makes the returnable handle sound with no scoped lifetimes: +/// implementors promise every `&str` returned by [`get`](Self::get) stays valid +/// and unchanged for as long as the value lives. Owned storage (`String`, +/// `Vec`) qualifies trivially; bindings can wrap refcounted foreign +/// strings (e.g. a `Py` keep-alive over Python's stable UTF-8 buffer) +/// for zero-copy owned encodes. +#[allow(clippy::len_without_is_empty)] // an empty batch is legal; nothing branches on it +pub trait Inputs: Send + Sync + 'static { + fn len(&self) -> usize; + fn get(&self, i: usize) -> &str; +} + +impl Inputs for String { + fn len(&self) -> usize { + 1 + } + fn get(&self, _i: usize) -> &str { + self + } +} +impl Inputs for Vec { + fn len(&self) -> usize { + self.as_slice().len() + } + fn get(&self, i: usize) -> &str { + &self[i] + } +} + +/// Conversion into [`Inputs`] for [`PipelineTokenizer::encode`]. Owned +/// inputs convert for free (a move); `&str`-family inputs pay **one copy** at +/// the boundary (measured at low single-digit % of encode cost). +pub trait IntoInputs { + type Inputs: Inputs; + fn into_inputs(self) -> Self::Inputs; +} + +impl IntoInputs for String { + type Inputs = String; + fn into_inputs(self) -> String { + self + } +} +impl IntoInputs for Vec { + type Inputs = Vec; + fn into_inputs(self) -> Vec { + self + } +} +impl IntoInputs for &str { + type Inputs = String; + fn into_inputs(self) -> String { + self.to_owned() + } +} +impl IntoInputs for &String { + type Inputs = String; + fn into_inputs(self) -> String { + self.clone() + } +} +impl IntoInputs for &[&str] { + type Inputs = Vec; + fn into_inputs(self) -> Vec { + self.iter().map(|s| (*s).to_owned()).collect() + } +} +/// One owned encode call's worth of work, shared with the pool workers via +/// `Arc`. Fully safe: the job owns its storage and a tokenizer handle, and the +/// chunks are byte ranges into the storage — nothing borrows the caller, so the +/// job outlives the `encode` call freely. Dropping the [`EncodeHandle`] sets +/// `cancelled`; workers stop claiming, in-flight chunks finish into the closed +/// channel, and the last `Arc` holder releases the storage. +struct JobCore { + storage: Box, + /// Ordered worker units, (seq, idx)-tagged for reassembly. + units: Vec, + /// Owned normalized segments for `SplitAtModel` units + /// ([`PrefixText::Norm`] points here). + norm_bufs: Vec, + /// Flat pool of pre-token spans for `SplitAtModel` units. + spans: Vec, + tokenizer: PipelineTokenizer, + cursor: AtomicUsize, + cancelled: AtomicBool, + tx: mpsc::Sender, +} + +/// One owned work unit: the `idx`-th piece of input `seq`. +struct Item { + seq: usize, + idx: usize, + work: Work, +} + +/// A unit's byte region within its source text: either a fixed range, or a +/// fused stride the worker boundary-snaps itself (no serial pre-scan). +#[derive(Clone)] +enum Region { + Full(Range), + Stride { + index: usize, + boundary: StrideBoundary, + /// The special-free segment (byte range within the source) this stride + /// tiles. Strides index from the segment start, so a rung-0-peeled text + /// segment is strided independently of the specials around it. + seg: Range, + }, +} + +impl Region { + /// Tile a special-free `seg` (byte range within its source) into work + /// regions: fixed ~`PARALLEL_MIN_TOTAL_BYTES` fused strides (workers + /// boundary-snap them), or a single `Full` region when the segment is too + /// small to be worth striding. The shared "stride-or-whole" decision for the + /// raw (`SplitRaw`) and normalized (`SplitNormalized`) paths. + fn tile(seg: Range, boundary: StrideBoundary) -> Vec { + let len = seg.len(); + if len >= 2 * PARALLEL_MIN_TOTAL_BYTES { + (0..len.div_ceil(PARALLEL_MIN_TOTAL_BYTES)) + .map(|index| Region::Stride { + index, + boundary, + seg: seg.clone(), + }) + .collect() + } else { + vec![Region::Full(seg)] + } + } + + /// Resolve to an absolute byte range within `text`. `None` when a stride's + /// window has no boundary — that stride yields an empty result, keeping the + /// per-input chunk count static for reassembly. A stride resolves against its + /// own `seg` sub-slice, then shifts the result back to `text` coordinates. + fn resolve(&self, text: &str, tags: &mut Vec) -> Option> { + match self { + Region::Full(r) => Some(r.clone()), + Region::Stride { + index, + boundary, + seg, + } => stride_range(&text[seg.clone()], *index, *boundary, tags) + .map(|r| seg.start + r.start..seg.start + r.end), + } + } + + /// Rough byte size, for LPT ordering. + fn approx_len(&self) -> usize { + match self { + Region::Full(r) => r.len(), + Region::Stride { seg, .. } => PARALLEL_MIN_TOTAL_BYTES.min(seg.len()), + } + } +} + +/// The pipeline stage a unit enters at (and the text its region indexes into). +enum Work { + /// Full pipeline (`encode_one_with`) over a raw region of input `seq`. + Raw(Region), + /// Pre-tokenize + model (`encode_normalized`) over a region of the + /// already-normalized buffer `JobCore::norm_bufs[buf]` (`SplitNormalized`). + Normalized { buf: usize, region: Region }, + /// Model-only over a span-group (`SplitAtModel`): `text` resolves against + /// input `seq` (`Raw`) or `JobCore::norm_bufs` (`Norm`); `spans` indexes + /// `JobCore::spans`. + Model { + text: PrefixText, + spans: Range, + }, +} + +impl JobCore { + /// Claim and run one unit; `false` when the cursor is exhausted (or the job + /// cancelled). A panicking unit is delivered as that input's `Err` so the + /// consumer never hangs on a chunk that will not arrive. + fn run_one(&self, scratch: &mut EncodeState) -> bool { + if self.cancelled.load(Ordering::Relaxed) { + return false; + } + let i = self.cursor.fetch_add(1, Ordering::Relaxed); + if i >= self.units.len() { + return false; + } + let unit = &self.units[i]; + let seq = unit.seq; + let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match &unit.work { + Work::Raw(region) => { + let text = self.storage.get(seq); + match region.resolve(text, &mut scratch.tags) { + Some(range) => self.tokenizer.encode_one_with(&text[range], scratch), + None => Ok(Vec::new()), + } + } + Work::Normalized { buf, region } => { + let text = &self.norm_bufs[*buf]; + match region.resolve(text, &mut scratch.tags) { + Some(range) => { + let mut output = Vec::with_capacity(range.len() / 4); + self.tokenizer + .encode_normalized(&text[range], scratch, &mut output) + .map(|()| output) + } + None => Ok(Vec::new()), + } + } + Work::Model { text, spans } => { + let text = match text { + PrefixText::Raw(r) => &self.storage.get(seq)[r.clone()], + PrefixText::Norm { buf, range } => &self.norm_bufs[*buf][range.clone()], + }; + self.tokenizer + .encode_spans(text, &self.spans[spans.clone()], scratch) + } + })) + .unwrap_or_else(|_| Err("encode worker panicked".into())); + // A dropped receiver (job dropped early) just discards the result. + let _ = self.tx.send(UnitResult { + seq, + idx: unit.idx, + tokens: res, + }); + true + } + + fn cancel(&self) { + self.cancelled.store(true, Ordering::Relaxed); + } +} + +/// An owned, escapable encode — what [`PipelineTokenizer::encode`] returns. +/// Workers keep encoding in the background while the caller holds this; results +/// surface through: +/// - the blocking `Iterator`, yielding `(input index, result)` in **completion +/// order** as each input finishes (and *assisting*: the calling thread claims +/// and encodes pending chunks instead of idling), +/// - [`EncodeHandle::wait_for_completion`](Self::wait_for_completion), which +/// blocks and returns results placed back into **input order**. +/// +/// Dropping the job cancels all unclaimed work; chunks already being encoded +/// finish into the void. Worker panics surface as the affected input's `Err`. +pub struct EncodeHandle { + inner: HandleInner, +} + +impl EncodeHandle { + fn ready(results: Vec>>) -> Self { + Self { + inner: HandleInner::Ready(results.into_iter().enumerate()), + } + } + + fn streaming(core: Arc, rx: mpsc::Receiver, counts: Vec) -> Self { + Self { + inner: HandleInner::Streaming { + core, + state: StreamState::new(rx, counts), + }, + } + } + + /// Number of inputs the handle will process + fn len(&self) -> usize { + match &self.inner { + HandleInner::Ready(it) => it.len(), + HandleInner::Streaming { state, .. } => state.slots.len(), + } + } + + /// Block until all inputs are encoded, results are ordered in input order. + /// Fails fast: returns the first error it receives. + pub fn wait_for_completion(self) -> Result>> { + // XXX: `Vec::new` does not allocate when capacity == 0 + let mut out: Vec> = vec![Vec::new(); self.len()]; + for (seq, res) in self { + out[seq] = res?; + } + Ok(out) + } + +} + +impl Iterator for EncodeHandle { + /// `(input index, that input's tokens or error)`. The index is the position + /// in the input batch (always `0` for a single input); results arrive in + /// completion order, not input order. + type Item = (usize, Result>); + fn next(&mut self) -> Option { + match &mut self.inner { + HandleInner::Ready(it) => it.next(), + HandleInner::Streaming { core, state } => state.next_completed(core), + } + } +} + +impl Drop for EncodeHandle { + fn drop(&mut self) { + // Cancel whatever hasn't been claimed. Harmless after a full drain (the + // cursor is already exhausted); on an early drop it stops the workers + // from encoding into the void. + if let HandleInner::Streaming { core, .. } = &self.inner { + core.cancel(); + } + } } impl TryFrom<&Tokenizer> for PipelineTokenizer { @@ -297,6 +1188,25 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { let legacy_av = tok.get_added_vocabulary(); let mut added_tokens: Vec<_> = legacy_av.get_added_tokens_decoder().iter().collect(); added_tokens.sort_by_key(|(id, _)| **id); + // A stride cut must never land inside — or in the affix-stripped + // whitespace of — an added token (each half would then mis-frame it). + // Rung 0 always peels the *raw*-matched (`normalized == false`) tokens + // before any striding, so those are safe regardless of the boundary. + // Only `normalized`-form tokens survive into strided text; disqualify + // striding if any of them either (a) carries an internal cut of the + // pre-tokenizer's boundary (e.g. a `_→0` number transition), or (b) is an + // affix token (`lstrip`/`rstrip`) whose absorbed adjacent-whitespace run a + // whitespace-boundary cut could split. Precomputed once: this feeds the + // per-`encode` `plan()` path, so it must not iterate the vocab per call. + let normalized_added_token_blocks_stride = + pre_tokenizer.stride_boundary().is_some_and(|boundary| { + added_tokens + .iter() + .filter(|(_, t)| t.normalized) + .any(|(_, t)| { + t.lstrip || t.rstrip || has_internal_boundary(&t.content, boundary) + }) + }); let mut added_vocabulary = BucketAddedVocabulary::new(); added_vocabulary.add_tokens( added_tokens.into_iter().map(|(_, t)| BucketAddedToken { @@ -368,128 +1278,537 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { }; Ok(Self { - added_vocabulary, - normalizer: tok.get_normalizer().cloned(), - pre_tokenizer, - model, - _post_processor: tok.get_post_processor().cloned(), + inner: Arc::new(PipelineInner { + added_vocabulary, + normalizer: tok.get_normalizer().cloned(), + pre_tokenizer, + model, + _post_processor: tok.get_post_processor().cloned(), + normalized_added_token_blocks_stride, + }), }) } } impl PipelineTokenizer { - /// Stage gates for [`encode_upto`](Self::encode_upto), in execution order. Each - /// level runs every stage up to and including itself; `STAGE_MODEL` is a full - /// encode. `STAGE_FRAME` is the special-token scan + iteration only (the "other" - /// slice in the decomposition). - pub const STAGE_FRAME: u8 = 0; - pub const STAGE_NORMALIZE: u8 = 1; - pub const STAGE_SPLIT: u8 = 2; - pub const STAGE_MODEL: u8 = 3; + // Component accessors: introspection for callers (and for external benches + // that recompose partial pipelines — the library itself carries no staging + // or timing scaffolding). pub fn get_model(&self) -> &PipelineModel { - &self.model + &self.inner.model } - /// Encode `input` into token ids. + /// The added/special-token matcher the frame passes run with. + pub fn get_added_vocabulary(&self) -> &impl PipelinePatternMatcher { + &self.inner.added_vocabulary + } + + pub fn get_normalizer(&self) -> Option<&NormalizerWrapper> { + self.inner.normalizer.as_ref() + } + + pub fn get_pre_tokenizer(&self) -> &PipelinePreTokenizer { + &self.inner.pre_tokenizer + } + + /// Encode one or many sequences, returning an [`EncodeHandle`] — an owned, + /// escapable handle: workers keep encoding in the background while the caller + /// holds it, drains it (blocking `Iterator` / [`EncodeHandle::wait_for_completion`]). + /// **This is the default API.** /// - /// Special tokens are matched in two passes: - /// 1. on the raw input, - /// 2. then on each segment after normalization + /// The job is `'static`: it holds the input storage (see [`IntoInputs`] — + /// moving a `String`/`Vec` is zero-copy; + /// passing `&str` pays one copy at the boundary) and a cheap clone of this + /// tokenizer handle. Chunks are byte *ranges* into that storage, so there is no + /// unsafe and no join barrier on this path — dropping the job cancels unclaimed + /// work, and the last worker releases the storage. /// - /// This way, special / added tokens declared on raw or normalized text are both caught. - /// The remaining text is pre-tokenized and run through the model span by span. + /// Scheduling: cost gate below `PARALLEL_MIN_TOTAL_BYTES` (the job comes + /// back already complete, no pool touched); otherwise each input is decomposed + /// rung 0 first (peel specials) then per the sub-[`ParallelPlan`], into work + /// units drained from one shared claim cursor by pool tasks and by the job's + /// blocking faces (caller-assist). Results surface in input order; worker + /// panics surface as that input's `Err`. /// - /// todo: wire the post-processing - pub fn encode(&self, input: &str, _add_special_tokens: bool) -> Result> { - let mut output = Vec::new(); - let mut pre_tokens = Vec::new(); - let mut scratch = self.model.init_scratch(); + /// [`EncodeHandle::wait_for_completion`]: EncodeHandle::wait_for_completion + pub fn encode(&self, inputs: I) -> EncodeHandle { + let storage = inputs.into_inputs(); + let n_inputs = storage.len(); + let refs: Vec<&str> = (0..n_inputs).map(|i| storage.get(i)).collect(); + + // Cost gate: small work is encoded right here; the job comes back Ready. + let total_bytes: usize = refs.iter().map(|s| s.len()).sum(); + if total_bytes < PARALLEL_MIN_TOTAL_BYTES { + return EncodeHandle::ready(self.encode_serial(&refs)); + } + let Some(pool) = pool::rayon() else { + return EncodeHandle::ready(self.encode_serial(&refs)); + }; - self.encode_generic::<{ Self::STAGE_MODEL }>( - input, - &mut pre_tokens, - &mut scratch, - &mut output, - )?; + // Pick the parallel plan and build the (seq, idx)-tagged unit list. + // Everything is range-based, so nothing borrows `refs`/`storage`. + let mut counts = vec![0usize; n_inputs]; + let mut units: Vec = Vec::new(); + let mut norm_bufs: Vec = Vec::new(); + let mut spans: Vec = Vec::new(); + let mut preresolved: Vec = Vec::new(); + + // Rung 0 is the outermost pipeline stage and unconditionally safe: special + // tokens are framed before anything else, so cutting at each match and + // encoding the inter-special text independently reproduces the whole + // encode exactly. So EVERY plan begins here — peel specials (a matcher + // scan the worker runs anyway, ~0.1% of encode) and hand each special-free + // text segment to the config's deepest sub-rung. `plan()` names that + // *sub*-rung; rung 0 wraps them all, so a strided segment can never + // bisect a special (which is what lets the byte-level boundary include + // number transitions — see `safe_fsm_cut`). + let plan = self.plan(); + // The expensive sub-rungs (rung 2/3) run a serial normalize / pre-tokenize + // prefix on the calling thread — only worth it when the batch alone can't + // feed the pool. Rung 0/1 add no such prefix, so they always apply. + let prefix_worth_it = n_inputs < pool.current_num_threads(); + + for (seq, &input) in refs.iter().enumerate() { + // An input too small to split into even two strides gains nothing from + // intra-input decomposition — and running rung 0's matcher scan per + // input on the calling thread would just serialize overhead across a + // large batch of small inputs. Emit one whole unit; the worker frames + // its specials itself (in parallel). No striding here, so this needs + // no rung-0 peel to be safe. + if input.len() < 2 * PARALLEL_MIN_TOTAL_BYTES { + self.emit_full(seq, 0..input.len(), &mut counts, &mut units); + continue; + } + for segment in SpecialSegmentIterator::new(input, &self.inner.added_vocabulary, false) { + let chunk = match segment { + // Rung 0: a matched special resolves immediately, no worker. + Segment::SpecialToken(id) => { + let idx = counts[seq]; + counts[seq] += 1; + preresolved.push(UnitResult { + seq, + idx, + tokens: Ok(vec![PipelineToken { id }]), + }); + continue; + } + Segment::Text(chunk) => chunk, + }; + let base = offset_within(input, chunk); + let seg = base..base + chunk.len(); + let big_enough = chunk.len() >= 2 * PARALLEL_MIN_TOTAL_BYTES; + + match &plan { + // Rung 1: stride the raw segment; workers boundary-snap + // locally with no serial pre-scan. + ParallelPlan::SplitRaw(boundary) => { + for region in Region::tile(seg.clone(), *boundary) { + let idx = counts[seq]; + counts[seq] += 1; + units.push(Item { + seq, + idx, + work: Work::Raw(region), + }); + } + } + // Rung 2: serial-normalize the segment into an owned buffer, + // then stride the buffer (pre-tokenize + model in parallel). + ParallelPlan::SplitNormalized(boundary) if prefix_worth_it && big_enough => { + match self.inner.normalizer.as_ref().map_or_else( + || Ok(chunk.to_owned()), + |n| n.normalize(chunk).map(|c| c.into_owned()), + ) { + Ok(buf_str) => { + let buf = norm_bufs.len(); + norm_bufs.push(buf_str); + for region in Region::tile(0..norm_bufs[buf].len(), *boundary) { + let idx = counts[seq]; + counts[seq] += 1; + units.push(Item { + seq, + idx, + work: Work::Normalized { buf, region }, + }); + } + } + // Normalize error → a whole-segment unit surfaces it. + Err(_) => self.emit_full(seq, seg.clone(), &mut counts, &mut units), + } + } + // Rung 3: serial normalize + pre-tokenize the segment, package + // the model work as span-groups. + ParallelPlan::SplitAtModel if prefix_worth_it && big_enough => { + match self.model_prefix_chunk(chunk, input) { + Ok(prefix) => { + // Merge this segment's prefix into the job pools. + // The prefix has at most one norm buffer (local + // index 0); its job index is `buf_base`. + let buf_base = norm_bufs.len(); + let span_base = spans.len(); + norm_bufs.extend(prefix.norm_buf); + spans.extend_from_slice(&prefix.spans); + for unit in prefix.units { + let idx = counts[seq]; + counts[seq] += 1; + match unit { + PrefixUnit::Special(t) => preresolved.push(UnitResult { + seq, + idx, + tokens: Ok(vec![PipelineToken { id: t }]), + }), + PrefixUnit::Group { text, spans: sr } => { + // Norm buffers were merged at `buf_base`; + // Raw offsets are already input-absolute. + let text = match text { + PrefixText::Norm { buf, range } => { + PrefixText::Norm { + buf: buf + buf_base, + range, + } + } + raw => raw, + }; + units.push(Item { + seq, + idx, + work: Work::Model { + text, + spans: sr.start + span_base..sr.end + span_base, + }, + }); + } + } + } + } + Err(_) => self.emit_full(seq, seg.clone(), &mut counts, &mut units), + } + } + // `WholeInput`, or an expensive sub-rung whose prefix isn't + // worth it here (batch saturates the pool, or a small segment): + // one whole-segment unit, full pipeline in the worker. + _ => self.emit_full(seq, seg.clone(), &mut counts, &mut units), + } + } + } + + // Too little schedulable work to be worth the pool. (Preresolved specials + // need no worker, so only worker `units` count; falling back re-encodes + // serially, which is byte-identical and cheap at this size.) + if units.len() < 2 { + return EncodeHandle::ready(self.encode_serial(&refs)); + } + drop(refs); + + // LPT: claim the largest units first so no giant chunk lands last + // (straggler tail); the (seq, idx) tags keep reassembly order-free. + units.sort_by_key(|unit| { + std::cmp::Reverse(match &unit.work { + Work::Raw(region) => region.approx_len(), + Work::Normalized { region, .. } => region.approx_len(), + Work::Model { spans: sr, .. } => { + let group = &spans[sr.clone()]; + group + .last() + .map(|last| (last.end - group[0].start) as usize) + .unwrap_or(0) + } + }) + }); + + let (tx, rx) = mpsc::channel::(); + // Units the serial prefix already resolved (special tokens) go straight + // into the channel; reassembly absorbs them like any other chunk. + for msg in preresolved { + let _ = tx.send(msg); + } + let n_units = units.len(); + let core = Arc::new(JobCore { + storage: Box::new(storage), + units, + norm_bufs, + spans, + tokenizer: self.clone(), + cursor: AtomicUsize::new(0), + cancelled: AtomicBool::new(false), + tx, + }); + + // Spawn one cursor-drainer per potential worker ('static: each holds an + // Arc of the core). Drainers that find the cursor exhausted are cheap + // no-ops. + let drainers = n_units.min(pool.current_num_threads()); + for _ in 0..drainers { + let core = Arc::clone(&core); + pool.spawn(move || { + SCRATCH.with(|st| { + let scratch = &mut *st.borrow_mut(); + while core.run_one(scratch) {} + }) + }); + } + + EncodeHandle::streaming(core, rx, counts) + } + + /// Model-only kernel for the `SplitAtModel` plan: tokenize each pre-token + /// span of (already normalized) `text` in order. + fn encode_spans( + &self, + text: &str, + spans: &[Span], + state: &mut EncodeState, + ) -> Result> { + // ~4.3 input bytes per token measured on English corpora; /4 is a + // conservative reserve that avoids most growth reallocations. + let covered = spans + .last() + .map(|last| (last.end - spans[0].start) as usize) + .unwrap_or(0); + let mut output = Vec::with_capacity(covered / 4); + let scratch = self.inner.model.scratch_in(&mut state.model_scratch); + for span in spans { + self.inner + .model + .tokenize_pipeline(&text[span.range()], scratch, &mut output)?; + } Ok(output) } - /// Single source of truth for the encode pipeline, generic over how many stages - /// run. `STAGE` is a **const generic**, so `if STAGE >= …` folds at compile time and - /// the disabled stages are compiled out — the full specialization ([`STAGE_MODEL`], - /// which [`encode`](Self::encode) calls) is branchless and identical to a - /// hand-written full pipeline, while the benchmark drives lower `STAGE` values to - /// time each stage's marginal cost (the ablation ladder), e.g. - /// `model = t(MODEL) − t(SPLIT)`. No runtime gate, no `Instant` in the loop. - /// - /// [`STAGE_MODEL`]: Self::STAGE_MODEL - /// - /// `output` and the `pre_tokens` scratch are caller-owned so a benchmark can reuse - /// them across calls and observe both buffers to anchor the ablation levels — the - /// library itself stays free of any `black_box`/timing artifact. - #[doc(hidden)] // public only so `examples/fixture_bench.rs` can drive partial stages - pub fn encode_generic( + /// Encode every input serially on the calling thread, reusing the thread's + /// warm scratch across the whole batch (one result per input). + fn encode_serial(&self, inputs: &[&str]) -> Vec>> { + SCRATCH.with(|st| { + let state = &mut *st.borrow_mut(); + inputs + .iter() + .map(|&input| self.encode_one_with(input, state)) + .collect() + }) + } + + /// The post-normalize suffix of the pipeline: frame (special tokens on + /// already-normalized text) → pre-tokenize → model, appending to `output`. + /// This is the entry point workers use when normalization has already run + /// (rung 2 / the `SplitNormalized` plan), and the shared tail of + /// [`encode_one_with`](Self::encode_one_with). Reuses the caller's + /// `EncodeState` scratch (pre-token buffer + model scratch). + fn encode_normalized( &self, - input: &str, - pre_tokens: &mut Vec, - scratch: &mut PipelineModelScratch, + normalized: &str, + state: &mut EncodeState, output: &mut Vec, ) -> Result<()> { - // First, we extract all special tokens from the non-normalized input - for segment in SpecialSegmentIterator::new(input, &self.added_vocabulary, false) { + let scratch = self.inner.model.scratch_in(&mut state.model_scratch); + let pre_tokens = &mut state.pre_tokens; + for segment in SpecialSegmentIterator::new(normalized, &self.inner.added_vocabulary, true) { match segment { - Segment::SpecialToken(token) => { - output.push(PipelineToken { id: token }); + Segment::SpecialToken(token) => output.push(PipelineToken { id: token }), + Segment::Text(chunk) => { + pre_tokens.clear(); + self.inner.pre_tokenizer.pre_tokenize(chunk, pre_tokens)?; + for pre_token in pre_tokens.iter() { + self.inner.model.tokenize_pipeline( + &chunk[pre_token.range()], + scratch, + output, + )?; + } } + } + } + Ok(()) + } + + /// The full single-sequence kernel: frame (special tokens on raw text) → + /// normalize → [`encode_normalized`](Self::encode_normalized), reusing the + /// caller's `EncodeState` scratch. The entry point when a worker holds a raw + /// region (rung 0/1: inline, pool drainers, caller-assist); span-group model + /// work goes through [`encode_spans`](Self::encode_spans) instead. + /// + /// Note: the post-processor (special-token framing like BOS/EOS) is not + /// wired yet; when it lands it applies per input after chunk concatenation. + fn encode_one_with(&self, input: &str, state: &mut EncodeState) -> Result> { + state.reset(); + // ~4.3 input bytes per token measured on English corpora; /4 is a + // conservative reserve that avoids most growth reallocations. + let mut output = Vec::with_capacity(input.len() / 4); + + // Extract the special tokens declared on raw text, normalize each text + // segment, then run the post-normalize suffix. + for segment in SpecialSegmentIterator::new(input, &self.inner.added_vocabulary, false) { + match segment { + Segment::SpecialToken(token) => output.push(PipelineToken { id: token }), Segment::Text(chunk) => { - let normalized: Cow = if STAGE >= Self::STAGE_NORMALIZE { - match &self.normalizer { - Some(normalizer) => normalizer.normalize(chunk)?, - None => Cow::Borrowed(chunk), - } - } else { - Cow::Borrowed(chunk) + let normalized: Cow = match &self.inner.normalizer { + Some(normalizer) => normalizer.normalize(chunk)?, + None => Cow::Borrowed(chunk), }; + self.encode_normalized(&normalized, state, &mut output)?; + } + } + } + Ok(output) + } - // Extract special tokens from the normalized input - for segment in - SpecialSegmentIterator::new(&normalized, &self.added_vocabulary, true) - { - match segment { - Segment::SpecialToken(token) => { - output.push(PipelineToken { id: token }); - } - Segment::Text(normalized_chunk) => { - if STAGE >= Self::STAGE_SPLIT { - // Pre-tokenize the chunk of normalized text - pre_tokens.clear(); - self.pre_tokenizer - .pre_tokenize(normalized_chunk, pre_tokens)?; - if STAGE >= Self::STAGE_MODEL { - // Tokenize each chunk - for pre_token in pre_tokens.iter() { - self.model.tokenize_pipeline( - &normalized_chunk[pre_token.range()], - scratch, - output, - )?; - } - } - } - } + /// The sub-rung this config applies to each special-free text segment (rung 0 + /// always peels specials first — see [`encode`](Self::encode)). The cheapest + /// safe rung wins (see `PARALLEL_RUNTIME_DESIGN.md` §8). + fn plan(&self) -> ParallelPlan { + if let Some(boundary) = self.stride_boundary() { + // Rung 1: no serial prefix, so always worth it when available. + ParallelPlan::SplitRaw(boundary) + } else if matches!(self.inner.model, PipelineModel::WordLevel(_)) { + // A WordLevel "model step" is one hash lookup per pre-token, and its + // pre-tokenize is cheap too — the serial prefix that rungs 2/3 need + // would dwarf the parallelized part, so lean on batch parallelism. + ParallelPlan::WholeInput + } else if let Some(boundary) = self.normalized_stride_boundary() { + // Rung 2: the normalizer rules out a raw cut, but the pre-tokenizer + // can cut the (serially) normalized text — parallelize pretok+model. + ParallelPlan::SplitNormalized(boundary) + } else { + ParallelPlan::SplitAtModel + } + } + + /// The pre-tokenizer's cut boundary for **already-normalized** text, gated + /// only on the added vocabulary (no normalizer requirement — rung 2 runs the + /// normalizer serially first, so it's out of the correctness equation). The + /// one remaining hazard is a `normalized`-form added token — which rung 0's + /// raw peel can't remove — being split by a stride; the precomputed + /// [`normalized_added_token_blocks_stride`](PipelineInner::normalized_added_token_blocks_stride) + /// bool covers it (raw-matched tokens are already peeled, so they don't gate). + fn normalized_stride_boundary(&self) -> Option { + if self.inner.normalized_added_token_blocks_stride { + return None; + } + self.inner.pre_tokenizer.stride_boundary() + } + + /// Whether — and where — the config lets **raw** text be cut and each piece + /// encoded independently with the same result as the whole (rung 1). Composes + /// the per-stage split-safety declarations: the normalizer must preserve + /// boundaries ([`NormalizerWrapper::preserves_stride_boundaries`]) *and* the + /// added vocabulary + pre-tokenizer must permit a cut + /// ([`normalized_stride_boundary`](Self::normalized_stride_boundary)). When + /// the normalizer doesn't preserve boundaries, rung 2 (`SplitNormalized`) + /// takes over via `normalized_stride_boundary` after a serial normalize. + fn stride_boundary(&self) -> Option { + let norm_ok = self + .inner + .normalizer + .as_ref() + .is_none_or(NormalizerWrapper::preserves_stride_boundaries); + if !norm_ok { + return None; + } + self.normalized_stride_boundary() + } + + /// Emit one whole-segment `Work::Raw` unit (full pipeline in the worker) for + /// the byte `range` of input `seq`. The shared fallback for `WholeInput` and + /// for the expensive sub-rungs when their serial prefix isn't worth running. + fn emit_full( + &self, + seq: usize, + range: Range, + counts: &mut [usize], + units: &mut Vec, + ) { + let idx = counts[seq]; + counts[seq] += 1; + units.push(Item { + seq, + idx, + work: Work::Raw(Region::Full(range)), + }); + } + + /// Serial prefix of the [`ParallelPlan::SplitAtModel`] plan (rung 3) over one + /// **special-free** text segment (rung 0 already peeled raw specials): normalize + /// → frame(normalized specials) → pre-tokenize, packaging the remaining model + /// work as ~`PARALLEL_MIN_TOTAL_BYTES`-sized span-groups. `input` is the whole + /// source the borrowed (`PrefixText::Raw`) offsets resolve against; `chunk` is a + /// subslice of it. Mirrors the walker in [`encode_one_with`](Self::encode_one_with) + /// stage by stage — the two must agree on unit order for reassembly to be + /// byte-identical to the serial encode. + fn model_prefix_chunk(&self, chunk: &str, input: &str) -> Result { + let mut prefix = ModelPrefix { + units: Vec::new(), + norm_buf: None, + spans: Vec::new(), + }; + let mut pre_tokens: Vec = Vec::new(); + let normalized: Cow = match &self.inner.normalizer { + Some(normalizer) => normalizer.normalize(chunk)?, + None => Cow::Borrowed(chunk), + }; + let owned = matches!(normalized, Cow::Owned(_)); + for seg in SpecialSegmentIterator::new(&normalized, &self.inner.added_vocabulary, true) { + match seg { + Segment::SpecialToken(token) => prefix.units.push(PrefixUnit::Special(token)), + Segment::Text(nchunk) => { + pre_tokens.clear(); + self.inner + .pre_tokenizer + .pre_tokenize(nchunk, &mut pre_tokens)?; + if pre_tokens.is_empty() { + continue; + } + // Where this segment's text lives. When owned, all groups + // point at the single `norm_buf` (local index 0, rebased to + // the job pool at merge); when borrowed, at the raw input. + let base = if owned { + offset_within(&normalized, nchunk) + } else { + // Borrowed all the way down: nchunk is a subslice of the + // raw input. + offset_within(input, nchunk) + }; + let text_range = base..base + nchunk.len(); + // Group consecutive spans into ~target-byte units. + let mut g = 0; + while g < pre_tokens.len() { + let group_start = pre_tokens[g].start; + let mut e = g + 1; + while e < pre_tokens.len() + && ((pre_tokens[e - 1].end - group_start) as usize) + < PARALLEL_MIN_TOTAL_BYTES + { + e += 1; } + let span_range = prefix.spans.len()..prefix.spans.len() + (e - g); + prefix.spans.extend_from_slice(&pre_tokens[g..e]); + let text = if owned { + // Local index 0 — the one `norm_buf`; merge rebases it. + PrefixText::Norm { + buf: 0, + range: text_range.clone(), + } + } else { + PrefixText::Raw(text_range.clone()) + }; + prefix.units.push(PrefixUnit::Group { + text, + spans: span_range, + }); + g = e; } } - }; + } } - Ok(()) + if let Cow::Owned(s) = normalized { + prefix.norm_buf = Some(s); + } + Ok(prefix) } } /// What [`split`] does with each split it forms #[derive(Clone, Copy, PartialEq, Eq, Debug)] -pub enum SplitPolicy { +pub(crate) enum SplitPolicy { /// Drop it, emit no split Remove, /// Emit it whole as one split @@ -498,7 +1817,7 @@ pub enum SplitPolicy { Isolate, } -/// Splits `text` into same-class groups, emitting each as a [`Split`] +/// Splits `text` into same-class groups, emitting each as a [`Span`] /// according to its [`SplitPolicy`]. /// /// `classify` maps each char to a small `Copy + Eq` class, the current @@ -506,7 +1825,7 @@ pub enum SplitPolicy { /// class), and `policy` decides what becomes of it. Ranges are byte offsets /// into `text`. #[inline(always)] -pub fn split( +pub(crate) fn split( text: &str, out: &mut Vec, classify: impl Fn(char) -> C, @@ -551,7 +1870,7 @@ pub fn split( /// joins the run before it (`"the-final"` -> `["the-", "final"]`). /// - `MergedWithNext` cuts *before* each delimiter, so it joins the run after it /// (`"the-final"` -> `["the", "-final"]`). -pub fn split_delimiter( +pub(crate) fn split_delimiter( text: &str, out: &mut Vec, is_delim: impl Fn(char) -> bool, @@ -619,7 +1938,7 @@ pub fn split_delimiter( /// `Pattern::find_matches` produces). This is the pipeline-side equivalent of /// the fold in `NormalizedString::split`; the arms mirror it exactly. Empty and /// removed pieces are dropped. -pub fn split_matches( +pub(crate) fn split_matches( out: &mut Vec, matches: Vec<((usize, usize), bool)>, behavior: SplitDelimiterBehavior, @@ -726,6 +2045,21 @@ pub enum PipelineModel { WordPiece(PipelineWordPiece), } +impl PipelineModel { + /// Get this model's per-thread scratch out of `slot`, (re)building it when + /// the slot is empty or was built for a different model kind — the + /// thread-local `EncodeState` is shared across tokenizers on a thread. + fn scratch_in<'a>( + &self, + slot: &'a mut Option, + ) -> &'a mut PipelineModelScratch { + if !slot.as_ref().is_some_and(|s| s.matches(self)) { + *slot = Some(self.init_scratch()); + } + slot.as_mut().unwrap() + } +} + impl Model for PipelineModel { type Scratch = PipelineModelScratch; @@ -769,6 +2103,20 @@ pub enum PipelineModelScratch { Unigram(UnigramScratch), } +impl PipelineModelScratch { + /// Whether this scratch was built for `model`'s kind (see + /// [`PipelineModel::scratch_in`]). + fn matches(&self, model: &PipelineModel) -> bool { + matches!( + (model, self), + (PipelineModel::BPE(_), Self::BPE(_)) + | (PipelineModel::Unigram(_), Self::Unigram(_)) + | (PipelineModel::WordLevel(_), Self::WordLevel(_)) + | (PipelineModel::WordPiece(_), Self::WordPiece(_)) + ) + } +} + impl ModelScratch for PipelineModelScratch {} #[cfg(test)] @@ -779,6 +2127,654 @@ mod tests { use crate::pre_tokenizers::byte_level::ByteLevel; use crate::pre_tokenizers::sequence::Sequence; + /// Serial oracle: encode one sequence with a fresh scratch, no pool. + fn encode_one(tok: &PipelineTokenizer, input: &str) -> Result> { + tok.encode_one_with(input, &mut EncodeState::new()) + } + + /// Test-only convenience: drain a single-input handle to its one result. + /// (Not on the public API — throughput callers use the streaming `Iterator`.) + trait IntoSingle { + fn into_single(self) -> Result>; + } + impl IntoSingle for EncodeHandle { + fn into_single(self) -> Result> { + self.wait_for_completion() + .map(|all| all.into_iter().next().unwrap_or_default()) + } + } + + /// A chunk-safe pipeline: WordLevel model + `WhitespaceSplit` pre-tokenizer, no + /// normalizer, no added tokens. + fn chunk_safe_pipeline() -> PipelineTokenizer { + use crate::models::wordlevel::WordLevelBuilder; + use crate::pre_tokenizers::whitespace::WhitespaceSplit; + use ahash::AHashMap; + let vocab: AHashMap = ["aa", "bb", "cc", ""] + .iter() + .enumerate() + .map(|(i, w)| ((*w).to_string(), i as u32)) + .collect(); + let model = WordLevelBuilder::new() + .vocab(vocab) + .unk_token("".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(model); + tok.with_pre_tokenizer(Some(WhitespaceSplit)); + PipelineTokenizer::try_from(&tok).unwrap() + } + + /// A large chunk-safe input is actually split into multiple chunks, and the chunked + /// (parallel) encode is identical to the serial encode. + #[test] + fn intra_seq_splits_and_matches_serial() { + let tok = chunk_safe_pipeline(); + assert!( + tok.stride_boundary().is_some(), + "wordlevel + WhitespaceSplit must be raw-chunkable" + ); + + let big = "aa bb cc\n".repeat(4000); // ~36 KB, newline-delimited + assert!(big.len() > 2 * PARALLEL_MIN_TOTAL_BYTES); + + let strides = big.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES); + let mut tags = Vec::new(); + let chunks = (0..strides) + .filter_map(|j| stride_range(&big, j, boundary_at_whitespace, &mut tags)) + .count(); + assert!( + chunks > 1, + "expected the large input to split, got {}", + chunks + ); + + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let parallel = ids(tok.encode(big.as_str()).into_single().unwrap()); + let serial = ids(encode_one(&tok, big.as_str()).unwrap()); + assert_eq!(parallel, serial, "chunked encode must equal serial encode"); + + // The owned (returnable) job must agree too — move the storage in. + let owned = ids(tok.encode(big.clone()).into_single().unwrap()); + assert_eq!(owned, serial, "owned EncodeHandle must equal serial encode"); + } + + /// Whitespace-family strides cut *at* a whitespace character (space or + /// newline) and tile the input exactly. + #[test] + fn stride_cuts_at_whitespace() { + let input = "aa bb cc\n".repeat(4000); + let chunks = assert_strides_partition(&input, boundary_at_whitespace, |bytes, start| { + assert!( + bytes[start].is_ascii_whitespace(), + "cut must fall at a whitespace character" + ); + }); + assert!(chunks > 1, "expected several chunks, got {}", chunks); + } + + /// Every non-empty stride chunk must tile the input exactly (no gap, no + /// overlap — adjacent workers must agree on every cut with no + /// communication), with each chunk start satisfying the cut predicate. + /// Returns the number of non-empty chunks. + fn assert_strides_partition( + input: &str, + boundary: StrideBoundary, + check_start: impl Fn(&[u8], usize), + ) -> usize { + let bytes = input.as_bytes(); + let mut tags = Vec::new(); + let mut covered = 0; + let mut chunks = 0; + for j in 0..input.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES) { + if let Some(range) = stride_range(input, j, boundary, &mut tags) { + assert_eq!(range.start, covered, "chunks must tile the input"); + assert!(range.end > range.start); + if range.start > 0 { + check_start(bytes, range.start); + } + covered = range.end; + chunks += 1; + } + } + assert_eq!(covered, input.len(), "chunks must cover the whole input"); + chunks + } + + /// An input with no safe boundary at all degrades to stride 0 owning the + /// whole input (serial), every other stride empty. + #[test] + fn strides_without_boundaries_degrade_to_whole_input() { + let input = "a".repeat(4 * PARALLEL_MIN_TOTAL_BYTES); + let mut tags = Vec::new(); + for boundary in [boundary_at_whitespace as StrideBoundary, boundary_fsm] { + assert_eq!( + stride_range(&input, 0, boundary, &mut tags), + Some(0..input.len()) + ); + for j in 1..input.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES) { + assert_eq!( + stride_range(&input, j, boundary, &mut tags), + None, + "stride {}", + j + ); + } + } + } + + /// `encode` drains via both faces (`wait_for_completion` bulk collect and the + /// streaming `Iterator`), over a mixed batch (large splittable inputs + a small one), and + /// both must equal the serial per-input encode in input order. + fn assert_encode_matches_serial(tok: &PipelineTokenizer) { + let a = "aa bb cc\n".repeat(3000); // ~27 KB, splits into several chunks + let b = "bb cc aa\n".repeat(40); // small, stays one chunk + let c = "cc aa bb\n".repeat(2000); + let inputs: Vec<&str> = vec![a.as_str(), b.as_str(), c.as_str()]; + + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial: Vec> = inputs + .iter() + .map(|s| ids(encode_one(tok, s).unwrap())) + .collect(); + + // Borrowed-slice sugar (copies in), bulk collect. + let collected = tok.encode(&inputs[..]).wait_for_completion().unwrap(); + let collected_ids: Vec> = collected.into_iter().map(ids).collect(); + assert_eq!(collected_ids, serial, "wait_for_completion != serial"); + + // Borrowed-slice sugar, streaming iterator: yields (index, result) in + // completion order; scattering by index must reproduce input order. + let mut streamed = vec![Vec::new(); inputs.len()]; + for (seq, r) in tok.encode(&inputs[..]) { + streamed[seq] = ids(r.unwrap()); + } + assert_eq!(streamed, serial, "streaming Iterator != serial"); + + // Owned returnable job (the default API): the handle escapes the call and + // is drained afterwards, both faces. + let owned_inputs: Vec = inputs.iter().map(|s| (*s).to_owned()).collect(); + let job = tok.encode(owned_inputs.clone()); + let owned: Vec> = job + .wait_for_completion() + .unwrap() + .into_iter() + .map(ids) + .collect(); + assert_eq!(owned, serial, "owned wait_for_completion != serial"); + + let job = tok.encode(owned_inputs.clone()); + let mut owned_streamed = vec![Vec::new(); owned_inputs.len()]; + for (seq, r) in job { + owned_streamed[seq] = ids(r.unwrap()); + } + assert_eq!(owned_streamed, serial, "owned Iterator != serial"); + + // Dropping a job early cancels cleanly (no hang, no panic) and the pool + // stays usable. + let job = tok.encode(owned_inputs); + drop(job); + let again = ids(tok.encode(a.clone()).into_single().unwrap()); + assert_eq!( + again, serial[0], + "pool must stay usable after a dropped job" + ); + } + + #[test] + fn encode_streams_and_matches_serial() { + assert_encode_matches_serial(&chunk_safe_pipeline()); + } + + /// FSM-family cuts on mixed prose land at whitespace boundaries and tile the + /// input exactly. + #[test] + fn stride_cuts_at_fsm_boundaries() { + let input = "word, another. 且つ 更に\n".repeat(2000); + let chunks = assert_strides_partition(&input, boundary_fsm, |bytes, start| { + assert!( + bytes[start].is_ascii_whitespace(), + "cut must land at whitespace, got {:#x}", + bytes[start] + ); + }); + assert!(chunks > 1, "expected several chunks, got {}", chunks); + } + + /// The newline-after-word cut is what parallelizes space-sparse scripts: a + /// document of CJK lines (no ASCII spaces) still splits, because a newline + /// following a CJK letter is a safe FSM cut. The previous character is a + /// multibyte letter, so this also exercises the `Cont` walk-back. + #[test] + fn stride_cuts_cjk_at_newlines() { + let input = "吾輩は猫である名前はまだ無い\n".repeat(3000); // ~120 KB, no spaces + let chunks = assert_strides_partition(&input, boundary_fsm, |bytes, start| { + assert_eq!(bytes[start], b'\n', "CJK cut must land at a newline"); + }); + assert!( + chunks > 1, + "space-sparse CJK input must still split, got {}", + chunks + ); + } + + /// The GPT-2 byte-level regex qualifies for `SpaceRun` raw cuts, both bare + /// and in the llama-shaped `Sequence[Span, None]`. A whitespace-containing + /// added token disables raw cutting only when it is `normalized` (matched + /// after normalization, so invisible to rung 0's raw peel); a raw/special one + /// is peeled by rung 0 first, so a stride can never bisect it. + #[test] + fn space_run_gating() { + use crate::AddedToken; + let split = || { + SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + ) + .unwrap() + }; + let make = |token: Option| { + let mut tok = Tokenizer::new(crate::models::bpe::BPE::default()); + tok.with_pre_tokenizer(Some(split())); + if let Some(t) = token { + tok.add_tokens([t]).unwrap(); + } + PipelineTokenizer::try_from(&tok).unwrap() + }; + assert!(make(None).stride_boundary().is_some()); + // Raw special with whitespace: peeled by rung 0, so raw cuts still qualify. + assert!( + make(Some(AddedToken::from(" extra", true))) + .stride_boundary() + .is_some(), + "a raw special is peeled by rung 0 and must not disable raw cuts" + ); + // Normalized added token with whitespace: rung 0's raw peel can't see it, + // so a stride could bisect it — raw cuts must be disabled. + assert!( + make(Some(AddedToken::from("aa bb", false).normalized(true))) + .stride_boundary() + .is_none(), + "whitespace inside a normalized added token must disable raw cuts" + ); + + // The llama-3 shape: Sequence[Span(known regex), ByteLevel(no regex) -> None]. + let seq = PipelinePreTokenizer::Sequence(PipelineSequence::new(vec![ + PipelinePreTokenizer::Split(split()), + PipelinePreTokenizer::None, + ])); + assert!(seq.stride_boundary().is_some()); + // A pre-tokenizer without a provable boundary must not qualify. + let no_boundary = PipelinePreTokenizer::Punctuation(Punctuation::default()); + assert!(no_boundary.stride_boundary().is_none()); + } + + /// A large input under the GPT-2 regex pre-tokenizer takes the `SpaceRun` + /// raw-cut path and stays id-identical to the serial encode. (WordPiece + /// stands in for the model; the boundary behavior under test is the + /// pre-tokenizer's.) + #[test] + fn space_run_split_matches_serial() { + use crate::models::wordpiece::WordPieceBuilder; + use ahash::AHashMap; + let vocab: AHashMap = ["aa", "bb", "cc", ""] + .iter() + .enumerate() + .map(|(i, w)| ((*w).to_string(), i as u32)) + .collect(); + let model = WordPieceBuilder::new() + .vocab(vocab) + .unk_token("".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(model); + tok.with_pre_tokenizer(Some( + SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + ) + .unwrap(), + )); + let tok = PipelineTokenizer::try_from(&tok).unwrap(); + assert!(matches!(tok.plan(), ParallelPlan::SplitRaw(_))); + + let big = "aa bb,cc! aa\tbb cc\n\n".repeat(2000); // ~44 KB, mixed runs + let mut tags = Vec::new(); + let chunks = (0..big.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES)) + .filter_map(|j| stride_range(&big, j, boundary_fsm, &mut tags)) + .count(); + assert!(chunks > 1, "expected the input to split, got {}", chunks); + + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial = ids(encode_one(&tok, &big).unwrap()); + let by_ref = ids(tok.encode(big.as_str()).into_single().unwrap()); + assert_eq!(by_ref, serial, "SpaceRun != serial"); + let owned = ids(tok.encode(big).into_single().unwrap()); + assert_eq!(owned, serial, "SpaceRun owned != serial"); + } + + /// Directly settles "rung 0 peels specials, so a stride can't split one" — it + /// holds only for *raw* (`normalized == false`) tokens. Rung 0's raw pass + /// (`SpecialSegmentIterator(.., false)`) searches the non-normalized vocab, so + /// a `normalized`-form token is invisible to it and stays inside a text + /// segment; it is peeled only later, *per stride*, by the worker. That is the + /// residual the `normalized_added_token_blocks_stride` gate exists for. + #[test] + fn rung0_raw_peel_misses_normalized_tokens() { + use crate::AddedToken; + let mut tok = Tokenizer::new(crate::models::bpe::BPE::default()); + tok.add_tokens([AddedToken::from("mask", false).normalized(true)]) + .unwrap(); + let pipe = PipelineTokenizer::try_from(&tok).unwrap(); + let av = &pipe.inner.added_vocabulary; + let input = "aa mask bb"; + + // Rung-0 pass (normalized = false): the token is NOT peeled. + let raw_special = SpecialSegmentIterator::new(input, av, false) + .any(|s| matches!(s, Segment::SpecialToken(_))); + assert!( + !raw_special, + "rung 0's raw pass must not see a normalized-form token" + ); + // Post-normalization pass (normalized = true): now it IS peeled. + let norm_special = SpecialSegmentIterator::new(input, av, true) + .any(|s| matches!(s, Segment::SpecialToken(_))); + assert!(norm_special, "the normalized pass must peel it"); + } + + /// Affix-strip gating after rung 0. A *raw* (`normalized == false`) + /// `lstrip`/`rstrip` token is peeled by rung 0 with its whitespace absorbed, + /// so raw cuts still qualify. A `normalized`-form affix token can't be peeled + /// before striding, and a whitespace-boundary cut could split its absorbed + /// run, so it must disable raw cuts. (The old `has_affix_strip` gate rejected + /// *both*, and re-scanned the vocab on every `encode`.) + #[test] + fn affix_strip_gating() { + use crate::AddedToken; + let make = |token: AddedToken| { + let mut tok = Tokenizer::new(crate::models::bpe::BPE::default()); + tok.with_pre_tokenizer(Some( + SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + ) + .unwrap(), + )); + tok.add_tokens([token]).unwrap(); + PipelineTokenizer::try_from(&tok).unwrap() + }; + assert!( + make( + AddedToken::from("", true) + .normalized(false) + .lstrip(true) + ) + .stride_boundary() + .is_some(), + "a raw lstrip token is peeled by rung 0 and must not disable raw cuts" + ); + assert!( + make( + AddedToken::from("", false) + .normalized(true) + .lstrip(true) + ) + .stride_boundary() + .is_none(), + "a normalized lstrip token must disable raw cuts" + ); + } + + /// A raw (`normalized == false`) `lstrip` token in a large input: rung 0 peels + /// it (absorbing the preceding whitespace) and the surrounding segments stride; + /// the result must stay byte-identical to the serial encode. + #[test] + fn raw_affix_token_strides_byte_identical() { + use crate::AddedToken; + let mut tok = Tokenizer::new(crate::models::bpe::BPE::default()); + tok.with_pre_tokenizer(Some( + SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + ) + .unwrap(), + )); + tok.add_tokens([AddedToken::from("", true) + .normalized(false) + .lstrip(true)]) + .unwrap(); + let pipe = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + matches!(pipe.plan(), ParallelPlan::SplitRaw(_)), + "raw affix token must keep SplitRaw" + ); + // Two ~20 KB segments around one lstrip token → both stride. + let big = format!("{} {}", "word ".repeat(4000), "word ".repeat(4000)); + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial = ids(encode_one(&pipe, &big).unwrap()); + let par = ids(pipe.encode(big.as_str()).into_single().unwrap()); + assert_eq!(par, serial, "raw affix striding != serial"); + } + + /// Gate refinement guard: llama-3 has 256 `normalized == false` reserved + /// specials (`<|reserved_special_token_0|>`), each carrying an internal + /// number-transition (`_→0`) cut. The old gate checked *all* added tokens and + /// would disqualify raw striding (silent `SplitAtModel` regression); the + /// refined gate only checks `normalized` tokens (rung 0 peels the raw ones), + /// so llama-3 must keep `SplitRaw` with the number-transition boundary. + #[test] + fn byte_level_raw_specials_keep_split_raw() { + let oracle = Tokenizer::from_file("../data/llama-3-tokenizer.json").unwrap(); + let tok = PipelineTokenizer::try_from(&oracle).unwrap(); + assert!( + matches!(tok.plan(), ParallelPlan::SplitRaw(_)), + "llama-3's raw-only specials must not disqualify SplitRaw" + ); + } + + /// A non-chunk-safe config: regex `Span` pre-tokenizer (unknown regex, so + /// no raw cut qualifies) + WordPiece model, so the plan must + /// escalate to [`ParallelPlan::SplitAtModel`]. Optionally a `Lowercase` + /// normalizer (safe per-char, but the pretok already forces the escalation) + /// and a `` special token. + fn split_at_model_pipeline(lowercase: bool) -> PipelineTokenizer { + use crate::models::wordpiece::WordPieceBuilder; + use crate::normalizers::utils::Lowercase; + use crate::AddedToken; + use ahash::AHashMap; + let vocab: AHashMap = ["aa", "bb", "cc", ".", "", ""] + .iter() + .enumerate() + .map(|(i, w)| ((*w).to_string(), i as u32)) + .collect(); + let model = WordPieceBuilder::new() + .vocab(vocab) + .unk_token("".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(model); + // Punctuation keeps whitespace inside runs, so no raw cut qualifies — + // and it needs no regex backend (the default build has none). + tok.with_pre_tokenizer(Some(Punctuation::default())); + if lowercase { + tok.with_normalizer(Some(Lowercase)).unwrap(); + } + tok.add_special_tokens([AddedToken::from("", true)]) + .unwrap(); + PipelineTokenizer::try_from(&tok).unwrap() + } + + /// SplitAtModel: a large single input under a non-chunk-safe config must + /// actually split into several parallel span-groups and stay id-identical + /// to the serial encode, through both the scoped and the owned paths. + #[test] + fn split_at_model_matches_serial() { + let tok = split_at_model_pipeline(false); + assert!(matches!(tok.plan(), ParallelPlan::SplitAtModel)); + + let big = "aa.bb.cc.".repeat(3000); // ~27 KB, punctuation-delimited + let prefix = tok.model_prefix_chunk(&big, &big).unwrap(); + let groups = prefix + .units + .iter() + .filter(|u| matches!(u, PrefixUnit::Group { .. })) + .count(); + assert!(groups > 1, "expected several span-groups, got {}", groups); + assert!(prefix.norm_buf.is_none(), "no normalizer -> no owned buf"); + + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial = ids(encode_one(&tok, &big).unwrap()); + let by_ref = ids(tok.encode(big.as_str()).into_single().unwrap()); + assert_eq!(by_ref, serial, "SplitAtModel != serial"); + let owned = ids(tok.encode(big).into_single().unwrap()); + assert_eq!(owned, serial, "owned SplitAtModel != serial"); + } + + /// SplitAtModel with a rewriting normalizer (uppercase input + `Lowercase` + /// forces `Cow::Owned` → the `Norm` buffer path) and an interleaved special + /// token (preresolved units must keep their position in the stream). + #[test] + fn split_at_model_normalizer_and_specials_match_serial() { + let tok = split_at_model_pipeline(true); + assert!(matches!(tok.plan(), ParallelPlan::SplitAtModel)); + + let half = "AA.BB.cc.".repeat(2000); + let big = format!("{half}{}", "CC.aa.BB.".repeat(2000)); + // Rung 0 peels `` in the builder; the per-segment model prefix runs on + // special-free text and (with `Lowercase`) must own its normalized buffer. + let prefix = tok.model_prefix_chunk(&half, &half).unwrap(); + assert!( + prefix.norm_buf.is_some(), + "lowercased text must live in the owned norm buf" + ); + + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial = ids(encode_one(&tok, &big).unwrap()); + let by_ref = ids(tok.encode(big.as_str()).into_single().unwrap()); + assert_eq!(by_ref, serial, "parallel != serial (norm + specials)"); + let owned = ids(tok.encode(big).into_single().unwrap()); + assert_eq!(owned, serial, "owned != serial (norm + specials)"); + } + + /// A boundary-sensitive normalizer (`Strip`) makes the config not chunk-safe, so a large + /// input stays whole (one chunk) — no unsafe splitting. + #[test] + fn unsafe_normalizer_config_does_not_split() { + use crate::models::wordlevel::WordLevelBuilder; + use crate::normalizers::strip::Strip; + use crate::pre_tokenizers::whitespace::WhitespaceSplit; + use ahash::AHashMap; + let vocab: AHashMap = [("aa", 0u32), ("", 1)] + .iter() + .map(|(w, i)| ((*w).to_string(), *i)) + .collect(); + let model = WordLevelBuilder::new() + .vocab(vocab) + .unk_token("".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(model); + tok.with_pre_tokenizer(Some(WhitespaceSplit)); + tok.with_normalizer(Some(Strip::new(true, true))).unwrap(); + let pipe = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + pipe.stride_boundary().is_none(), + "Strip normalizer must disable raw-text chunking" + ); + // WordLevel model: the plan cannot escalate to SplitAtModel either. + assert!( + matches!(pipe.plan(), ParallelPlan::WholeInput), + "Strip + WordLevel must fall to batch-level parallelism only", + ); + } + + /// Rung 0 (always first): a config no lower rung can cut — `WordLevel` (so no + /// `SplitAtModel`) under a boundary-hostile `Strip` normalizer (so no + /// `SplitRaw`/`SplitNormalized`) picks the `WholeInput` sub-rung. But a single + /// large input dense with special tokens still parallelizes, because rung 0 + /// wraps every plan: the builder peels the specials (preresolved) and hands + /// each inter-special text segment to a worker. Must stay byte-identical to + /// the serial encode through both the borrowed and owned handles. + #[test] + fn rung0_specials_parallelize_whole_input() { + use crate::models::wordlevel::WordLevelBuilder; + use crate::normalizers::strip::Strip; + use crate::pre_tokenizers::whitespace::WhitespaceSplit; + use crate::AddedToken; + use ahash::AHashMap; + let vocab: AHashMap = [("aa", 0u32), ("bb", 1), ("cc", 2), ("", 3)] + .iter() + .map(|(w, i)| ((*w).to_string(), *i)) + .collect(); + let model = WordLevelBuilder::new() + .vocab(vocab) + .unk_token("".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(model); + tok.with_pre_tokenizer(Some(WhitespaceSplit)); + tok.with_normalizer(Some(Strip::new(true, true))).unwrap(); + tok.add_special_tokens([AddedToken::from("", true)]) + .unwrap(); + let tok = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + matches!(tok.plan(), ParallelPlan::WholeInput), + "Strip + WordLevel sub-rung must be WholeInput (rung 0 does the splitting)" + ); + + // ~24 KB, ~2000 interspersed specials → thousands of worker units. + let big = "aa bb cc ".repeat(2000); + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial = ids(encode_one(&tok, &big).unwrap()); + let by_ref = ids(tok.encode(big.as_str()).into_single().unwrap()); + assert_eq!(by_ref, serial, "rung-0 split (borrowed) != serial"); + let owned = ids(tok.encode(big.clone()).into_single().unwrap()); + assert_eq!(owned, serial, "rung-0 split (owned) != serial"); + } + + /// Rung 2 (`SplitNormalized`): a non-boundary-preserving normalizer + /// (`Prepend`) rules out `SplitRaw`, but the `WhitespaceSplit` pre-tokenizer + /// can cut the normalized text. The config serial-normalizes then + /// parallelizes pre-tokenize + model, and must stay byte-identical to the + /// serial encode (through both the borrowed and owned handles). + #[test] + fn split_normalized_matches_serial() { + use crate::models::wordpiece::WordPieceBuilder; + use crate::normalizers::prepend::Prepend; + use crate::pre_tokenizers::whitespace::WhitespaceSplit; + use ahash::AHashMap; + let vocab: AHashMap = ["▁aa", "aa", "bb", "cc", ""] + .iter() + .enumerate() + .map(|(i, w)| ((*w).to_string(), i as u32)) + .collect(); + let model = WordPieceBuilder::new() + .vocab(vocab) + .unk_token("".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(model); + tok.with_pre_tokenizer(Some(WhitespaceSplit)); + tok.with_normalizer(Some(Prepend::new("▁".to_string()))) + .unwrap(); + let tok = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + matches!(tok.plan(), ParallelPlan::SplitNormalized(_)), + "Prepend (unsafe normalizer) + WhitespaceSplit + WordPiece must pick rung 2" + ); + let big = "aa bb cc\n".repeat(4000); // ~36 KB, forces striding + let ids = |v: Vec| v.iter().map(|t| t.id).collect::>(); + let serial = ids(encode_one(&tok, &big).unwrap()); + let by_ref = ids(tok.encode(big.as_str()).into_single().unwrap()); + assert_eq!(by_ref, serial, "SplitNormalized (borrowed) != serial"); + let owned = ids(tok.encode(big.clone()).into_single().unwrap()); + assert_eq!(owned, serial, "SplitNormalized (owned) != serial"); + } + struct FixedMatcher(Vec<((usize, usize), u32)>); impl PipelinePatternMatcher for FixedMatcher { fn extract_next( diff --git a/tokenizers/tk-encode/src/tokenizer/pool.rs b/tokenizers/tk-encode/src/tokenizer/pool.rs new file mode 100644 index 000000000..ed3e479ea --- /dev/null +++ b/tokenizers/tk-encode/src/tokenizer/pool.rs @@ -0,0 +1,155 @@ +//! Process-global worker pool for the parallel encode path: a lazily-built, +//! library-private `rayon::ThreadPool`. +//! +//! Scheduling itself lives in `pipeline.rs` as the shared-cursor job model — +//! drainer tasks and the consuming thread claim chunks off one atomic cursor +//! (caller-assist), which benching showed is where the throughput wins live; a +//! bespoke thread pool added nothing over rayon once the job model carried the +//! assist (see `PARALLEL_RUNTIME_DESIGN.md` §7 for the removal decision and +//! numbers; the handrolled pool lives in git history if ever needed). +//! +//! Fork safety: a `pthread_atfork` child handler bumps [`POOL_GEN`]; the +//! dispatch path compares generations and abandons a stale pool *without +//! touching it* — its threads died with the fork and its locks may be held by +//! ghost threads, so it is intentionally leaked (bounded: one pool per fork) — +//! then lazily builds a fresh one. Callers that cannot get a pool (lock +//! unavailable, spawn failure, parallelism disabled) fall back to the inline +//! path, so a fork can cost throughput but never liveness. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Mutex, OnceLock}; + +use crate::utils::parallelism::{ + get_parallelism, mark_parallelism_used, num_threads_override, NUM_THREADS_ENV_VARIABLE, +}; + +/// Fork generation: bumped in the child by the `pthread_atfork` handler (a +/// single atomic store — async-signal-safe). Pools stamped with an older +/// generation are stale and must never be touched again. +static POOL_GEN: AtomicU64 = AtomicU64::new(0); + +#[cfg(unix)] +fn register_fork_handler() { + static REGISTERED: std::sync::Once = std::sync::Once::new(); + REGISTERED.call_once(|| { + unsafe extern "C" fn child_after_fork() { + POOL_GEN.fetch_add(1, Ordering::SeqCst); + } + // Registration failure means fork detection is off — same behavior as + // before this module existed; nothing to do about it here. + unsafe { + let _ = libc::pthread_atfork(None, None, Some(child_after_fork)); + } + }); +} + +#[cfg(not(unix))] +fn register_fork_handler() {} + +/// Abandon the current pool (without touching it) so the next dispatch builds +/// a fresh one — used by `set_num_threads` to apply a new size, and by the +/// fork handler through the same generation mechanism. +pub(crate) fn invalidate() { + POOL_GEN.fetch_add(1, Ordering::SeqCst); +} + +/// Worker count, per the precedence in [`crate::utils::parallelism`]: +/// `set_num_threads` (live) > `TOKENIZERS_NUM_THREADS` (cached at first use) > +/// the machine's available parallelism. +fn num_threads() -> usize { + if let Some(n) = num_threads_override() { + return n; + } + static N: OnceLock = OnceLock::new(); + *N.get_or_init(|| { + std::env::var(NUM_THREADS_ENV_VARIABLE) + .ok() + .and_then(|v| v.parse().ok()) + .filter(|&n| n > 0) + .unwrap_or_else(|| { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + }) + }) +} + +/// The pool for one encode dispatch, or `None` when encoding must run inline: +/// parallelism disabled, thread spawn failure, or the cell lock unavailable +/// (held mid-fork by a ghost thread, or poisoned) — `try_lock` is what +/// guarantees this path can never hang. Cheap on the happy path: one atomic +/// load and one uncontended `try_lock`. +pub(crate) fn rayon() -> Option<&'static rayon::ThreadPool> { + /// The generation-stamped, intentionally-leaked pool. + static CELL: Mutex> = Mutex::new(None); + + if !get_parallelism() { + return None; + } + register_fork_handler(); + let mut guard = CELL.try_lock().ok()?; + let generation = POOL_GEN.load(Ordering::Acquire); + if let Some((pool, stamp)) = *guard { + if stamp == generation { + return Some(pool); + } + // Stale after a fork: abandon the old pool without touching it (its + // threads died with the fork; its locks may be held by ghosts) and + // build a fresh one for this generation. + } + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(num_threads()) + .thread_name(|i| format!("tk-encode-{i}")) + .build() + .ok()?; + let pool: &'static rayon::ThreadPool = Box::leak(Box::new(pool)); + *guard = Some((pool, generation)); + // Feed the Python bindings' fork guard: an *unconfigured* forked child of + // this process goes serial (v1 semantics, prevents DataLoader thread + // oversubscription); explicitly configured parallelism survives forks at + // full speed via the generation rebuild. + mark_parallelism_used(); + Some(pool) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// One sequential test: both scenarios mutate the process-global pool + /// generation, so running them as separate (concurrent) tests would flake. + /// + /// (a) A fork-generation bump abandons the current pool and lazily builds a + /// fresh one; a stable generation reuses the same pool. Building marks + /// parallelism as used (fork-guard telemetry). + /// (b) `set_num_threads` takes priority over env/default and applies + /// *live*: the existing pool is abandoned and the next dispatch builds one + /// at the requested size; `0` resets to env-or-default resolution. + #[test] + fn pool_generation_and_num_threads_control() { + use crate::utils::parallelism::set_num_threads; + + // (a) fork generation + let a = rayon().unwrap() as *const rayon::ThreadPool; + let a2 = rayon().unwrap() as *const rayon::ThreadPool; + assert_eq!(a, a2, "stable generation must reuse the pool"); + POOL_GEN.fetch_add(1, Ordering::SeqCst); + let b = rayon().unwrap() as *const rayon::ThreadPool; + assert_ne!(a, b, "generation bump must rebuild the pool"); + let b2 = rayon().unwrap() as *const rayon::ThreadPool; + assert_eq!(b, b2); + assert!(crate::utils::parallelism::has_parallelism_been_used()); + + // (b) live worker-count control + let default_threads = rayon().unwrap().current_num_threads(); + let target = if default_threads == 3 { 2 } else { 3 }; + set_num_threads(target); + assert_eq!( + rayon().unwrap().current_num_threads(), + target, + "setter must override env/default and rebuild live" + ); + set_num_threads(0); // reset to env-or-default + assert_eq!(rayon().unwrap().current_num_threads(), default_threads); + } +} diff --git a/tokenizers/tk-encode/src/utils/mod.rs b/tokenizers/tk-encode/src/utils/mod.rs index 1f2ed7802..29d4d9fc1 100644 --- a/tokenizers/tk-encode/src/utils/mod.rs +++ b/tokenizers/tk-encode/src/utils/mod.rs @@ -16,7 +16,7 @@ pub use no_regex::SysRegex; // Recognize known GPT pre-tokenization regexes and route them to atomsplit's native (unrolled) FSM. mod unrolled_regex; -pub use unrolled_regex::{gpt_fsm, is_deepseek, GptFsm, GptFsmPattern}; +pub use unrolled_regex::{gpt_fsm, is_deepseek, is_deepseek_member, GptFsm, GptFsmPattern}; pub mod byte_level; pub mod iter; diff --git a/tokenizers/tk-encode/src/utils/parallelism.rs b/tokenizers/tk-encode/src/utils/parallelism.rs index ea2fd331a..98eabe393 100644 --- a/tokenizers/tk-encode/src/utils/parallelism.rs +++ b/tokenizers/tk-encode/src/utils/parallelism.rs @@ -1,32 +1,58 @@ +//! Process-wide parallelism controls. //! -//! This module defines helpers to allow optional Rayon usage. +//! The **programmatic setters are the primary interface**; the environment +//! variables are the backwards-compatibility channel. Precedence, highest +//! first: //! +//! 1. [`set_parallelism`] / [`set_num_threads`] — always live; a threads +//! change rebuilds the encode pool on its next use. +//! 2. `TOKENIZERS_PARALLELISM` / `TOKENIZERS_NUM_THREADS` — read **once, at +//! first use** (set them before the first encode; to change behavior after +//! that, use the setters). +//! 3. Sound defaults: parallelism on, one worker per available core +//! (cgroup-aware via `std::thread::available_parallelism`). +//! +//! This module also defines the `Maybe*` helpers the legacy paths use for +//! optional Rayon usage. use rayon::iter::IterBridge; use rayon::prelude::*; use rayon_cond::CondIterator; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU8; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::sync::OnceLock; // Re-export rayon current_num_threads pub use rayon::current_num_threads; pub const ENV_VARIABLE: &str = "TOKENIZERS_PARALLELISM"; +pub const NUM_THREADS_ENV_VARIABLE: &str = "TOKENIZERS_NUM_THREADS"; static USED_PARALLELISM: AtomicBool = AtomicBool::new(false); static PARALLELISM: AtomicU8 = AtomicU8::new(0); +/// Programmatic worker-count override; `0` means "not set". +static NUM_THREADS: AtomicUsize = AtomicUsize::new(0); -/// Check if the TOKENIZERS_PARALLELISM env variable has been explicitly set +/// Check if parallelism has been explicitly configured, via either setter or +/// env var. (The Python bindings' fork guard leaves configured processes +/// alone and only forces *unconfigured* forked children serial.) pub fn is_parallelism_configured() -> bool { std::env::var(ENV_VARIABLE).is_ok() || get_override_parallelism().is_some() } -/// Check if at some point we used a parallel iterator +/// Check if parallelism has been used at some point (a parallel iterator, or +/// the encode worker pool). pub fn has_parallelism_been_used() -> bool { USED_PARALLELISM.load(Ordering::SeqCst) } +/// Record that parallelism was exercised (feeds the fork guard above). +pub(crate) fn mark_parallelism_used() { + USED_PARALLELISM.store(true, Ordering::SeqCst); +} + /// Get internally set parallelism fn get_override_parallelism() -> Option { match PARALLELISM.load(Ordering::SeqCst) { @@ -48,19 +74,43 @@ fn get_env_parallelism() -> bool { } } +/// Whether encoding may parallelize: [`set_parallelism`] wins, then the +/// `TOKENIZERS_PARALLELISM` env var (cached at first use), then on. pub fn get_parallelism() -> bool { if let Some(parallel) = get_override_parallelism() { parallel } else { - get_env_parallelism() + // The env var is the backwards-compat channel, frozen at first use; + // runtime changes go through `set_parallelism`. + static ENV: OnceLock = OnceLock::new(); + *ENV.get_or_init(get_env_parallelism) } } -/// Set the value for `TOKENIZERS_PARALLELISM` for the current process +/// Enable or disable parallelism for the current process. Always live; takes +/// priority over the `TOKENIZERS_PARALLELISM` env var. pub fn set_parallelism(val: bool) { PARALLELISM.store(if val { 2 } else { 1 }, Ordering::SeqCst); } +/// Set the number of worker threads for the parallel encode pool. Takes +/// priority over the `TOKENIZERS_NUM_THREADS` env var; `0` resets to the +/// env-var-or-default resolution. Always live: an existing pool is abandoned +/// and rebuilt at the new size on the next encode (the abandoned pool is +/// leaked — bounded by the number of calls, so don't call this in a loop). +pub fn set_num_threads(num_threads: usize) { + NUM_THREADS.store(num_threads, Ordering::SeqCst); + crate::tokenizer::pool::invalidate(); +} + +/// The programmatic worker-count override, if one was set. +pub(crate) fn num_threads_override() -> Option { + match NUM_THREADS.load(Ordering::SeqCst) { + 0 => None, + n => Some(n), + } +} + /// Allows to convert into an iterator that can be executed either parallelly or serially. /// /// The choice is made according to the currently set `TOKENIZERS_PARALLELISM` environment variable. diff --git a/tokenizers/tk-encode/src/utils/unrolled_regex.rs b/tokenizers/tk-encode/src/utils/unrolled_regex.rs index 212240a03..947243c14 100644 --- a/tokenizers/tk-encode/src/utils/unrolled_regex.rs +++ b/tokenizers/tk-encode/src/utils/unrolled_regex.rs @@ -106,6 +106,14 @@ pub fn is_deepseek(p0: &str, p1: &str, p2: &str) -> bool { p0 == DS_NUM && p1 == DS_CJK && p2 == DS_BIG } +/// True iff `pattern` is one of deepseek's three member regexes. None of them is a standalone GPT +/// FSM, so they'd otherwise demand a system-regex backend at construction — but they only ever run +/// inside a recognized deepseek `Sequence`, which `fsm_deepseek` drives regex-less. Lets such a +/// member `Split` build without a backend (its standalone `pre_tokenize` still errors, as before). +pub fn is_deepseek_member(pattern: &str) -> bool { + pattern == DS_NUM || pattern == DS_CJK || pattern == DS_BIG +} + #[cfg(test)] mod tests { use super::*; diff --git a/tokenizers/tk-encode/tests/pipeline_oracle.rs b/tokenizers/tk-encode/tests/pipeline_oracle.rs index 7c8781f1c..655eff19a 100644 --- a/tokenizers/tk-encode/tests/pipeline_oracle.rs +++ b/tokenizers/tk-encode/tests/pipeline_oracle.rs @@ -7,8 +7,20 @@ use std::convert::TryFrom; -use tk_encode::pipeline::PipelineTokenizer; -use tk_encode::Tokenizer; +use tk_encode::pipeline::{EncodeHandle, PipelineToken, PipelineTokenizer}; +use tk_encode::{Result, Tokenizer}; + +/// Test-only convenience: drain a single-input handle to its one result (the +/// pipeline's public API is the streaming `Iterator` / `wait_for_completion`). +trait IntoSingle { + fn into_single(self) -> Result>; +} +impl IntoSingle for EncodeHandle { + fn into_single(self) -> Result> { + self.wait_for_completion() + .map(|all| all.into_iter().next().unwrap_or_default()) + } +} fn load(corpus: &str) -> (Tokenizer, PipelineTokenizer, String) { let oracle = Tokenizer::from_file("../data/bert-wiki.json").unwrap(); @@ -41,7 +53,7 @@ fn check_chunks(corpus: &str, target_bytes: usize) { for chunk in make_chunks(&text, target_bytes) { let expected = oracle.encode(chunk.as_str(), false).unwrap(); let got: Vec = pipeline - .encode(&chunk, false) + .encode(&chunk).into_single() .unwrap() .iter() .map(|t| t.id) @@ -55,6 +67,210 @@ fn check_chunks(corpus: &str, target_bytes: usize) { } } +/// Batch parallel floor: a multi-sequence `encode` (fanned out on the pool) must +/// be result-identical to both the serial pipeline `encode` (order-preserving +/// parallelism) and the reference oracle. Also checks the two `EncodeHandle` +/// consumers agree: `wait_for_completion` (bulk collect) and `Iterator` +/// (input-ordered streaming). Runs the whole corpus as one batch so multiple +/// documents land on different worker threads. +fn check_batch(corpus: &str, target_bytes: usize) { + let (oracle, pipeline, text) = load(corpus); + let chunks = make_chunks(&text, target_bytes); + let refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + + let batch = pipeline + .encode(&refs[..]).wait_for_completion() + .unwrap(); + assert_eq!(batch.len(), chunks.len(), "batch length mismatch"); + + // Iterator face yields (index, result) in completion order; scattering by + // index must reproduce the same per-input results. + let mut streamed: Vec> = vec![Vec::new(); refs.len()]; + for (seq, r) in pipeline.encode(&refs[..]) { + streamed[seq] = r.unwrap().iter().map(|t| t.id).collect(); + } + + for (i, (chunk, got)) in chunks.iter().zip(&batch).enumerate() { + let got_ids: Vec = got.iter().map(|t| t.id).collect(); + + // wait_for_completion == Iterator + assert_eq!(got_ids, streamed[i], "wait_for_completion != Iterator"); + + // batch == serial pipeline + let serial: Vec = pipeline + .encode(chunk).into_single() + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!( + got_ids, serial, + "batch != serial on {:?}", + chunk.chars().take(80).collect::(), + ); + + // batch == reference oracle + let expected = oracle.encode(chunk.as_str(), false).unwrap(); + assert_eq!( + expected.get_ids(), + got_ids.as_slice(), + "batch != oracle on {:?}", + chunk.chars().take(80).collect::(), + ); + } +} + +/// Exercises the cost gate's *inline* branch: a multi-sequence `encode` whose +/// summed size is below the fan-out threshold runs on the calling thread (not +/// the pool) and must still match the oracle. The corpus `check_batch` always +/// clears the threshold and fans out, so this covers the other side. +fn check_small_batch(corpus: &str) { + let (oracle, pipeline, text) = load(corpus); + // A handful of short lines — well under the fan-out byte threshold. + let inputs: Vec<&str> = text + .lines() + .filter(|l| !l.trim().is_empty()) + .take(4) + .collect(); + if inputs.len() < 2 { + return; + } + let got = pipeline + .encode(&inputs[..]).wait_for_completion() + .unwrap(); + assert_eq!(got.len(), inputs.len()); + for (input, tokens) in inputs.iter().zip(&got) { + let expected = oracle.encode(*input, false).unwrap(); + let ids: Vec = tokens.iter().map(|t| t.id).collect(); + assert_eq!( + expected.get_ids(), + ids.as_slice(), + "small-batch (inline) mismatch on {input:?}", + ); + } +} + +/// Intra-sequence chunking: a single large document is encoded through the fused stride +/// path — split at newline boundaries into chunks, each run through the full pipeline in +/// parallel, and concatenated in order. The result must equal the oracle (which, combined +/// with the serial `check_chunks` above, transitively confirms chunked == serial == oracle). +/// The bert-wiki config is chunk-safe (per-char normalizer + whitespace-delimiting +/// pre-tokenizer), so this exercises the split; an unsafe config would stay whole. +fn check_intra_seq(corpus: &str) { + let (oracle, pipeline, text) = load(corpus); + // One document large enough to be split into several chunks, but capped so the oracle + // reference encode stays cheap. + let mut doc = String::new(); + for line in text.lines().filter(|l| !l.trim().is_empty()) { + if !doc.is_empty() { + doc.push('\n'); + } + doc.push_str(line); + if doc.len() >= 64 * 1024 { + break; + } + } + if doc.len() < 16 * 1024 { + return; // corpus too small to exercise the parallel path + } + let got: Vec = pipeline + .encode(doc.as_str()).into_single() + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + let expected = oracle.encode(doc.as_str(), false).unwrap(); + assert_eq!( + expected.get_ids(), + got.as_slice(), + "intra-seq parallel encode != oracle", + ); +} + +/// Byte-level BPE (llama-3): the `SpaceRun` raw-cut path must be id-identical +/// to the reference oracle across large real documents. Exercised over English +/// (dense with non-ws→space cut candidates) and Japanese (multibyte seams — +/// cuts may only land beside ASCII), through both encode paths. +fn check_llama3(corpus: &str) { + let oracle = Tokenizer::from_file("../data/llama-3-tokenizer.json").unwrap(); + let pipeline = PipelineTokenizer::try_from(&oracle).unwrap(); + let text = std::fs::read_to_string(corpus).unwrap(); + let mut doc = String::new(); + for line in text.lines().filter(|l| !l.trim().is_empty()) { + if !doc.is_empty() { + doc.push('\n'); + } + doc.push_str(line); + if doc.len() >= 256 * 1024 { + break; + } + } + if doc.len() < 32 * 1024 { + return; + } + let expected = oracle.encode(doc.as_str(), false).unwrap(); + let single: Vec = pipeline + .encode(doc.as_str()).into_single() + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!( + expected.get_ids(), + single.as_slice(), + "llama-3 single-doc SpaceRun encode != oracle" + ); + let owned: Vec = pipeline + .encode(doc) + .into_single() + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!( + expected.get_ids(), + owned.as_slice(), + "llama-3 owned SpaceRun encode != oracle" + ); +} + +#[test] +fn llama3_intra_seq_english() { + check_llama3("../data/big.txt"); +} + +#[test] +fn llama3_intra_seq_japanese() { + check_llama3("../data/unigram_wagahaiwa_nekodearu.txt"); +} + +/// Byte-level BPE on a whitespace-free document (base64-like) must stay +/// byte-identical to the oracle. This input has no whitespace, so its only safe +/// stride cuts are the byte-level letter↔number transitions — re-enabled now +/// that rung 0 always peels specials first, making a mid-special cut impossible. +/// So this exercises the number-transition boundary striding (not a whole-input +/// fallback) and pins it id-identical to tiktoken. +#[test] +fn llama3_intra_seq_whitespace_free() { + let oracle = Tokenizer::from_file("../data/llama-3-tokenizer.json").unwrap(); + let pipeline = PipelineTokenizer::try_from(&oracle).unwrap(); + let doc = "aGVsbG8x3Wb29ybGQ7abc123def456ghi789jkl0mno".repeat(5000); + assert!(!doc.bytes().any(|b| b.is_ascii_whitespace())); + let expected = oracle.encode(doc.as_str(), false).unwrap(); + let single: Vec = pipeline + .encode(doc.as_str()) + .into_single() + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!( + expected.get_ids(), + single.as_slice(), + "whitespace-free byte-level encode != oracle" + ); +} + macro_rules! corpus_tests { ($($name:ident => $file:literal),* $(,)?) => { $( @@ -64,9 +280,25 @@ macro_rules! corpus_tests { super::check_chunks($file, 1024); } #[test] + fn intra_seq() { + super::check_intra_seq($file); + } + #[test] fn chunks_10kb() { super::check_chunks($file, 10 * 1024); } + #[test] + fn batch_1kb() { + super::check_batch($file, 1024); + } + #[test] + fn batch_10kb() { + super::check_batch($file, 10 * 1024); + } + #[test] + fn small_batch_inline() { + super::check_small_batch($file); + } } )* };