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..effbc04ab --- /dev/null +++ b/bindings/python/src/pipeline.rs @@ -0,0 +1,172 @@ +//! 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: we only ever read from the underlying buffers, +// and the lifetime of the buffers is guaranteed by storing +// the Py in _owners, which is kept alive as long as +// a reference is held, because Python refcount. +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..fbe3e3b76 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}; @@ -55,7 +55,7 @@ const MEM_CHUNKS_PER_FIXTURE: usize = 2; // `ADDED_SPECIAL` (normalized:false) is matched on the raw pass; `ADDED_NORMALIZED` // (normalized:true) on the normalized pass. The strings are distinctive markers that do // not occur in the language/modality corpora, so they leave those results untouched. The -// `added_*` fixtures are built from exactly these strings (see the dataset FIXTURES.md). +// `added_*` fixtures are built from exactly these strings. const ADDED_SPECIAL: &[&str] = &["<|xs0|>", "<|xs1|>", "<|xs2|>", "<|xs3|>", "<|xs4|>"]; const ADDED_NORMALIZED: &[&str] = &[ "widgetron", @@ -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..58f3ed248 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1,8 +1,41 @@ -use std::cell::RefCell; +//! 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 — each input the moment it finishes (and *assisting* — the +//! caller drains pending work instead of idling) — and +//! [`EncodeHandle::wait_for_completion`], which blocks and returns all results +//! collected back into 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. There is no result +//! channel: each worker writes its unit into a shared write-once `Slot` and +//! decrements the input's `pending` counter; the consumer reads an input's slots +//! once `pending` hits zero. 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`. +//! +//! **Where the work is split.** Specials are peeled first (the outermost stage, +//! so a cut there is always safe; ~0.1% of encode, and it keeps strided segments +//! special-free). Each special-free **segment** is then split at the earliest +//! pipeline stage its config allows (the `ParallelPlan` ladder). Earlier = more +//! parallel; 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, UnsafeCell}; use std::convert::TryInto; +use std::ops::Range; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Arc; + +use rayon::prelude::*; 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 +58,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 +96,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 +141,35 @@ impl PreTokenizer for PipelinePreTokenizer { } } +impl PipelinePreTokenizer { + /// Return the [`StrideBoundary`] predicate for this pre-tokenizer, if it has one. + /// The predicate determines the safe cutting points in the input, such that: + /// ```text + /// let (a, b) = input.split(input.len() / 2); + /// pre_tokenize(input) == pre_tokenize(a) + pre_tokenize(b) + /// ``` + 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 +218,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 +235,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 +267,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, @@ -266,15 +318,853 @@ 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. +/// Main tokenizer struct +/// +/// Thread-safe because immutable after construction +/// and cheap clones thanks to arc wrapped internals +#[derive(Clone)] pub struct PipelineTokenizer { + inner: Arc, +} + +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 `Raw`/`Normalized`. Only + /// `normalized` tokens count — the special split peels raw-matched ones + /// before any striding. Precomputed at build; `plan()` reads it per `encode`, + /// so it must not re-scan the vocab. + normalized_added_token_blocks_stride: bool, +} + +// comptime verification that PipelineTokenizer 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). +#[derive(Default)] +pub(crate) struct EncodeState { + pre_tokens: Vec, + /// Model-specific heap buffers. Deliberately not cleared by `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 +/// `Pretokenized`, 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) -> 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)? + }; + 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) { + 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 char is a safe cut — these +/// splitters drop whitespace and never emit a token spanning it, so each side +/// drops its own run. The cut joins the right chunk and stays on a char boundary +/// (a ws lead is `WS`, its continuation `Cont`). +fn boundary_at_whitespace(window_tags: &[u8]) -> Option { + (1..window_tags.len()).find(|&i| in_mask(window_tags[i], mask::WS)) +} + +/// Side-independent cut for the atomsplit-FSM byte-level families (gpt2 / cl100k +/// / o200k / deepseek), with the cut char `cur` joining the right chunk. Four +/// safe `prev`→`cur` transitions (each verified branch-by-branch against +/// `atomsplit::regexes`): non-ws→whitespace, word ([`NEWLINE_PREV`])→newline, +/// non-number ([`NUM_START_PREV`])→number, number→letter. The number transitions +/// can fall inside a special, so they are sound *only* because specials are +/// peeled before striding (a strided segment is special-free); normalized-form +/// added tokens still need the stride guard. +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 `text` contains any `boundary` cut at all, deciding within a bounded prefix. +/// +/// Used by [`PipelineTokenizer::plan`] to reject a cut-based plan for input that has no cuts. +/// Only the first few strides are examined: a text whose opening `PROBE` bytes hold no cut is +/// treated as uncuttable, which is right for the case this exists for (Chinese/Japanese prose +/// has none anywhere) and merely gives up some parallelism in the pathological case of a +/// document that is boundary-free only at the front. Bounded so the probe stays O(1) in the +/// input: it runs once per `encode`, before any work is dispatched. +fn cut_exists(text: &str, boundary: StrideBoundary) -> bool { + const PROBE: usize = 4 * PARALLEL_MIN_TOTAL_BYTES; + // `boundary_in_window` reads one byte of left context (`block_lo - 1`), so `lo` must be + // at least 1 -- passing 0 underflows and hangs the char-boundary walk-back. Byte 0 is + // never a cut anyway: stride 0 starts at 0 unconditionally. + if text.len() < 2 { + return false; + } + let hi = text.len().min(PROBE); + let hi = (1..=hi).rev().find(|&i| text.is_char_boundary(i)).unwrap_or(1); + boundary_in_window(text, 1, hi, boundary).is_some() +} + +/// 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). The special split peels raw-matched tokens before 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) -> Option { + // Per-thread classify scratch: overwritten (clear + resize + classify) on + // every block before it is read, so nothing carries across calls — it is + // kept only to reuse the allocation. Distinct from the pre-tokenizer's own + // classify scratch (`classify_into_spans`); the two never share a buffer. + thread_local! { + static TAGS: RefCell> = const { RefCell::new(Vec::new()) }; + } + const BLOCK: usize = 1024; + debug_assert!(1 <= lo && lo <= hi && hi <= text.len()); + TAGS.with(|cell| { + let tags = &mut *cell.borrow_mut(); + 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 + }) +} + +/// How one special-free segment is split, chosen per config (the special split +/// already peeled the specials). The pipeline splits at the earliest stage the +/// config allows; earlier is more parallel. Never wrong, only less parallel. +#[derive(Clone, Copy, Debug)] +enum ParallelPlan { + /// Stride the raw text; each chunk runs the full pipeline (`encode_one_with`). + /// Needs a boundary-preserving normalizer. + Raw(StrideBoundary), + /// Normalize the segment serially, then stride the normalized buffer and + /// parallelize pre-tokenize + model (`encode_normalized`). The serial + /// normalize takes the normalizer out of the equation, so this needs only a + /// splittable pre-tokenizer. Used when the normalizer rules out `Raw`. + Normalized(StrideBoundary), + /// Normalize + pre-tokenize serially, then parallelize the model over + /// ~`PARALLEL_MIN_TOTAL_BYTES`-sized span-groups. Valid for any config. + Pretokenized, + /// No intra-input split. + Whole, +} + +/// One special-free segment reduced to model-ready work by the `Pretokenized` +/// plan: normalize, frame normalized specials, and pre-tokenize are done; the +/// model stage remains, packaged as span-groups. Index/range-based, no borrows. +struct PretokenizedSegment { + /// Ordered pipeline units for this segment. + units: Vec, + /// The owned normalized segment ([`SegmentText::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 SegmentUnit { + /// A matched special/added token + Special(u32), + /// A group of consecutive pre-token spans over one text segment. + Group { + text: SegmentText, + spans: Range, + }, +} + +/// Where a span-group's text lives. Spans are relative to the resolved text. +enum SegmentText { + /// 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()); +} + +/// A write-once result cell for one unit. Exactly one worker writes it — the one +/// the cursor handed the unit to — then publishes by decrementing the input's +/// `pending` counter. The consumer reads it only after seeing `pending == 0` +/// (acquire), which happens-after every write. That single-writer / published- +/// read contract is what makes the `unsafe impl Sync` sound. +struct Slot(UnsafeCell>>>); +unsafe impl Sync for Slot {} +impl Slot { + fn empty() -> Self { + Self(UnsafeCell::new(None)) + } + /// SAFETY: caller holds the sole claim on this unit (via the cursor). + unsafe fn set(&self, tokens: Result>) { + *self.0.get() = Some(tokens); + } + /// Consumer-side take, valid once the owning input has completed. + fn take(&self) -> Option>> { + // SAFETY: input complete; the single consumer is the only reader now. + unsafe { &mut *self.0.get() }.take() + } +} + +enum HandleInner { + /// Fully computed; yields `(input index, result)` in input order. + Ready(std::iter::Enumerate>>>), + /// Filled by pool workers writing shared slots; the calling thread assists + /// and surfaces each input in completion order the moment its chunks are in. + Streaming { + core: Arc, + state: StreamState, + }, +} + +/// Completion-order streaming cursor over a job's shared slots. No channel and no +/// per-unit messages: the calling thread assists the pool (claims and encodes +/// pending units) until the shared cursor is drained, and surfaces each input +/// in completion order (via `completed_order`) the moment its chunks are all in. +struct StreamState { + /// Next completion slot to surface — an index into + /// [`JobCore::completed_order`], not an input index: inputs are yielded in + /// the order they *finish*. + next_k: usize, + /// Total inputs. + n: usize, + /// Set once the cursor is exhausted: nothing left to assist, so the blocking + /// faces stop trying to claim and just wait on the last in-flight units. + assist_done: bool, +} + +impl StreamState { + fn new(n: usize) -> Self { + Self { + next_k: 0, + n, + assist_done: false, + } + } + + /// Surface the next input to *complete* (completion order), assisting the + /// pool while chunks are still in flight. `None` once every input is out. + /// The worker that finishes an input logs its `seq` in `completed_order`; we + /// read that log in order, so a fast input never waits behind a slow + /// earlier-index one. + fn next_completed(&mut self, core: &JobCore) -> Option<(usize, Result>)> { + if self.next_k >= self.n { + return None; + } + loop { + let seq = core.completed_order[self.next_k].load(Ordering::Acquire); + if seq != NOT_DONE { + self.next_k += 1; + return Some((seq, core.take_result(seq))); + } + // the caller has to wait for its results, so instead of spinning idle or parking the thread, + // we do some work which is ~equivalent to waiting: added benefit of not having + // to add extra machinery to wake the caller or implement the spin loop; we simply + // reuse our worker's job function + if !self.assist_done { + if SCRATCH.with(|st| core.run_one(&mut st.borrow_mut())) { + continue; + } + self.assist_done = true; + } + std::hint::spin_loop(); + } + } +} + +/// Sentinel for an unfilled [`JobCore::completed_order`] slot. +const NOT_DONE: usize = usize::MAX; + +/// 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` 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 via move; `&str`-family inputs pay one copy at +/// the boundary +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 their slots, +/// and the last `Arc` holder releases the storage. +/// Isolates a hot atomic on its own cache line. `JobCore.cursor` is `fetch_add`'d +/// once per unit by every worker; without padding it could share a line with the +/// read-only fields workers also read per unit, so each claim would invalidate +/// their caches (false sharing). +#[repr(align(64))] +struct CachePadded(T); + +struct JobCore { + storage: Box, + /// Ordered worker units, (seq, idx)-tagged for reassembly. + units: Vec, + /// Owned normalized segments for `Pretokenized` units + /// ([`SegmentText::Norm`] points here). + norm_bufs: Vec, + /// Flat pool of pre-token spans for `Pretokenized` units. + spans: Vec, + tokenizer: PipelineTokenizer, + cursor: CachePadded, + cancelled: AtomicBool, + /// `slots[seq][idx]` — one write-once result cell per chunk of each input. + slots: Vec>, + /// `pending[seq]` — chunks of input `seq` not yet written; zero ⇒ the input + /// is complete and its slots are safe to read. Distinct counters per input, + /// so worker completions don't contend on one line (only the cursor does). + pending: Vec, + /// Completion log: the worker that drives an input's `pending` to zero + /// claims the next `completed_slot` and stores its `seq` here, so the + /// consumer can surface inputs in completion order. `NOT_DONE` until filled. + completed_order: Vec, + /// Claim counter for the next free `completed_order` slot. + completed_slot: CachePadded, +} + +/// 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 in the source) this stride tiles. + /// Strides index from the segment start, so each 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 (`Raw`) and normalized (`Normalized`) paths. + fn tile(seg: Range, boundary: StrideBoundary) -> impl Iterator { + let strided = seg.len() >= 2 * PARALLEL_MIN_TOTAL_BYTES; + let n = if strided { + seg.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES) + } else { + 1 + }; + (0..n).map(move |index| { + if strided { + Region::Stride { + index, + boundary, + seg: seg.clone(), + } + } else { + Region::Full(seg.clone()) + } + }) + } + + /// 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) -> Option> { + match self { + Region::Full(r) => Some(r.clone()), + Region::Stride { + index, + boundary, + seg, + } => stride_range(&text[seg.clone()], *index, *boundary) + .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]` (`Normalized`). + Normalized { buf: usize, region: Region }, + /// Model-only over a span-group (the `Pretokenized` plan): `text` resolves + /// against input `seq` ([`SegmentText::Raw`]) or `JobCore::norm_bufs` + /// ([`SegmentText::Norm`]); `spans` indexes `JobCore::spans`. + Pretokenized { + text: SegmentText, + 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.0.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) { + 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) { + 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::Pretokenized { text, spans } => { + let text = match text { + SegmentText::Raw(r) => &self.storage.get(seq)[r.clone()], + SegmentText::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())); + // Publish: write our slot, then decrement the input's pending count. + // SAFETY: the cursor handed unit `i` to this thread alone, so we are the + // sole writer of `slots[seq][idx]`. + unsafe { self.slots[seq][unit.idx].set(res) }; + // AcqRel so the thread that drives `pending` to zero has acquired every + // other chunk's write (release sequence on `pending`); it then logs the + // completion, whose release publishes the whole input to the consumer. + if self.pending[seq].fetch_sub(1, Ordering::AcqRel) == 1 { + self.mark_done(seq); + } + true + } + + /// Log `seq` in `completed_order` — called once, by whoever drives its + /// `pending` to zero (a worker, or the builder for all-special / empty + /// inputs). Claims the next completion slot and releases `seq` into it. + fn mark_done(&self, seq: usize) { + let pos = self.completed_slot.0.fetch_add(1, Ordering::Relaxed); + self.completed_order[pos].store(seq, Ordering::Release); + } + + /// Concatenate input `seq`'s chunks in idx order, or its first error. Called + /// once per input by the single consumer, after `pending[seq] == 0`. A large + /// many-chunk input (a strided single document) would spend most of its + /// wall-time in this copy while the pool idles, so it is committed in + /// parallel: disjoint offsets → disjoint writes into the output buffer. + fn take_result(&self, seq: usize) -> Result> { + let row = &self.slots[seq]; + // One chunk (small / whole-unit inputs): hand back its Vec, no concat. + if row.len() == 1 { + return row[0].take().unwrap_or(Ok(Vec::new())); + } + // Move each chunk out (cheap — just Vec headers); first error wins. + let mut chunks: Vec> = Vec::with_capacity(row.len()); + for slot in row { + match slot.take() { + Some(Ok(t)) => chunks.push(t), + Some(Err(e)) => return Err(e), + None => chunks.push(Vec::new()), + } + } + let total: usize = chunks.iter().map(Vec::len).sum(); + + // Small enough (or no pool): serial concat. + const PAR_COMMIT_MIN: usize = 64 * 1024; + let Some(pool) = (total >= PAR_COMMIT_MIN).then(pool::rayon).flatten() else { + let mut out = Vec::with_capacity(total); + for c in &chunks { + out.extend_from_slice(c); + } + return Ok(out); + }; + + // Parallel scatter. Prefix-sum offsets, then each chunk copies into its + // own disjoint window of the output. + let mut offsets = Vec::with_capacity(chunks.len()); + let mut acc = 0usize; + for c in &chunks { + offsets.push(acc); + acc += c.len(); + } + let mut out: Vec = Vec::with_capacity(total); + let base = out.as_mut_ptr() as usize; + pool.install(|| { + chunks + .par_iter() + .zip(offsets.par_iter()) + .for_each(|(chunk, &off)| { + // SAFETY: offsets are a prefix sum of the chunk lengths, so + // the windows are disjoint and all lie within `total` + // (reserved above); no two tasks touch the same element. + unsafe { + std::ptr::copy_nonoverlapping( + chunk.as_ptr(), + (base as *mut PipelineToken).add(off), + chunk.len(), + ); + } + }); + }); + // SAFETY: the scatter wrote every element of `0..total`. + unsafe { out.set_len(total) }; + Ok(out) + } + + 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** — each input the moment it finishes (and *assisting*: the calling +/// thread claims and encodes pending chunks instead of idling), +/// - [`EncodeHandle::wait_for_completion`], which blocks and returns all results +/// collected back into **input order**. +/// +/// Dropping the job cancels all unclaimed work; chunks already being encoded +/// finish into slots nobody reads. Worker panics surface as the 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, n: usize) -> Self { + Self { + inner: HandleInner::Streaming { + core, + state: StreamState::new(n), + }, + } + } + + /// Number of inputs the handle will process + fn len(&self) -> usize { + match &self.inner { + HandleInner::Ready(it) => it.len(), + HandleInner::Streaming { state, .. } => state.n, + } + } + + /// 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 — each input the moment its chunks are all in. + 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 { @@ -282,8 +1172,8 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { /// Build a pipeline from an existing [`Tokenizer`], cloning its components. /// - /// The base [`Tokenizer`] carries the legacy [`crate::AddedVocabulary`]; the pipeline uses the - /// fast bucket [`BucketAddedVocabulary`], so we rebuild it from the tokenizer's added tokens. + /// The base [`Tokenizer`] carries the legacy `crate::AddedVocabulary`; the pipeline uses the + /// fast bucket `BucketAddedVocabulary`, so we rebuild it from the tokenizer's added tokens. /// Adding them in id order preserves ids (tokens present in the model reuse their model id, the /// rest keep their dense order), so the pipeline emits the same ids as the reference tokenizer. fn try_from(tok: &Tokenizer) -> Result { @@ -297,6 +1187,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). + // The special split 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 +1277,540 @@ 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. - /// - /// Special tokens are matched in two passes: - /// 1. on the raw input, - /// 2. then on each segment after normalization - /// - /// 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. + /// 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`]: + /// workers keep encoding in the background while the caller + /// holds it. /// - /// 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(); + /// The job is `'static`: it holds the input storage and a cheap clone of this + /// tokenizer handle. Dropping the job cancels unclaimed + /// work, and the last worker releases the storage. + 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(); + + 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, - )?; + 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<(usize, usize, u32)> = Vec::new(); + + // Probed on the longest input: it dominates the wall time, and a short one is + // unrepresentative (a 20-byte first input would decide the plan for a 4 MB second). + let plan = self.plan(refs.iter().copied().max_by_key(|s| s.len()).unwrap_or("")); + // `Normalized`/`Pretokenized` earn their serial prefix only for a + // segment larger than a fair per-thread share of the batch. Below that, + // batch parallelism balances it as one whole unit; the prefix would just + // serialize work. This also splits a lone huge segment inside an + // otherwise-large batch (a straggler otherwise). `Raw` never pays a + // prefix, so it ignores this. + let fair_share = total_bytes / pool.current_num_threads(); + + for (seq, &input) in refs.iter().enumerate() { + if input.len() < 2 * PARALLEL_MIN_TOTAL_BYTES { + self.emit_full(seq, 0..input.len(), &mut counts, &mut units); + continue; + } + for piece in SpecialSegmentIterator::new(input, &self.inner.added_vocabulary, false) { + let segment = match piece { + Segment::SpecialToken(id) => { + let idx = counts[seq]; + counts[seq] += 1; + preresolved.push((seq, idx, id)); + continue; + } + Segment::Text(segment) => segment, + }; + let base = offset_within(input, segment); + let range = base..base + segment.len(); + let big_enough = segment.len() >= 2 * PARALLEL_MIN_TOTAL_BYTES; + + match &plan { + ParallelPlan::Raw(boundary) => { + for region in Region::tile(range.clone(), *boundary) { + let idx = counts[seq]; + counts[seq] += 1; + units.push(Item { + seq, + idx, + work: Work::Raw(region), + }); + } + } + ParallelPlan::Normalized(boundary) + if big_enough && segment.len() > fair_share => + { + match self.inner.normalizer.as_ref().map_or_else( + || Ok(segment.to_owned()), + |n| n.normalize(segment).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, range.clone(), &mut counts, &mut units), + } + } + ParallelPlan::Pretokenized if big_enough && segment.len() > fair_share => { + match self.pretokenize_segment(segment, input) { + Ok(pretokenized) => { + // Merge this segment's pretokenized units into the job pools. + // It 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(pretokenized.norm_buf); + spans.extend_from_slice(&pretokenized.spans); + for unit in pretokenized.units { + let idx = counts[seq]; + counts[seq] += 1; + match unit { + SegmentUnit::Special(t) => preresolved.push((seq, idx, t)), + SegmentUnit::Group { text, spans: sr } => { + // Norm buffers were merged at `buf_base`; + // Raw offsets are already input-absolute. + let text = match text { + SegmentText::Norm { buf, range } => { + SegmentText::Norm { + buf: buf + buf_base, + range, + } + } + raw => raw, + }; + units.push(Item { + seq, + idx, + work: Work::Pretokenized { + text, + spans: sr.start + span_base..sr.end + span_base, + }, + }); + } + } + } + } + Err(_) => self.emit_full(seq, range.clone(), &mut counts, &mut units), + } + } + // `Whole`, or a prefix plan not worth it here (batch saturates + // the pool, or the segment is small): one whole-segment unit, + // full pipeline in the worker. + _ => self.emit_full(seq, range.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::Pretokenized { spans: sr, .. } => { + let group = &spans[sr.clone()]; + group + .last() + .map(|last| (last.end - group[0].start) as usize) + .unwrap_or(0) + } + }) + }); + + let n_units = units.len(); + let n_inputs = counts.len(); + // Per-(seq, idx) write-once cells, per-input pending counters, and the + // completion log (one slot per input, filled as inputs finish). + let slots: Vec> = counts + .iter() + .map(|&c| (0..c).map(|_| Slot::empty()).collect()) + .collect(); + let pending: Vec = counts.iter().map(|&c| AtomicUsize::new(c)).collect(); + let completed_order: Vec = + (0..n_inputs).map(|_| AtomicUsize::new(NOT_DONE)).collect(); + let core = Arc::new(JobCore { + storage: Box::new(storage), + units, + norm_bufs, + spans, + tokenizer: self.clone(), + cursor: CachePadded(AtomicUsize::new(0)), + cancelled: AtomicBool::new(false), + slots, + pending, + completed_order, + completed_slot: CachePadded(AtomicUsize::new(0)), + }); + // Specials were resolved in the serial prefix; write their cells now + // (single thread, before any worker runs) and retire them from + // `pending`. An all-special input completes here and is logged. + for (seq, idx, id) in preresolved { + // SAFETY: no worker has spawned yet; this thread is the sole writer. + unsafe { core.slots[seq][idx].set(Ok(vec![PipelineToken { id }])) }; + if core.pending[seq].fetch_sub(1, Ordering::Relaxed) == 1 { + core.mark_done(seq); + } + } + // Empty inputs (no chunks) are complete from the start. + for seq in 0..n_inputs { + if counts[seq] == 0 { + core.mark_done(seq); + } + } + + // 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, n_inputs) + } + + /// Model-only kernel for the `Pretokenized` 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`. + /// Workers use this for the `Normalized` plan (normalization already ran); + /// it is also the shared tail of `encode_one_with`. + /// Reuses the caller's [`EncodeState`] 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`, reusing the + /// caller's [`EncodeState`] scratch. The entry point for a worker holding a raw + /// region (the `Raw` and `Whole` plans); span-group model work goes through + /// `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, - )?; - } - } - } - } + /// How this config splits each special-free segment (specials are peeled + /// first). The cheapest safe split wins. + /// + /// `sample` is the text about to be encoded. A boundary the *config* allows is not the + /// same as a boundary the *input* contains, and the difference is not academic: + /// punctuation-terminated CJK has no cut at all (no spaces anywhere, and + /// [`NEWLINE_PREV`] rejects a newline after punctuation because the punct rule + /// ` ?[…]+[\r\n]*` absorbs it). Striding such an input scans it twice for cuts that do + /// not exist and then hands the whole thing to one worker — measured at 0.28x the plain + /// serial encode, and it does not parallelise at any thread count. So probe the input + /// before committing to a cut-based plan. + /// + /// Every plan is byte-exact (the oracles pin all of them), so this only ever trades + /// throughput — a mis-probe cannot change the ids. + fn plan(&self, sample: &str) -> ParallelPlan { + if let Some(boundary) = self.stride_boundary() { + if cut_exists(sample, boundary) { + // `Raw`: no serial prefix, so always worth it when the input can be cut. + return ParallelPlan::Raw(boundary); + } + } + if matches!(self.inner.model, PipelineModel::WordLevel(_)) { + // WordLevel's model + pre-tokenize are cheap, so a serial prefix + // would dwarf the parallel part. Lean on batch parallelism instead. + return ParallelPlan::Whole; + } + if let Some(boundary) = self.normalized_stride_boundary() { + // Normalizer rules out a raw cut, but the pre-tokenizer can cut the + // (serially) normalized text. Probed on the raw sample: normalization does not + // manufacture whitespace, so a text with no cut raw has none normalized either. + if cut_exists(sample, boundary) { + return ParallelPlan::Normalized(boundary); + } + } + // No usable cut: pre-tokenize serially and parallelise the model over span groups. + // The serial prefix goes through the normal pre-tokenizer, so it is `bitsplit` + // wherever `bitsplit` is wired (gpt2, cl100k) -- which is what makes this fallback + // cheap enough to prefer over a degenerate stride. + ParallelPlan::Pretokenized + } + + /// The pre-tokenizer's cut boundary for **already-normalized** text, gated + /// only on the added vocabulary (the `Normalized` plan normalizes serially + /// first, so the normalizer is out of the equation). The one remaining hazard + /// is a `normalized`-form added token — which the special split's raw pass + /// can't remove — split by a stride; the precomputed + /// `normalized_added_token_blocks_stride` + /// covers it (raw-matched tokens are already peeled). + 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 (the `Raw` plan). + /// Composes the per-stage split-safety: the normalizer must preserve + /// boundaries ([`NormalizerWrapper::preserves_stride_boundaries`]) *and* the + /// added vocabulary + pre-tokenizer must permit a cut + /// (`normalized_stride_boundary`). If the + /// normalizer doesn't preserve boundaries, the `Normalized` plan takes over + /// 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`. Shared fallback for `Whole` and for the + /// prefix plans 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::Pretokenized` plan over one + /// special-free segment (specials already peeled): normalize → frame + /// normalized specials → pre-tokenize, packaging the model work as + /// ~`PARALLEL_MIN_TOTAL_BYTES`-sized span-groups. `input` is the whole source + /// the borrowed ([`SegmentText::Raw`]) offsets resolve against; `segment` is a + /// subslice of it. Mirrors the walker in `encode_one_with` + /// stage by stage — the two must agree on unit order for reassembly to be + /// byte-identical to the serial encode. + fn pretokenize_segment(&self, segment: &str, input: &str) -> Result { + let mut pretokenized = PretokenizedSegment { + 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(segment)?, + None => Cow::Borrowed(segment), + }; + let owned = matches!(normalized, Cow::Owned(_)); + for piece in SpecialSegmentIterator::new(&normalized, &self.inner.added_vocabulary, true) { + match piece { + Segment::SpecialToken(token) => pretokenized.units.push(SegmentUnit::Special(token)), + Segment::Text(ntext) => { + pre_tokens.clear(); + self.inner + .pre_tokenizer + .pre_tokenize(ntext, &mut pre_tokens)?; + if pre_tokens.is_empty() { + continue; + } + // Where this piece'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, ntext) + } else { + // Borrowed all the way down: ntext is a subslice of the + // raw input. + offset_within(input, ntext) + }; + let text_range = base..base + ntext.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 = pretokenized.spans.len()..pretokenized.spans.len() + (e - g); + pretokenized.spans.extend_from_slice(&pre_tokens[g..e]); + let text = if owned { + // Local index 0 — the one `norm_buf`; merge rebases it. + SegmentText::Norm { + buf: 0, + range: text_range.clone(), + } + } else { + SegmentText::Raw(text_range.clone()) + }; + pretokenized.units.push(SegmentUnit::Group { + text, + spans: span_range, + }); + g = e; } } - }; + } } - Ok(()) + if let Cow::Owned(s) = normalized { + pretokenized.norm_buf = Some(s); + } + Ok(pretokenized) } } -/// What [`split`] does with each split it forms +/// 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,15 +1819,15 @@ pub enum SplitPolicy { Isolate, } -/// Splits `text` into same-class groups, emitting each as a [`Split`] -/// according to its [`SplitPolicy`]. +/// 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 +/// [`classify`] maps each char to a small `Copy + Eq` class, the current /// split ends whenever the class changes (or on every char of an `Isolate` /// 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, @@ -542,16 +1863,16 @@ pub fn split( } /// Splits `text` around a single delimiter predicate, honoring the full -/// [`SplitDelimiterBehavior`] contract. The pipeline-side equivalent of +/// `SplitDelimiterBehavior` contract. The pipeline-side equivalent of /// `NormalizedString::split(pattern, behavior)` for a char predicate. /// -/// The three non-merging behaviors reduce to a [`SplitPolicy`] on the delimiter -/// class and reuse [`split`]. The two merge variants are their own single pass: +/// The three non-merging behaviors reduce to a `SplitPolicy` on the delimiter +/// class and reuse `split`. The two merge variants are their own single pass: /// - `MergedWithPrevious` cuts the split *after* each delimiter, so a delimiter /// 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, @@ -611,7 +1932,7 @@ pub fn split_delimiter( }); } -/// Applies a [`SplitDelimiterBehavior`] to a match segmentation and appends the +/// Applies a `SplitDelimiterBehavior` to a match segmentation and appends the /// resulting pieces to `out`. /// /// `matches` is the `(offsets, is_match)` sequence covering the whole input, @@ -619,7 +1940,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 +2047,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,16 +2105,714 @@ 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)] mod tests { use super::*; + + /// A `plan()` probe input that plainly contains cuts (spaces after letters), so these + /// tests keep asserting what the *config* allows rather than what a sample happens to + /// hold -- `cut_exists` is covered separately by `uncuttable_input_skips_striding`. + const CUTTABLE: &str = "the quick brown fox jumps over the lazy dog and runs away fast "; use crate::models::bpe::BPE; use crate::models::wordpiece::WordPiece; 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 chunks = (0..strides) + .filter_map(|j| stride_range(&big, j, boundary_at_whitespace)) + .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 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) { + 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); + for boundary in [boundary_at_whitespace as StrideBoundary, boundary_fsm] { + assert_eq!(stride_range(&input, 0, boundary), Some(0..input.len())); + for j in 1..input.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES) { + assert_eq!( + stride_range(&input, j, boundary), + 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 reproduces the batch. + 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 the special split's raw pass); a raw/special one + /// is peeled first, so a stride can never bisect it. + /// The bug this guards: a config CAN cut (gpt2 exposes `boundary_fsm`) but the INPUT + /// cannot. Chinese prose has no space anywhere and ends its lines with `。`, and a + /// newline after punctuation is not a legal cut (the punct rule absorbs it), so striding + /// found zero cuts, gave one worker the whole document, and scanned the document twice + /// on the way. Measured at 0.28x the serial encode and flat across thread counts. + #[test] + fn uncuttable_input_skips_striding() { + let split = SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + ) + .unwrap(); + let mut tok = Tokenizer::new(crate::models::bpe::BPE::default()); + tok.with_pre_tokenizer(Some(split)); + let pipe = PipelineTokenizer::try_from(&tok).unwrap(); + + // The config itself allows a cut -- that is why the old code chose `Raw`. + assert!(pipe.stride_boundary().is_some()); + + // Punctuation-terminated Chinese: no space, and every newline follows `。`. + let cjk = "汉字汉字汉字汉字汉字。\n".repeat(2000); + assert!( + !cut_exists(&cjk, pipe.stride_boundary().unwrap()), + "`。` then newline must not count as a cut" + ); + assert!( + matches!(pipe.plan(&cjk), ParallelPlan::Pretokenized), + "uncuttable input must not be strided" + ); + + // A newline after a Han *letter* is a legal cut, so that input still strides. + let cuttable = "汉字汉字汉字汉字汉字\n".repeat(2000); + assert!(matches!(pipe.plan(&cuttable), ParallelPlan::Raw(_))); + + // And ordinary spaced text is unaffected. + assert!(matches!(pipe.plan(CUTTABLE), ParallelPlan::Raw(_))); + } + + #[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 the special split, so raw cuts still qualify. + assert!( + make(Some(AddedToken::from(" extra", true))) + .stride_boundary() + .is_some(), + "a raw special is peeled by the special split and must not disable raw cuts" + ); + // Normalized added token with whitespace: the special split's raw pass 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(CUTTABLE), ParallelPlan::Raw(_))); + + let big = "aa bb,cc! aa\tbb cc\n\n".repeat(2000); // ~44 KB, mixed runs + let chunks = (0..big.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES)) + .filter_map(|j| stride_range(&big, j, boundary_fsm)) + .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 "peeling specials means a stride can't split one" — it + /// holds only for *raw* (`normalized == false`) tokens. The special split'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 special_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"; + + // Raw 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, + "the 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 the special split. A *raw* (`normalized == false`) + /// `lstrip`/`rstrip` token is peeled by the special split 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 the special split 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: the special split 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(CUTTABLE), ParallelPlan::Raw(_)), + "raw affix token must keep Raw" + ); + // 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 `Pretokenized` regression); the + /// refined gate only checks `normalized` tokens (the special split peels the raw ones), + /// so llama-3 must keep `Raw` 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(CUTTABLE), ParallelPlan::Raw(_)), + "llama-3's raw-only specials must not disqualify Raw" + ); + } + + /// 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::Pretokenized`. 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() + } + + /// Pretokenized: 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(CUTTABLE), ParallelPlan::Pretokenized)); + + let big = "aa.bb.cc.".repeat(3000); // ~27 KB, punctuation-delimited + let pretokenized = tok.pretokenize_segment(&big, &big).unwrap(); + let groups = pretokenized + .units + .iter() + .filter(|u| matches!(u, SegmentUnit::Group { .. })) + .count(); + assert!(groups > 1, "expected several span-groups, got {}", groups); + assert!(pretokenized.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, "Pretokenized != serial"); + let owned = ids(tok.encode(big).into_single().unwrap()); + assert_eq!(owned, serial, "owned Pretokenized != serial"); + } + + /// Pretokenized 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(CUTTABLE), ParallelPlan::Pretokenized)); + + let half = "AA.BB.cc.".repeat(2000); + let big = format!("{half}{}", "CC.aa.BB.".repeat(2000)); + // The special split peels `` in the builder; the per-segment model prefix runs on + // special-free text and (with `Lowercase`) must own its normalized buffer. + let pretokenized = tok.pretokenize_segment(&half, &half).unwrap(); + assert!( + pretokenized.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 Pretokenized either. + assert!( + matches!(pipe.plan(CUTTABLE), ParallelPlan::Whole), + "Strip + WordLevel must fall to batch-level parallelism only", + ); + } + + /// A config no plan can split — `WordLevel` (so no `Pretokenized`) under a + /// boundary-hostile `Strip` normalizer (so no `Raw`/`Normalized`) picks + /// `Whole`. A single large input dense with special tokens still + /// parallelizes: the special split (always first) peels the specials + /// (preresolved) and hands each inter-special segment to a worker. Must stay + /// byte-identical to the serial encode, both borrowed and owned handles. + #[test] + fn 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(CUTTABLE), ParallelPlan::Whole), + "Strip + WordLevel plan must be Whole (the special split 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, "special-split (borrowed) != serial"); + let owned = ids(tok.encode(big.clone()).into_single().unwrap()); + assert_eq!(owned, serial, "special-split (owned) != serial"); + } + + /// The `Normalized` plan: a non-boundary-preserving normalizer + /// (`Prepend`) rules out `Raw`, 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(CUTTABLE), ParallelPlan::Normalized(_)), + "Prepend (unsafe normalizer) + WhitespaceSplit + WordPiece must pick `Normalized`" + ); + 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, "Normalized (borrowed) != serial"); + let owned = ids(tok.encode(big.clone()).into_single().unwrap()); + assert_eq!(owned, serial, "Normalized (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..82778de5b --- /dev/null +++ b/tokenizers/tk-encode/src/tokenizer/pool.rs @@ -0,0 +1,152 @@ +//! Process-global worker pool for the parallel encode path: a lazily-built, +//! library-private `rayon::ThreadPool`. +//! +//! Scheduling lives in `pipeline.rs` as the shared-cursor job model: drainer +//! tasks and the consuming thread claim chunks off one atomic cursor. This pool +//! only provides the worker threads. +//! +//! 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..01dc2e8a1 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 — safe because the +/// special split always peels specials first, so a stride can't bisect one. This +/// exercises number-transition 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); + } } )* };