diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index eb647501b..58f3ed248 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -521,6 +521,27 @@ 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` @@ -1316,7 +1337,9 @@ impl PipelineTokenizer { let mut spans: Vec = Vec::new(); let mut preresolved: Vec<(usize, usize, u32)> = Vec::new(); - let plan = self.plan(); + // 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 @@ -1613,21 +1636,43 @@ impl PipelineTokenizer { /// How this config splits each special-free segment (specials are peeled /// first). The cheapest safe split wins. - fn plan(&self) -> ParallelPlan { + /// + /// `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() { - // `Raw`: no serial prefix, so always worth it when available. - ParallelPlan::Raw(boundary) - } else if matches!(self.inner.model, PipelineModel::WordLevel(_)) { + 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. - ParallelPlan::Whole - } else if let Some(boundary) = self.normalized_stride_boundary() { + 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. - ParallelPlan::Normalized(boundary) - } else { - ParallelPlan::Pretokenized + // (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 @@ -2079,6 +2124,11 @@ 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; @@ -2314,6 +2364,45 @@ mod tests { /// 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; @@ -2389,7 +2478,7 @@ mod tests { .unwrap(), )); let tok = PipelineTokenizer::try_from(&tok).unwrap(); - assert!(matches!(tok.plan(), ParallelPlan::Raw(_))); + 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)) @@ -2499,7 +2588,7 @@ mod tests { .unwrap(); let pipe = PipelineTokenizer::try_from(&tok).unwrap(); assert!( - matches!(pipe.plan(), ParallelPlan::Raw(_)), + matches!(pipe.plan(CUTTABLE), ParallelPlan::Raw(_)), "raw affix token must keep Raw" ); // Two ~20 KB segments around one lstrip token → both stride. @@ -2521,7 +2610,7 @@ mod tests { let oracle = Tokenizer::from_file("../data/llama-3-tokenizer.json").unwrap(); let tok = PipelineTokenizer::try_from(&oracle).unwrap(); assert!( - matches!(tok.plan(), ParallelPlan::Raw(_)), + matches!(tok.plan(CUTTABLE), ParallelPlan::Raw(_)), "llama-3's raw-only specials must not disqualify Raw" ); } @@ -2564,7 +2653,7 @@ mod tests { #[test] fn split_at_model_matches_serial() { let tok = split_at_model_pipeline(false); - assert!(matches!(tok.plan(), ParallelPlan::Pretokenized)); + 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(); @@ -2590,7 +2679,7 @@ mod tests { #[test] fn split_at_model_normalizer_and_specials_match_serial() { let tok = split_at_model_pipeline(true); - assert!(matches!(tok.plan(), ParallelPlan::Pretokenized)); + assert!(matches!(tok.plan(CUTTABLE), ParallelPlan::Pretokenized)); let half = "AA.BB.cc.".repeat(2000); let big = format!("{half}{}", "CC.aa.BB.".repeat(2000)); @@ -2637,7 +2726,7 @@ mod tests { ); // WordLevel model: the plan cannot escalate to Pretokenized either. assert!( - matches!(pipe.plan(), ParallelPlan::Whole), + matches!(pipe.plan(CUTTABLE), ParallelPlan::Whole), "Strip + WordLevel must fall to batch-level parallelism only", ); } @@ -2671,7 +2760,7 @@ mod tests { .unwrap(); let tok = PipelineTokenizer::try_from(&tok).unwrap(); assert!( - matches!(tok.plan(), ParallelPlan::Whole), + matches!(tok.plan(CUTTABLE), ParallelPlan::Whole), "Strip + WordLevel plan must be Whole (the special split does the splitting)" ); @@ -2712,7 +2801,7 @@ mod tests { .unwrap(); let tok = PipelineTokenizer::try_from(&tok).unwrap(); assert!( - matches!(tok.plan(), ParallelPlan::Normalized(_)), + 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