diff --git a/.github/scripts/render_pipeline_bench.py b/.github/scripts/render_pipeline_bench.py
index ed7e92c71..41670db0b 100644
--- a/.github/scripts/render_pipeline_bench.py
+++ b/.github/scripts/render_pipeline_bench.py
@@ -47,6 +47,12 @@
# `added_split` = added/special-token scan (AddedVocabulary), `pre_tokenize` =
# pre-tokenizer split — two distinct splitting costs. `post` = special-token
# id-frame splice (~0 for models with no post-processor). The five sum to `total`.
+#
+# `total` is the ladder's own full encode, NOT the public `encode` call the throughput
+# phase times: that additionally acquires a pooled scratch and returns an owned Vec.
+# The subtitle names that exclusion rather than the stack absorbing it, because the
+# difference is part real cost and part pooled-vs-fresh cache state, and no measurement in
+# the bench separates the two.
STAGES = [("added_split", "added-token"), ("normalize", "normalize"),
("pre_tokenize", "pre-tokenize"), ("model", "model"), ("post", "post")]
STAGE_INK = {"added_split": "#7a5ea8", "normalize": "#2a9d8f",
@@ -640,12 +646,20 @@ def stage_chart_svg(model, subtitle_base, meta, baseline_label):
sink = STAGE_INK
rows = [r for r in model["results"] if "stage_ns_per_byte" in r]
+ # A row's stages are shares of that row's `total`, and they can sum past 100% when a
+ # rung measured at or below the noise floor and clamped to zero. Scaling the plot by the
+ # widest row keeps such a bar on the canvas and visibly past the 100% gridline, instead
+ # of running off the right edge as though the chart were broken.
+ span = max([1.0] + [sum(r["stage_ns_per_byte"].get(k, 0.0) for k, _ in STAGES)
+ / (r["stage_ns_per_byte"]["total"] or 1.0) for r in rows])
+ scale = PLOT_W / span
+
top = 84
col_x = GUTTER + PLOT_W + PAD_R + COL_W - 16
body = [f'total ns/B · ×speedup',
f''
- f'share of pipeline encode time (each bar = 100%) · label = share% · ns/B']
+ f'share of pipeline encode time (100% = that fixture\'s total) · label = share% · ns/B']
y = top
for key, title in GROUPS:
group_rows = sorted((r for r in rows if r["group"] == key),
@@ -665,7 +679,7 @@ def stage_chart_svg(model, subtitle_base, meta, baseline_label):
for skey, _ in STAGES:
val = s.get(skey, 0.0)
frac = val / total
- seg = frac * PLOT_W # each bar fills PLOT_W (100% of this fixture's total)
+ seg = frac * scale # 100% of this fixture's total is `scale` wide
if seg > 0.4:
body.append(f'')
@@ -689,7 +703,7 @@ def stage_chart_svg(model, subtitle_base, meta, baseline_label):
grid = []
for pct in (0, 25, 50, 75, 100):
- gx = GUTTER + pct / 100 * PLOT_W
+ gx = GUTTER + pct / 100 * scale
grid.append(f'')
grid.append(f'= 0.02)
+ # The route is part of the reading, not decoration: each one decomposes into stages
+ # differently, and the split bar only means the shipped splitter because every rung
+ # drives this route. Older JSON has no `route` key, so say so rather than guess.
+ route = model.get("route") or "route not recorded"
+ subtitle = f'{model["shape"]} · {route} route · {mix_txt}'
+ # Two things the reader would otherwise have to infer. The rungs are independent
+ # medians, so one at or below the noise floor clamps to zero and the parts then sum
+ # past 100%, which is visible as a bar overrunning the axis. And `total` is the
+ # ladder's own encode, without the pooled scratch and returned Vec of a real call.
+ sums = [sum(r["stage_ns_per_byte"].get(k, 0.0) for k, _ in STAGES)
+ / r["stage_ns_per_byte"]["total"]
+ for r in rows if r["stage_ns_per_byte"].get("total")]
+ if sums:
+ subtitle += f' · parts sum {100 * sum(sums) / len(sums):.0f}% (rung noise)'
+ if any("alloc" in r["stage_ns_per_byte"] for r in rows):
+ subtitle += " · excludes encode() call overhead"
return svg_doc(ink, height, f'{model["model"]} — Pipeline encode stage mix',
subtitle, "".join(grid) + "".join(body) + legend, meta, subtitle_base)
@@ -884,9 +915,9 @@ def x(v):
def pretok_compare_md(model):
- """Per-fixture 'classify + fsm vs a regex engine' table — the pre-tokenize split
- vs onig/fancy/pcre2/logos on the model's own regex, WITH and WITHOUT SIMD.
- Rendered only for regex pre-tokenizers (a reference is non-null)."""
+ """Per-fixture table of the shipped pre-tokenize split against onig, fancy, pcre2 and
+ logos on the model's own regex. Rendered only for regex pre-tokenizers (a reference is
+ non-null)."""
engines = ("onig", "fancy", "pcre2", "logos")
def has_ref(r):
@@ -898,13 +929,15 @@ def has_ref(r):
cell = lambda v: fnum(v, "{:.1f}") # noqa: E731 — "—" for null
def ratio(v, sp, cp):
return f"{v / sp:.1f}× / {v / cp:.1f}×" if (v is not None and sp > 0 and cp > 0) else "—"
- cols = 4 + 2 * len(engines) # cls×2 + pipe×2 + one abs + one ×vs per engine
- md = ["", "**Pre-tokenize: `classify + fsm` vs regex engines** — ns/byte, lower better. The fsm is "
- "the scalar jump-table in both pipe columns; **SIMD / scalar is the classify pass** (regex "
- "pre-tokenizers have no SIMD fsm). `×vs` = engine ÷ our pipeline (SIMD / scalar classify); "
- "`onig` & `pcre2` (JIT) are C, `fancy` is pure-Rust fancy-regex, `logos` is a compile-time DFA "
- "lexer (approximate grammar; n/a for deepseek).", "",
- "| Fixture | classify SIMD | classify scalar | pipe (SIMD cls + fsm) | pipe (scalar cls + fsm) | "
+ cols = 4 + 2 * len(engines) # cls×2 + split×2 + one abs + one ×vs per engine
+ md = ["", "**Pre-tokenize: our split vs regex engines.** ns/byte, lower better. `split` is the "
+ "splitter this model's route actually runs, straight off the stage ladder. "
+ "`scalar-cls` swaps **only the classify pass** for its scalar version. It is not the split "
+ "without SIMD, since a regex-shaped fsm is a boundary-mask scanner and carries SIMD of its "
+ "own. `×vs` = engine ÷ our split (shipped / scalar-classify); `onig` & `pcre2` (JIT) are C, "
+ "`fancy` is pure-Rust fancy-regex, `logos` is a compile-time DFA lexer (approximate grammar; "
+ "n/a for deepseek).", "",
+ "| Fixture | classify SIMD | classify scalar | split (shipped) | split (scalar-cls) | "
+ " | ".join(engines) + " | " + " | ".join(f"×vs {e}" for e in engines) + " |",
"|---" + "|---:" * cols + "|"]
for r in sorted(rows, key=lambda r: (r["group"], r["fixture"])):
diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml
index fb8947f3f..dce664880 100644
--- a/tokenizers/atomsplit/Cargo.toml
+++ b/tokenizers/atomsplit/Cargo.toml
@@ -43,6 +43,10 @@ harness = false
name = "classify"
harness = false
+[[bench]]
+name = "literal"
+harness = false
+
[[bench]]
name = "class_runs"
harness = false
diff --git a/tokenizers/atomsplit/benches/literal.rs b/tokenizers/atomsplit/benches/literal.rs
new file mode 100644
index 000000000..c165b3f6c
--- /dev/null
+++ b/tokenizers/atomsplit/benches/literal.rs
@@ -0,0 +1,67 @@
+//! `Literal::matches` (one search per match) against `Literal::matches_into` (one scan per
+//! text), across the delimiter densities pre-tokenizers see: a space about every six bytes of
+//! English, a `▁` per word after a SentencePiece replace, and a pattern the text does not
+//! contain at all (where `memmem`'s skip-ahead is at its best and the scan must keep up).
+//! Run: cargo bench --bench literal
+use atomsplit::literal::Literal;
+use std::hint::black_box;
+use std::time::Instant;
+
+fn best_ns_per_byte(text_len: usize, mut pass: impl FnMut() -> usize) -> f64 {
+ let iters = (16_000_000 / text_len.max(1)).clamp(4, 400) as u32;
+ for _ in 0..3 {
+ black_box(pass());
+ }
+ let mut best = f64::INFINITY;
+ for _ in 0..9 {
+ let t = Instant::now();
+ for _ in 0..iters {
+ black_box(pass());
+ }
+ best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text_len) as f64);
+ }
+ best
+}
+
+fn main() {
+ let manifest = env!("CARGO_MANIFEST_DIR");
+ let english =
+ std::fs::read_to_string(format!("{manifest}/../data/big.txt")).unwrap_or_default();
+ let english: String = english.chars().take(2_000_000).collect();
+ if english.is_empty() {
+ println!("empty corpus — data/big.txt missing? (make test downloads it)");
+ return;
+ }
+ let metaspaced = english.replace(' ', "\u{2581}");
+
+ println!(
+ "{:<22} {:>9} {:>12} {:>12} {:>9} {:>12}",
+ "", "matches", "iterator", "batch", "speedup", "count-only"
+ );
+ for (label, text, pattern) in [
+ ("space in English", english.as_bytes(), " "),
+ ("▁ in metaspaced", metaspaced.as_bytes(), "\u{2581}"),
+ ("▁ absent", english.as_bytes(), "\u{2581}"),
+ ] {
+ let literal = Literal::new(pattern.as_bytes()).unwrap();
+ let mut offsets: Vec = Vec::with_capacity(text.len());
+ let mut buffer = vec![0u32; text.len() + 4];
+
+ let count = literal.matches(text).count();
+ let iterator = best_ns_per_byte(text.len(), || {
+ offsets.clear();
+ offsets.extend(literal.matches(text));
+ offsets.len()
+ });
+ let batch = best_ns_per_byte(text.len(), || literal.matches_into(text, &mut buffer));
+ let counting = best_ns_per_byte(text.len(), || literal.count_matches(text));
+
+ println!(
+ "{label:<22} {count:>9} {:>7.2} GB/s {:>7.2} GB/s {:>8.2}x {:>7.2} GB/s",
+ 1.0 / iterator,
+ 1.0 / batch,
+ iterator / batch,
+ 1.0 / counting
+ );
+ }
+}
diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs
index 45c8ca286..5d6b950d4 100644
--- a/tokenizers/atomsplit/src/fsm.rs
+++ b/tokenizers/atomsplit/src/fsm.rs
@@ -6,8 +6,8 @@
//! [`class_runs_into`]: on aarch64/wasm the SIMD movemask boundary-extractor + homogeneous-chunk
//! early-out (in `simd_fsm`), elsewhere the scalar run-end core ([`emit_class_spans`]). The
//! regex-shaped ones ([`fsm_cl100k`] / [`fsm_o200k`] / [`fsm_tekken`] / [`fsm_deepseek`] /
-//! [`fsm_byte_level`]) are scalar jump-tables (only the class family's [`class_runs_into`] has a SIMD
-//! path).
+//! [`fsm_byte_level`]) are scalar jump-tables; byte_level additionally has a boundary-mask SIMD
+//! form ([`scan_byte_level_masked`], aarch64 only).
pub(crate) use crate::classify::{Atom, char_len, classify, in_mask, mask};
// Atom-tag aliases, shared with the per-tokenizer FSM submodules (`fsm/*.rs`) via `use super::*`.
@@ -237,11 +237,17 @@ pub fn emit_class_spans(
mod byte_level;
mod cl100k;
mod deepseek;
+mod masked;
mod o200k;
-pub use byte_level::fsm_byte_level;
-pub use cl100k::{fsm_cl100k, fsm_cl100k_cap};
-pub use deepseek::fsm_deepseek;
-pub use o200k::{fsm_o200k, fsm_tekken};
+pub use byte_level::{fsm_byte_level, scan_byte_level};
+pub use cl100k::{fsm_cl100k, fsm_cl100k_cap, scan_cl100k_cap};
+pub use deepseek::{fsm_deepseek, scan_deepseek};
+pub use masked::{
+ fsm_byte_level_masked, fsm_cl100k_cap_masked, fsm_deepseek_masked, fsm_o200k_masked,
+ fsm_tekken_masked, scan_byte_level_masked, scan_cl100k_cap_masked, scan_deepseek_masked,
+ scan_o200k_masked, scan_tekken_masked,
+};
+pub use o200k::{fsm_o200k, fsm_tekken, scan_o200k, scan_tekken};
// ── Composition recipes ────────────────────────────────────────────────────────────────────────
// Each pre-tokenizer = (classify → fsm shape + params). `tags` and `out` are caller-owned
@@ -339,49 +345,3 @@ impl ByteLevel {
fsm_byte_level(text, tags, out)
}
}
-
-/// `Split(char, Removed)` — the only pre-tokenizer that keys on a *literal char* rather than an atom
-/// class, so it scans bytes directly (no classify pass). UTF-8 is self-synchronizing, so the
-/// delimiter's byte pattern only matches on char boundaries.
-pub struct CharDelimiterSplit(pub char);
-impl CharDelimiterSplit {
- /// Split on the literal char (Removed); writes spans into `out` (len ≥ `text.len()`), returns count.
- #[inline]
- #[must_use]
- pub fn pre_tokenize(&self, text: &[u8], _tags: &mut [u8], out: &mut [Span]) -> usize {
- debug_assert!(out.len() >= text.len());
- let mut buf = [0u8; 4];
- let delim = self.0.encode_utf8(&mut buf).as_bytes();
- let (n, dl) = (text.len(), delim.len());
- let (mut start, mut i, mut w) = (0usize, 0usize, 0usize);
- while i + dl <= n {
- // memchr the first delimiter byte, then confirm the full pattern. memchr (already a
- // workspace dep) beats a scalar scan 1.4–23× here — the gap widening as the delimiter
- // gets rarer over large inputs, since its SIMD skips whole 16/32/64-byte strides.
- match memchr::memchr(delim[0], &text[i..n - dl + 1]) {
- Some(off) if text[i + off..i + off + dl] == *delim => {
- let m = i + off;
- if m > start {
- out[w] = Span {
- start: start as u32,
- end: m as u32,
- }; // gap before the delimiter (Removed)
- w += 1;
- }
- i = m + dl;
- start = i;
- }
- Some(off) => i += off + 1, // first byte matched mid-pattern; keep scanning
- None => break,
- }
- }
- if start < n {
- out[w] = Span {
- start: start as u32,
- end: n as u32,
- };
- w += 1;
- }
- w
- }
-}
diff --git a/tokenizers/atomsplit/src/fsm/byte_level.rs b/tokenizers/atomsplit/src/fsm/byte_level.rs
index 73f84ac9f..479a4b1de 100644
--- a/tokenizers/atomsplit/src/fsm/byte_level.rs
+++ b/tokenizers/atomsplit/src/fsm/byte_level.rs
@@ -8,10 +8,21 @@ use super::*;
#[must_use]
pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
- let end = text.len();
- // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
- // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
- let tags = &tags[..end];
+ let mut w = 0usize;
+ scan_byte_level(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of the
+/// byte-level regex. [`scan_byte_level`] loops it over the whole text; the masked scanner
+/// ([`super::scan_byte_level_masked`]) re-derives tokens with it where its batch masks are not
+/// trustworthy.
+#[inline(always)]
+pub(super) fn advance_byte_level(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
// `\s+(?!\S)|\s+`: the whole run at EOF, else leave the last ws char for the next ` ?`-prefixed run.
let ws = |i: usize| -> usize {
let re = run_end(tags, i, end, mask::WS);
@@ -25,57 +36,57 @@ pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
if last > i { last } else { re }
}
};
-
- let mut i = 0;
- let mut w = 0usize;
- while i < end {
- let start = i;
- match tags[i] & 0x0F {
- LET => i = run_end(tags, i, end, mask::LETTER), // ` ?\p{L}+` (space taken by the Space arm)
- NW | NO => i = run_end(tags, i, end, mask::NUMBER), // ` ?\p{N}+` — UNBOUNDED
- MRK | CON | PUN | SYM | NMO | CTL => i = run_end(tags, i, end, mask::NOT_WS_L_N), // ` ?[^…]+`
- // `'s|'t|'re|'ve|'m|'ll|'d` (case-sensitive), else `[^\s\p{L}\p{N}]+` (apostrophe ∈ that set)
- APO => {
- let adv = match (text.get(i + 1), text.get(i + 2)) {
- (Some(b's' | b't' | b'm' | b'd'), _) => 2,
- (Some(b'r'), Some(b'e'))
- | (Some(b'v'), Some(b'e'))
- | (Some(b'l'), Some(b'l')) => 3,
- _ => 0,
- };
- i = if adv > 0 {
- i + adv
- } else {
- run_end(tags, i, end, mask::NOT_WS_L_N)
- };
+ match tags[i] & 0x0F {
+ LET => run_end(tags, i, end, mask::LETTER), // ` ?\p{L}+` (space taken by the Space arm)
+ NW | NO => run_end(tags, i, end, mask::NUMBER), // ` ?\p{N}+` — UNBOUNDED
+ MRK | CON | PUN | SYM | NMO | CTL => run_end(tags, i, end, mask::NOT_WS_L_N), // ` ?[^…]+`
+ // `'s|'t|'re|'ve|'m|'ll|'d` (case-sensitive), else `[^\s\p{L}\p{N}]+` (apostrophe ∈ that set)
+ APO => {
+ let adv = match (text.get(i + 1), text.get(i + 2)) {
+ (Some(b's' | b't' | b'm' | b'd'), _) => 2,
+ (Some(b'r'), Some(b'e')) | (Some(b'v'), Some(b'e')) | (Some(b'l'), Some(b'l')) => 3,
+ _ => 0,
+ };
+ if adv > 0 {
+ i + adv
+ } else {
+ run_end(tags, i, end, mask::NOT_WS_L_N)
}
- // Space: the ` ?` prefix — attach one space to a following letter / number / "other" run,
- // else it's whitespace (rules `\s+(?!\S)|\s+`, which leave one space for the next run).
- SPC => {
- let a = i + 1; // Space is ASCII (0x20)
- i = match tags.get(a).map(|&t| t & 0x0F) {
- Some(LET) => run_end(tags, a, end, mask::LETTER),
- Some(NW) | Some(NO) => run_end(tags, a, end, mask::NUMBER),
- Some(t) if in_mask(t, mask::NOT_WS_L_N) => {
- run_end(tags, a, end, mask::NOT_WS_L_N)
- }
- _ => ws(i),
- };
- }
- // WsOther / Newline: whitespace only — the ` ?` prefix is a literal 0x20, so tabs/newlines
- // never prefix a run.
- WSO | NLN => i = ws(i),
- // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively.
- _ => i += char_len(text[i]),
}
- // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1).
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: i as u32,
+ // Space: the ` ?` prefix — attach one space to a following letter / number / "other" run,
+ // else it's whitespace (rules `\s+(?!\S)|\s+`, which leave one space for the next run).
+ SPC => {
+ let a = i + 1; // Space is ASCII (0x20)
+ match tags.get(a).map(|&t| t & 0x0F) {
+ Some(LET) => run_end(tags, a, end, mask::LETTER),
+ Some(NW) | Some(NO) => run_end(tags, a, end, mask::NUMBER),
+ Some(t) if in_mask(t, mask::NOT_WS_L_N) => run_end(tags, a, end, mask::NOT_WS_L_N),
+ _ => ws(i),
}
- };
- w += 1;
+ }
+ // WsOther / Newline: whitespace only — the ` ?` prefix is a literal 0x20, so tabs/newlines
+ // never prefix a run.
+ WSO | NLN => ws(i),
+ // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively.
+ _ => i + char_len(text[i]),
+ }
+}
+
+/// The scan under [`fsm_byte_level`]: hands each token to `emit` the moment it is cut,
+/// so a caller can consume tokens in place instead of collecting a span buffer first.
+pub fn scan_byte_level(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) {
+ debug_assert!(tags.len() >= text.len());
+ let end = text.len();
+ // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
+ // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
+ let tags = &tags[..end];
+ let mut i = 0;
+ while i < end {
+ let e = advance_byte_level(text, tags, i, end);
+ emit(Span {
+ start: i as u32,
+ end: e as u32,
+ });
+ i = e;
}
- w
}
diff --git a/tokenizers/atomsplit/src/fsm/cl100k.rs b/tokenizers/atomsplit/src/fsm/cl100k.rs
index 65976ecef..5450fd284 100644
--- a/tokenizers/atomsplit/src/fsm/cl100k.rs
+++ b/tokenizers/atomsplit/src/fsm/cl100k.rs
@@ -14,17 +14,26 @@ pub fn fsm_cl100k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
#[must_use]
pub fn fsm_cl100k_cap(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize {
debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
- cl100k(text, tags, out, digit_cap)
+ let mut w = 0usize;
+ scan_cl100k_cap(text, tags, digit_cap, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
}
-fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize {
- // Leading-atom values, as `const` so the `match` below is a dense jump table (not an if-cascade):
- // the dispatch is O(1) and a token never pays for a rule it can't start (e.g. non-number tokens
- // never test the number rule — which is what the POC's const-gating removed by hand; here it's free).
- let end = text.len();
- // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
- // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
- let tags = &tags[..end];
+/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of
+/// the cl100k-family regex. [`scan_cl100k_cap`] loops it over the whole text; the masked scanner
+/// re-derives tokens with it where its batch masks are not trustworthy.
+#[inline(always)]
+pub(super) fn advance_cl100k_cap(
+ text: &[u8],
+ tags: &[u8],
+ i: usize,
+ end: usize,
+ digit_cap: usize,
+) -> usize {
let letters = |a: usize| run_end(tags, a, end, mask::LETTER);
// rule 4 body: `[^\s\p{L}\p{N}]+[\r\n]*` from `sp0` (any leading space already consumed). Returns
// the run end, or `sp0` if there is no "other" run there (caller then treats it as whitespace).
@@ -40,79 +49,84 @@ fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize
// rules 5-7 (`\s*[\r\n] | \s+(?!\S) | \s+`) → the shared `ws_tail`.
let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) };
- let mut i = 0;
- let mut w = 0usize;
- while i < end {
- let start = i;
- let b = text[i];
- match tags[i] & 0x0F {
- // rule 2: `\p{L}+`
- LET => i = letters(i),
- // rule 3: `\p{N}{1,cap}` (cap = 3 cl100k, 1 Qwen2, MAX for `\p{N}+`)
- NW | NO => {
- let (mut p, mut cnt) = (i, 0);
- while p < end && cnt < digit_cap && in_mask(tags[p], mask::NUMBER) {
- p += char_len(text[p]);
- cnt += 1;
- }
- i = p;
- }
- // Space: rule 2 (space prefix + `\p{L}+`) | rule 4 (` ` + "other") | rules 5-7
- SPC => {
- let a = i + 1; // Space is ASCII (0x20)
- i = if a < end && (tags[a] & 0x0F) == LET {
- letters(a)
- } else {
- let p = other(a);
- if p > a { p } else { ws(i) }
- };
+ let b = text[i];
+ match tags[i] & 0x0F {
+ // rule 2: `\p{L}+`
+ LET => letters(i),
+ // rule 3: `\p{N}{1,cap}` (cap = 3 cl100k, 1 Qwen2, MAX for `\p{N}+`)
+ NW | NO => {
+ let (mut p, mut cnt) = (i, 0);
+ while p < end && cnt < digit_cap && in_mask(tags[p], mask::NUMBER) {
+ p += char_len(text[p]);
+ cnt += 1;
}
- // WsOther: rule 2 (prefix + `\p{L}+`) | whitespace (never rule 4 — not in NOT_WS_L_N)
- WSO => {
- let a = i + char_len(b);
- i = if a < end && (tags[a] & 0x0F) == LET {
- letters(a)
- } else {
- ws(i)
- };
+ p
+ }
+ // Space: rule 2 (space prefix + `\p{L}+`) | rule 4 (` ` + "other") | rules 5-7
+ SPC => {
+ let a = i + 1; // Space is ASCII (0x20)
+ if a < end && (tags[a] & 0x0F) == LET {
+ letters(a)
+ } else {
+ let p = other(a);
+ if p > a { p } else { ws(i) }
}
- // Newline: whitespace (rule 5 ends at the last newline)
- NLN => i = ws(i),
- // Apostrophe: rule 1 (contraction) | rule 2 (prefix + `\p{L}+`) | rule 4
- APO => {
- let adv = contraction(text, i); // rule 1: `'s 't 're 've 'm 'll 'd` (case-insensitive)
- i = if adv > 0 {
- i + adv
- } else {
- let a = i + 1; // Apostrophe is ASCII (0x27)
- if a < end && (tags[a] & 0x0F) == LET {
- letters(a)
- } else {
- other(i)
- } // c ∈ NOT_WS_L_N ⇒ > i
- };
+ }
+ // WsOther: rule 2 (prefix + `\p{L}+`) | whitespace (never rule 4 — not in NOT_WS_L_N)
+ WSO => {
+ let a = i + char_len(b);
+ if a < end && (tags[a] & 0x0F) == LET {
+ letters(a)
+ } else {
+ ws(i)
}
- // Mark | Connector | Punct | SymOther | NumericOther | Control (all in NOT_WS_L_N):
- // rule 2 (prefix + `\p{L}+`) | rule 4
- MRK | CON | PUN | SYM | NMO | CTL => {
- let a = i + char_len(b);
- i = if a < end && (tags[a] & 0x0F) == LET {
+ }
+ // Newline: whitespace (rule 5 ends at the last newline)
+ NLN => ws(i),
+ // Apostrophe: rule 1 (contraction) | rule 2 (prefix + `\p{L}+`) | rule 4
+ APO => {
+ let adv = contraction(text, i); // rule 1: `'s 't 're 've 'm 'll 'd` (case-insensitive)
+ if adv > 0 {
+ i + adv
+ } else {
+ let a = i + 1; // Apostrophe is ASCII (0x27)
+ if a < end && (tags[a] & 0x0F) == LET {
letters(a)
} else {
other(i)
- }; // c ∈ NOT_WS_L_N ⇒ > i
+ } // c ∈ NOT_WS_L_N ⇒ > i
}
- // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively.
- _ => i += char_len(b),
}
- // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1).
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: i as u32,
- }
- };
- w += 1;
+ // Mark | Connector | Punct | SymOther | NumericOther | Control (all in NOT_WS_L_N):
+ // rule 2 (prefix + `\p{L}+`) | rule 4
+ MRK | CON | PUN | SYM | NMO | CTL => {
+ let a = i + char_len(b);
+ if a < end && (tags[a] & 0x0F) == LET {
+ letters(a)
+ } else {
+ other(i)
+ } // c ∈ NOT_WS_L_N ⇒ > i
+ }
+ // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively.
+ _ => i + char_len(b),
+ }
+}
+
+/// The scan under [`fsm_cl100k_cap`]: hands each token to `emit` the moment it is cut,
+/// so a caller can consume tokens in place instead of collecting a span buffer first.
+pub fn scan_cl100k_cap(text: &[u8], tags: &[u8], digit_cap: usize, mut emit: impl FnMut(Span)) {
+ debug_assert!(tags.len() >= text.len());
+ let end = text.len();
+ // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
+ // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
+ let tags = &tags[..end];
+ let mut i = 0;
+ while i < end {
+ let e = advance_cl100k_cap(text, tags, i, end, digit_cap);
+ emit(Span {
+ start: i as u32,
+ end: e as u32,
+ });
+ i = e;
}
- w
}
diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs
index f54afb8fa..b130ebf0c 100644
--- a/tokenizers/atomsplit/src/fsm/deepseek.rs
+++ b/tokenizers/atomsplit/src/fsm/deepseek.rs
@@ -39,81 +39,187 @@ fn ds_is_cjk_at(text: &[u8], p: usize) -> bool {
/// `\w` but categorically `\p{S}`) take the `[\p{P}\p{S}]` path, not the letter run.
#[must_use]
pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
- debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
- // Leading-atom values as `const` → the `match` is a dense jump table (see `cl100k`). The Split
- // precedence (digits → CJK → big-regex alts) is preserved because the atom partition is disjoint.
- // `Mark` refined as an Other_Alphabetic symbol (Ⓘ …): coarse `LETTER_MARK`, but categorically `\p{S}`
- // — excluded from `[\p{L}\p{M}]`, routed to the `[\p{P}\p{S}]+` run instead (see `punct`).
- let end = text.len();
- // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
- // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
- let tags = &tags[..end];
- // maximal `[\p{L}\p{M}]+` run from `a`, stopping at CJK-range chars (Split-2 took those), ZWJ/ZWNJ
- // (not `\p{L}∪\p{M}` — see `ds_breaks`), and Other_Alphabetic symbols (`ASM`, categorically `\p{S}`).
- // BYTE-wise (`p += 1`, continuation bytes stay in-run, `ds_breaks` only fires at a lead) — the
- // `char_len`-per-char form was ~2× slower (see `run_end`'s note). Hot inner loop of the latin path.
- let letter_run = |a: usize| -> usize {
- let mut p = a;
- // ZWJ/ASM are now tags (no text peek); only the CJK-range exclusion still peeks text.
- while p < end {
- let t = tags[p];
- if t == CONT
- || (in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ && !ds_is_cjk_at(text, p))
- {
- p += 1;
- } else {
- break;
- }
- }
- p
- };
- // is `text[a]` the start of a deepseek letter/mark char (alt-2 run body / space-prefix target)?
- let is_lm = |a: usize| {
- a < end
- && in_mask(tags[a], mask::LETTER_MARK)
- && tags[a] != ASM
- && tags[a] != ZWJ
- && !ds_is_cjk_at(text, a)
- };
- // Split-3 alt-3 tail `[\p{P}\p{S}]+[\r\n]*` from `sp0` (a leading space is already consumed); `sp0`
- // if there is no punct/sym run there. STOPS at CJK-range chars — Split-1 isolated those, so a CJK
- // punct (・) is never merged into a non-CJK punct run (`!・` → `!`, `・`, not `!・`).
- let punct = |sp0: usize| -> usize {
- let mut p = sp0;
- while p < end
- && (in_mask(tags[p], mask::PUNCT_SYM) || tags[p] == ASM)
- && !ds_is_cjk_at(text, p)
+ debug_assert!(out.len() >= text.len());
+ let mut w = 0usize;
+ scan_deepseek(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// Maximal `[\p{L}\p{M}]+` run from `a`, stopping at CJK-range chars (Split-2 took those),
+/// ZWJ/ZWNJ (not `\p{L}∪\p{M}`), and Other_Alphabetic symbols (`ASM`, categorically `\p{S}`).
+/// BYTE-wise (`p += 1`, continuation bytes stay in-run) — the `char_len`-per-char form was ~2×
+/// slower (see `run_end`'s note). Hot inner loop of the latin path.
+#[inline(always)]
+fn letter_run(text: &[u8], tags: &[u8], a: usize, end: usize) -> usize {
+ let mut p = a;
+ // ZWJ/ASM are tags (no text peek); only the CJK-range exclusion still peeks text.
+ while p < end {
+ let t = tags[p];
+ if t == CONT
+ || (in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ && !ds_is_cjk_at(text, p))
{
+ p += 1;
+ } else {
+ break;
+ }
+ }
+ p
+}
+
+/// Is `text[a]` the start of a deepseek letter/mark char (alt-2 run body / space-prefix target)?
+#[inline(always)]
+fn is_lm(text: &[u8], tags: &[u8], a: usize, end: usize) -> bool {
+ a < end
+ && in_mask(tags[a], mask::LETTER_MARK)
+ && tags[a] != ASM
+ && tags[a] != ZWJ
+ && !ds_is_cjk_at(text, a)
+}
+
+/// Split-3 alt-3 tail `[\p{P}\p{S}]+[\r\n]*` from `sp0` (a leading space is already consumed);
+/// `sp0` if there is no punct/sym run there. STOPS at CJK-range chars — Split-1 isolated those,
+/// so a CJK punct (・) is never merged into a non-CJK punct run (`!・` → `!`, `・`, not `!・`).
+#[inline(always)]
+fn punct(text: &[u8], tags: &[u8], sp0: usize, end: usize) -> usize {
+ let mut p = sp0;
+ while p < end && (in_mask(tags[p], mask::PUNCT_SYM) || tags[p] == ASM) && !ds_is_cjk_at(text, p)
+ {
+ p += char_len(text[p]);
+ }
+ if p > sp0 {
+ while p < end && tags[p] == NLN {
p += char_len(text[p]);
}
- if p > sp0 {
- while p < end && tags[p] == NLN {
+ }
+ p
+}
+
+/// Split-3 alts d/e/f (whitespace). Unlike cl100k: a ws run FOLLOWED BY a digit/CJK is its own
+/// Sequence piece (Split-1/2 isolated the next match) → `\s+(?!\S)` takes the WHOLE run; only a
+/// following letter/punct (same Split-3 piece) leaves the last ws char for its ` ?`/`[^…]?`
+/// prefix.
+#[inline(always)]
+fn ds_ws(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
+ let re = run_end(tags, i, end, mask::WS);
+ let next_isolated = re < end && (in_mask(tags[re], mask::NUMBER) || ds_is_cjk_at(text, re));
+ if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) {
+ i + r + 1
+ } else if re == end || next_isolated {
+ re // whole ws run is one token
+ } else {
+ let mut last = re - 1;
+ while last > i && text[last] & 0xC0 == 0x80 {
+ last -= 1;
+ }
+ if last > i { last } else { re }
+ }
+}
+
+/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of
+/// the deepseek Sequence. The two multi-emit paths of [`scan_deepseek`] decompose per token: a
+/// gap run followed by letters ends at its LAST gap char (the next dispatch, on that char,
+/// takes the prefix-plus-letters path), and a CJK run is one same-kind sub-run per call. The
+/// masked scanner re-derives tokens with this where its batch masks are not trustworthy.
+#[inline(always)]
+pub(super) fn advance_deepseek(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
+ if ds_is_cjk_at(text, i) {
+ let is_letter = in_mask(tags[i], mask::LETTER_MARK);
+ let mut p = i + 3; // CJK-range chars are all 3-byte (leads E3..E9)
+ while p < end && ds_is_cjk_at(text, p) && in_mask(tags[p], mask::LETTER_MARK) == is_letter {
+ p += 3;
+ }
+ return p;
+ }
+ if matches!(tags[i] & 0x0F, NMO | CTL) || tags[i] == ZWJ {
+ let (mut p, mut last) = (i, i);
+ while p < end && (matches!(tags[p] & 0x0F, NMO | CTL) || tags[p] == ZWJ) {
+ last = p;
+ p += char_len(text[p]);
+ }
+ return if is_lm(text, tags, p, end) {
+ if last > i {
+ last // gap sans the prefix char
+ } else {
+ letter_run(text, tags, p, end) // prefix char + `[\p{L}\p{M}]+`
+ }
+ } else {
+ p
+ };
+ }
+ let b = text[i];
+ match tags[i] & 0x0F {
+ NW | NO => {
+ let (mut p, mut cnt) = (i, 0);
+ while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) {
p += char_len(text[p]);
+ cnt += 1;
}
+ p
}
- p
- };
- // Split-3 alts d/e/f (whitespace). Unlike cl100k: a ws run FOLLOWED BY a digit/CJK is its own
- // Sequence piece (Split-1/2 isolated the next match) → `\s+(?!\S)` takes the WHOLE run; only a
- // following letter/punct (same Split-3 piece) leaves the last ws char for its ` ?`/`[^…]?` prefix.
- let ws = |i: usize| -> usize {
- let re = run_end(tags, i, end, mask::WS);
- let next_isolated = re < end && (in_mask(tags[re], mask::NUMBER) || ds_is_cjk_at(text, re));
- if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) {
- i + r + 1
- } else if re == end || next_isolated {
- re // whole ws run is one token
- } else {
- let mut last = re - 1;
- while last > i && text[last] & 0xC0 == 0x80 {
- last -= 1;
+ LET | MRK => {
+ if tags[i] == ASM {
+ punct(text, tags, i, end)
+ } else {
+ letter_run(text, tags, i, end)
+ }
+ }
+ SPC => {
+ let a = i + 1; // Space is ASCII (0x20)
+ if is_lm(text, tags, a, end) {
+ letter_run(text, tags, a, end)
+ } else if a < end && ds_is_cjk_at(text, a) {
+ ds_ws(text, tags, i, end)
+ } else {
+ let p = punct(text, tags, a, end);
+ if p > a { p } else { ds_ws(text, tags, i, end) }
+ }
+ }
+ WSO => {
+ let a = i + char_len(b);
+ if is_lm(text, tags, a, end) {
+ letter_run(text, tags, a, end)
+ } else {
+ ds_ws(text, tags, i, end)
}
- if last > i { last } else { re }
}
- };
+ NLN => ds_ws(text, tags, i, end),
+ CON | PUN | APO | SYM => {
+ if b.is_ascii_punctuation() && i + 1 < end && text[i + 1].is_ascii_alphabetic() {
+ let mut p = i + 1;
+ while p < end && text[p].is_ascii_alphabetic() {
+ p += 1;
+ }
+ p
+ } else {
+ punct(text, tags, i, end) // c ∈ PUNCT_SYM ⇒ > i
+ }
+ }
+ _ => i + char_len(b),
+ }
+}
+
+/// The scan under [`fsm_deepseek`]: hands each token to `emit` the moment it is cut,
+/// so a caller can consume tokens in place instead of collecting a span buffer first.
+pub fn scan_deepseek(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) {
+ debug_assert!(tags.len() >= text.len());
+ // Leading-atom values as `const` → the `match` is a dense jump table (see `cl100k`). The Split
+ // precedence (digits → CJK → big-regex alts) is preserved because the atom partition is disjoint.
+ // `Mark` refined as an Other_Alphabetic symbol (Ⓘ …): coarse `LETTER_MARK`, but categorically `\p{S}`
+ // — excluded from `[\p{L}\p{M}]`, routed to the `[\p{P}\p{S}]+` run instead (see `punct`).
+ let end = text.len();
+ // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
+ // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
+ let tags = &tags[..end];
+ let letter_run = |a: usize| letter_run(text, tags, a, end);
+ let is_lm = |a: usize| is_lm(text, tags, a, end);
+ let punct = |sp0: usize| punct(text, tags, sp0, end);
+ let ws = |i: usize| ds_ws(text, tags, i, end);
let mut i = 0;
- let mut w = 0usize;
while i < end {
let start = i;
let b = text[i];
@@ -129,13 +235,10 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
{
p += 3;
}
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: p as u32,
- }
- };
- w += 1;
+ emit(Span {
+ start: start as u32,
+ end: p as u32,
+ });
i = p;
continue;
}
@@ -150,31 +253,22 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
}
if is_lm(p) {
if last > i {
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: last as u32,
- }
- }; // gap sans prefix char
- w += 1;
+ emit(Span {
+ start: start as u32,
+ end: last as u32,
+ }); // gap sans prefix char
}
let e = letter_run(p);
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: last as u32,
- end: e as u32,
- }
- }; // prefix char + `[\p{L}\p{M}]+`
- w += 1;
+ emit(Span {
+ start: last as u32,
+ end: e as u32,
+ }); // prefix char + `[\p{L}\p{M}]+`
i = e;
} else {
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: p as u32,
- }
- }; // whole gap run is one piece
- w += 1;
+ emit(Span {
+ start: start as u32,
+ end: p as u32,
+ }); // whole gap run is one piece
i = p;
}
continue;
@@ -234,14 +328,9 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
// Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively.
_ => i += char_len(b),
}
- // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1).
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: i as u32,
- }
- };
- w += 1;
+ emit(Span {
+ start: start as u32,
+ end: i as u32,
+ });
}
- w
}
diff --git a/tokenizers/atomsplit/src/fsm/masked.rs b/tokenizers/atomsplit/src/fsm/masked.rs
new file mode 100644
index 000000000..aa6f1795a
--- /dev/null
+++ b/tokenizers/atomsplit/src/fsm/masked.rs
@@ -0,0 +1,440 @@
+//! Boundary-mask scanners: the SIMD replacement for the per-run scalar walk of the regex-shaped
+//! FSMs (one scheme module per regex family, sharing this walker).
+//!
+//! A scalar FSM advances one token at a time: a rule dispatch, then a per-byte run scan. A
+//! masked scanner instead classifies 64 bytes at once into per-class bitmasks (bit k set =
+//! "byte k is a letter") and computes every token start in the batch with a few dozen
+//! branch-free u64 operations. That works because these regexes are class-run languages with
+//! one or two chars of context: a token starts exactly at a class change that is not an
+//! absorbed prefix, at a whitespace-run edge, or at a contraction edge, and each of those
+//! conditions is a shifted-mask expression. The boundary algebra is transcribed per scheme from
+//! gigatoken's `src/pretokenize/fast/` scanners (MIT), with one structural difference:
+//! gigatoken classifies raw bytes in-batch and falls back to a per-char loop whenever a batch
+//! contains a non-ASCII byte, while these scanners build their masks from the
+//! [`crate::classify`] tag stream, which is already SIMD and covers all of Unicode. A
+//! continuation byte's tag is [`Atom::Cont`]; the fill step gives it its char's class, so byte
+//! adjacency equals char adjacency and the same algebra applies to non-ASCII batches.
+//!
+//! # Trust boundaries
+//!
+//! Each scheme's `batch_masks` returns `(boundary, bad)` for one 64-byte batch. A `boundary`
+//! bit is a proven token start. A `bad` bit means the algebra cannot decide that byte (batch-
+//! edge straddles, char-counted rules over multi-byte chars, run-contextual classes; each
+//! scheme documents its own list). `boundary & bad` is always 0, and no span is emitted across
+//! a bad zone: the walker re-derives tokens there with the scheme's scalar `advance`, the same
+//! rules the plain scan runs, and resumes on masks at the next batch. The scalar scans stay the
+//! ground truth; the `masked_*` tests (tests/fsm.rs) pin byte-exactness at every batch-edge
+//! offset.
+//!
+//! # Targets
+//!
+//! The per-target work is confined to [`block`] (64-byte predicate masks): aarch64 NEON,
+//! x86_64 SSE2 and wasm32 simd128 have kernels; every other target delegates the
+//! `scan_*_masked` entry points to the scalar scans, so correctness never depends on a SIMD
+//! path being present.
+//!
+//! Inputs must be well-formed UTF-8 (the crate-level contract); the fill step relies on
+//! continuation runs of at most 3 bytes.
+
+use super::*;
+
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+mod block;
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+mod byte_level;
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+mod cl100k;
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+mod deepseek;
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+mod o200k;
+
+/// [`fsm_byte_level`] over the masked scanner: writes spans into `out` (len >= `text.len()`)
+/// and returns the count.
+#[must_use]
+pub fn fsm_byte_level_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
+ debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
+ let mut w = 0usize;
+ scan_byte_level_masked(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// The masked twin of [`scan_byte_level`]: same tokens, same emit order. Targets without a
+/// [`block`] kernel delegate to the scalar scan.
+pub fn scan_byte_level_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) {
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ walk(&byte_level::ByteLevelMasked, text, tags, emit);
+ #[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ )))]
+ scan_byte_level(text, tags, emit);
+}
+
+/// [`fsm_cl100k_cap`] over the masked scanner: writes spans into `out` (len >= `text.len()`)
+/// and returns the count.
+#[must_use]
+pub fn fsm_cl100k_cap_masked(
+ text: &[u8],
+ tags: &[u8],
+ out: &mut [Span],
+ digit_cap: usize,
+) -> usize {
+ debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
+ let mut w = 0usize;
+ scan_cl100k_cap_masked(text, tags, digit_cap, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// The masked twin of [`scan_cl100k_cap`]: same tokens, same emit order. Digit caps other than
+/// 1, 3 and `usize::MAX` (none ship today) and targets without a [`block`] kernel delegate to
+/// the scalar scan.
+pub fn scan_cl100k_cap_masked(text: &[u8], tags: &[u8], digit_cap: usize, emit: impl FnMut(Span)) {
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ {
+ if matches!(digit_cap, 1 | 3 | usize::MAX) {
+ walk(&cl100k::Cl100kMasked { digit_cap }, text, tags, emit);
+ } else {
+ scan_cl100k_cap(text, tags, digit_cap, emit);
+ }
+ }
+ #[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ )))]
+ scan_cl100k_cap(text, tags, digit_cap, emit);
+}
+
+/// [`fsm_o200k`] over the masked scanner: writes spans into `out` (len >= `text.len()`) and
+/// returns the count.
+#[must_use]
+pub fn fsm_o200k_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
+ debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
+ let mut w = 0usize;
+ scan_o200k_masked(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// [`fsm_tekken`] over the masked scanner: writes spans into `out` (len >= `text.len()`) and
+/// returns the count.
+#[must_use]
+pub fn fsm_tekken_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
+ debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
+ let mut w = 0usize;
+ scan_tekken_masked(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// The masked twin of [`scan_o200k`]: same tokens, same emit order. Targets without a
+/// [`block`] kernel delegate to the scalar scan.
+pub fn scan_o200k_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) {
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ walk(&o200k::O200kMasked::, text, tags, emit);
+ #[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ )))]
+ scan_o200k(text, tags, emit);
+}
+
+/// The masked twin of [`scan_tekken`]: same tokens, same emit order. Targets without a
+/// [`block`] kernel delegate to the scalar scan.
+pub fn scan_tekken_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) {
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ walk(&o200k::O200kMasked::, text, tags, emit);
+ #[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ )))]
+ scan_tekken(text, tags, emit);
+}
+
+/// [`fsm_deepseek`] over the masked scanner: writes spans into `out` (len >= `text.len()`) and
+/// returns the count.
+#[must_use]
+pub fn fsm_deepseek_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
+ debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
+ let mut w = 0usize;
+ scan_deepseek_masked(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// The masked twin of [`scan_deepseek`]: same tokens, same emit order. Targets without a
+/// [`block`] kernel delegate to the scalar scan.
+pub fn scan_deepseek_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) {
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ walk(&deepseek::DeepSeekMasked, text, tags, emit);
+ #[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ )))]
+ scan_deepseek(text, tags, emit);
+}
+
+/// One masked scheme: the batch classifier and the scalar rules the walker falls back on.
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+trait MaskedFsm {
+ /// `(boundary, bad)` for `text[scan..scan + 64]`: `boundary` bit k = proven token start at
+ /// `scan + k`, `bad` bit k = the walker must re-derive byte `scan + k` with
+ /// [`Self::advance`]. `boundary & bad` must be 0. Callers guarantee `scan + 64 <
+ /// text.len()` (one lookahead tag is readable).
+ fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64);
+
+ /// Scalar ground truth: end of the token starting at `i` (`i < end`, `i` on a token
+ /// boundary).
+ fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize;
+}
+
+/// The batch walker: consume proven token starts batch by batch, re-derive bad zones with the
+/// scheme's scalar rules. `pending` is always the start of the open (not yet emitted) token.
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+fn walk(scheme: &impl MaskedFsm, text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) {
+ let end = text.len();
+ let tags = &tags[..end];
+ let mut pending = 0usize;
+ let mut scan = 0usize;
+ // Each batch reads the lookahead tag at scan + 64 (the `\s+(?!\S)` bit-63 rules), so the
+ // last <= 64 bytes always go through the scalar tail below.
+ while scan + 64 < end {
+ if scan + 64 <= pending {
+ // A scalar re-derivation ran past this whole batch.
+ scan += 64;
+ continue;
+ }
+ let (boundary, mut bad) = scheme.batch_masks(text, tags, scan);
+ let mut m = boundary;
+ if pending > scan {
+ // Bits at or below `pending` are inside or at the start of the open token; they are
+ // stale leftovers of a scalar re-derivation that entered this batch.
+ let done = pending - scan + 1;
+ m = if done >= 64 {
+ 0
+ } else {
+ m & (u64::MAX << done)
+ };
+ }
+ // Interleave: consume proven starts below the next bad zone, re-derive the zone with
+ // the scalar rules, repeat. A span must never be emitted across an unresolved zone, so
+ // starts above one cannot pair with `pending` from below it.
+ loop {
+ let zone = if bad == 0 {
+ 64
+ } else {
+ bad.trailing_zeros() as usize
+ };
+ while m != 0 {
+ let j = m.trailing_zeros() as usize;
+ if j >= zone {
+ break;
+ }
+ let p = scan + j;
+ if p > pending {
+ emit(Span {
+ start: pending as u32,
+ end: p as u32,
+ });
+ pending = p;
+ }
+ m &= m - 1;
+ }
+ if bad == 0 {
+ break;
+ }
+ // The zone's contiguous extent; the scalar rules resolve through its end (their
+ // tokens may overshoot it, or the whole batch — later bits fall to the guards).
+ let zone_end = zone + ((!(bad >> zone)).trailing_zeros() as usize).min(64 - zone);
+ while pending < scan + zone_end {
+ let e = scheme.advance(text, tags, pending, end);
+ emit(Span {
+ start: pending as u32,
+ end: e as u32,
+ });
+ pending = e;
+ }
+ bad = if zone_end >= 64 {
+ 0
+ } else {
+ bad & (u64::MAX << zone_end)
+ };
+ }
+ scan += 64;
+ }
+ while pending < end {
+ let e = scheme.advance(text, tags, pending, end);
+ emit(Span {
+ start: pending as u32,
+ end: e as u32,
+ });
+ pending = e;
+ }
+}
+
+// ── shared u64 helpers (platform-independent; the scheme modules compose these) ────────────────
+
+/// The two continuation-run masks the fill steps need: `c2` bit k = bytes k and k-1 are both
+/// continuations, `c3` = bytes k, k-1, k-2 all are.
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+#[inline(always)]
+fn cont_runs(c: u64) -> (u64, u64) {
+ let c2 = c & (c << 1);
+ (c2, c2 & (c << 2))
+}
+
+/// Fill: every continuation byte of a char whose lead is in `m` joins `m`, so byte adjacency
+/// equals char adjacency (UTF-8 chars are at most 4 bytes: 3 hops cover every continuation).
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+#[inline(always)]
+fn fill(m: u64, c: u64, c2: u64, c3: u64) -> u64 {
+ m | ((m << 1) & c) | ((m << 2) & c2) | ((m << 3) & c3)
+}
+
+/// Smear `seed` upward (toward higher bits) through contiguous set bits of `within`, in log
+/// steps (via gigatoken's `cl100k_family.rs`, MIT).
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+#[inline(always)]
+fn smear_up(seed: u64, within: u64) -> u64 {
+ let mut a = seed;
+ let mut m = within;
+ let mut sh = 1u32;
+ while sh < 64 {
+ a |= (a << sh) & m;
+ m &= m << sh;
+ sh <<= 1;
+ }
+ a
+}
+
+/// Token-start bits inside ASCII digit runs for `\p{N}{1,3}`: each run splits into 3-char
+/// tokens, so boundaries sit at run start + 3k (via gigatoken's `mask.rs`, MIT). Callers keep
+/// multi-byte digit chars out of `d`: their grouping is char-counted, and byte hops would
+/// misphase it.
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+#[inline(always)]
+fn digit_run_splits3(d: u64) -> u64 {
+ let mut b = d & !(d << 1); // run starts
+ // A start at p re-arms at p+3 while the run continues: hop condition c = "p..p+3 all
+ // digits". Log-doubling covers 64-bit runs in 5 steps.
+ let mut c = d & (d >> 1) & (d >> 2) & (d >> 3);
+ let mut sh = 3u32;
+ while sh < 64 {
+ b |= (b & c) << sh;
+ c &= c >> sh;
+ sh <<= 1;
+ }
+ b
+}
+
+/// `x << n`, saturating to 0 at `n >= 64` (`trailing_zeros` on an empty mask yields 64).
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+#[inline(always)]
+fn shl_sat(x: u64, n: u32) -> u64 {
+ if n >= 64 { 0 } else { x << n }
+}
+
+/// The lead of the char containing `tags[p]`: walk back over continuation tags (at most 3 on
+/// well-formed UTF-8; `p` itself may already be the lead).
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+#[inline(always)]
+fn char_lead(tags: &[u8], mut p: usize) -> usize {
+ while p > 0 && tags[p] & 0x0F == CONT {
+ p -= 1;
+ }
+ p
+}
diff --git a/tokenizers/atomsplit/src/fsm/masked/block.rs b/tokenizers/atomsplit/src/fsm/masked/block.rs
new file mode 100644
index 000000000..3e06e7c89
--- /dev/null
+++ b/tokenizers/atomsplit/src/fsm/masked/block.rs
@@ -0,0 +1,427 @@
+//! Per-target 64-byte block classifiers for the masked scanners. A block loads 64 consecutive
+//! bytes into four SIMD registers once; each method then answers one per-byte predicate for the
+//! whole block as a u64 bitmask (bit k = byte k passes). Everything downstream of these masks is
+//! platform-independent u64 arithmetic in the scheme modules, so this file is the entire
+//! per-target surface.
+//!
+//! Three targets have a kernel: aarch64 (NEON, baseline), x86_64 (SSE2, baseline) and wasm32
+//! with `simd128`. Any other target never reaches this module — the `scan_*_masked` entry
+//! points delegate to the scalar scans there.
+
+#![allow(dead_code)] // arch-gated: each build compiles one target's kernel, and schemes land
+// one by one, so some predicates are unused until their scheme arrives.
+
+// ── aarch64 / NEON ──────────────────────────────────────────────────────────────────────────────
+#[cfg(target_arch = "aarch64")]
+pub(crate) use neon::Block;
+
+#[cfg(target_arch = "aarch64")]
+mod neon {
+ use core::arch::aarch64::*;
+
+ /// 64 bytes in four NEON registers. `tag` methods fold the refinement nibble away first
+ /// (`& 0x0F`), `full`/byte methods compare the raw byte.
+ pub(crate) struct Block {
+ v: [uint8x16_t; 4],
+ }
+
+ /// simdjson's arm64 movemask: 4 mask vectors (64 lanes of 0x00/0xFF) to one u64, bit i =
+ /// lane i. The 4-`addp` reduction is pinned as asm (via gigatoken's `mask.rs`, MIT): written
+ /// with `vpaddq_u8`, LLVM rewrites the pairwise adds into uzp1/uzp2/orr triples and the call
+ /// grows from 9 to 17 vector ops.
+ #[inline(always)]
+ unsafe fn movemask64(v0: uint8x16_t, v1: uint8x16_t, v2: uint8x16_t, v3: uint8x16_t) -> u64 {
+ // SAFETY: pure NEON register arithmetic, no memory access beyond the 16-byte constant.
+ unsafe {
+ const W: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
+ let w = vld1q_u8(W.as_ptr());
+ let mut a0 = vandq_u8(v0, w);
+ let a1 = vandq_u8(v1, w);
+ let a2 = vandq_u8(v2, w);
+ let a3 = vandq_u8(v3, w);
+ core::arch::asm!(
+ "addp {a0:v}.16b, {a0:v}.16b, {a1:v}.16b",
+ "addp {a2:v}.16b, {a2:v}.16b, {a3:v}.16b",
+ "addp {a0:v}.16b, {a0:v}.16b, {a2:v}.16b",
+ "addp {a0:v}.16b, {a0:v}.16b, {a0:v}.16b",
+ a0 = inout(vreg) a0,
+ a1 = in(vreg) a1,
+ a2 = inout(vreg) a2 => _,
+ a3 = in(vreg) a3,
+ options(pure, nomem, nostack, preserves_flags),
+ );
+ vgetq_lane_u64::<0>(vreinterpretq_u64_u8(a0))
+ }
+ }
+
+ impl Block {
+ /// Load `bytes[at..at + 64]`.
+ ///
+ /// # Safety
+ ///
+ /// `at + 64 <= bytes.len()`. NEON loads are alignment-free.
+ #[inline(always)]
+ pub(crate) unsafe fn load(bytes: &[u8], at: usize) -> Self {
+ debug_assert!(at + 64 <= bytes.len());
+ // SAFETY: the fn contract puts all four 16-byte loads in bounds.
+ unsafe {
+ let p = bytes.as_ptr().add(at);
+ Self {
+ v: [
+ vld1q_u8(p),
+ vld1q_u8(p.add(16)),
+ vld1q_u8(p.add(32)),
+ vld1q_u8(p.add(48)),
+ ],
+ }
+ }
+ }
+
+ #[inline(always)]
+ fn mask(&self, f: impl Fn(uint8x16_t) -> uint8x16_t) -> u64 {
+ // SAFETY: register arithmetic only.
+ unsafe { movemask64(f(self.v[0]), f(self.v[1]), f(self.v[2]), f(self.v[3])) }
+ }
+
+ #[inline(always)]
+ fn any(&self, f: impl Fn(uint8x16_t) -> uint8x16_t) -> bool {
+ // SAFETY: register arithmetic only.
+ unsafe {
+ let o = vorrq_u8(
+ vorrq_u8(f(self.v[0]), f(self.v[1])),
+ vorrq_u8(f(self.v[2]), f(self.v[3])),
+ );
+ vmaxvq_u8(o) != 0
+ }
+ }
+
+ /// Bytes whose low nibble equals `k`.
+ #[inline(always)]
+ pub(crate) fn eq_tag(&self, k: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe { vceqq_u8(vandq_u8(v, vdupq_n_u8(0x0F)), vdupq_n_u8(k)) })
+ }
+
+ /// Bytes in `lo..=lo + span` (the raw byte — text blocks).
+ #[inline(always)]
+ pub(crate) fn range_full(&self, lo: u8, span: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe { vcleq_u8(vsubq_u8(v, vdupq_n_u8(lo)), vdupq_n_u8(span)) })
+ }
+
+ /// Bytes whose low nibble is in `lo..=lo + span`.
+ #[inline(always)]
+ pub(crate) fn range_tag(&self, lo: u8, span: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ vcleq_u8(
+ vsubq_u8(vandq_u8(v, vdupq_n_u8(0x0F)), vdupq_n_u8(lo)),
+ vdupq_n_u8(span),
+ )
+ })
+ }
+
+ /// Bytes equal to `k` (the raw byte — refined tags, or text bytes).
+ #[inline(always)]
+ pub(crate) fn eq_full(&self, k: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe { vceqq_u8(v, vdupq_n_u8(k)) })
+ }
+
+ /// Any byte with low nibble in `lo..=lo + span`? (Cheaper than `range_tag` when only
+ /// presence matters — no movemask.)
+ #[inline(always)]
+ pub(crate) fn any_range_tag(&self, lo: u8, span: u8) -> bool {
+ // SAFETY: register arithmetic only.
+ self.any(|v| unsafe {
+ vcleq_u8(
+ vsubq_u8(vandq_u8(v, vdupq_n_u8(0x0F)), vdupq_n_u8(lo)),
+ vdupq_n_u8(span),
+ )
+ })
+ }
+
+ /// Any byte equal to `k`?
+ #[inline(always)]
+ pub(crate) fn any_eq_full(&self, k: u8) -> bool {
+ // SAFETY: register arithmetic only.
+ self.any(|v| unsafe { vceqq_u8(v, vdupq_n_u8(k)) })
+ }
+
+ /// ASCII letters `[A-Za-z]` (text blocks).
+ #[inline(always)]
+ pub(crate) fn ascii_alpha(&self) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ vcleq_u8(
+ vsubq_u8(vorrq_u8(v, vdupq_n_u8(0x20)), vdupq_n_u8(b'a')),
+ vdupq_n_u8(25),
+ )
+ })
+ }
+
+ /// ASCII punctuation (the four `is_ascii_punctuation` ranges, text blocks).
+ #[inline(always)]
+ pub(crate) fn ascii_punct(&self) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ let r = |v: uint8x16_t, lo: u8, span: u8| {
+ vcleq_u8(vsubq_u8(v, vdupq_n_u8(lo)), vdupq_n_u8(span))
+ };
+ vorrq_u8(
+ vorrq_u8(r(v, 0x21, 0x0E), r(v, 0x3A, 0x06)),
+ vorrq_u8(r(v, 0x5B, 0x05), r(v, 0x7B, 0x03)),
+ )
+ })
+ }
+ }
+}
+
+// ── x86_64 / SSE2 ──────────────────────────────────────────────────────────────────────────────
+#[cfg(target_arch = "x86_64")]
+pub(crate) use sse2::Block;
+
+#[cfg(target_arch = "x86_64")]
+mod sse2 {
+ use core::arch::x86_64::*;
+
+ /// 64 bytes in four SSE2 registers; `pmovmskb` is the native movemask, 16 bits per register.
+ /// SSE2 is baseline on x86_64, so no runtime detection is needed. (SSE2 has no unsigned
+ /// byte compare: `x <= span` is done as `max(x, span) == span`.)
+ pub(crate) struct Block {
+ v: [__m128i; 4],
+ }
+
+ #[inline(always)]
+ fn movemask64(v0: __m128i, v1: __m128i, v2: __m128i, v3: __m128i) -> u64 {
+ // SAFETY: register arithmetic only; SSE2 is baseline on x86_64.
+ unsafe {
+ (_mm_movemask_epi8(v0) as u16 as u64)
+ | ((_mm_movemask_epi8(v1) as u16 as u64) << 16)
+ | ((_mm_movemask_epi8(v2) as u16 as u64) << 32)
+ | ((_mm_movemask_epi8(v3) as u16 as u64) << 48)
+ }
+ }
+
+ #[inline(always)]
+ fn le(x: __m128i, span: u8) -> __m128i {
+ // SAFETY: register arithmetic only.
+ unsafe {
+ let s = _mm_set1_epi8(span as i8);
+ _mm_cmpeq_epi8(_mm_max_epu8(x, s), s)
+ }
+ }
+
+ impl Block {
+ /// Load `bytes[at..at + 64]`.
+ ///
+ /// # Safety
+ ///
+ /// `at + 64 <= bytes.len()`. Unaligned loads (`loadu`).
+ #[inline(always)]
+ pub(crate) unsafe fn load(bytes: &[u8], at: usize) -> Self {
+ debug_assert!(at + 64 <= bytes.len());
+ // SAFETY: the fn contract puts all four 16-byte loads in bounds.
+ unsafe {
+ let p = bytes.as_ptr().add(at);
+ Self {
+ v: [
+ _mm_loadu_si128(p.cast()),
+ _mm_loadu_si128(p.add(16).cast()),
+ _mm_loadu_si128(p.add(32).cast()),
+ _mm_loadu_si128(p.add(48).cast()),
+ ],
+ }
+ }
+ }
+
+ #[inline(always)]
+ fn mask(&self, f: impl Fn(__m128i) -> __m128i) -> u64 {
+ movemask64(f(self.v[0]), f(self.v[1]), f(self.v[2]), f(self.v[3]))
+ }
+
+ #[inline(always)]
+ pub(crate) fn eq_tag(&self, k: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ _mm_cmpeq_epi8(
+ _mm_and_si128(v, _mm_set1_epi8(0x0F)),
+ _mm_set1_epi8(k as i8),
+ )
+ })
+ }
+
+ #[inline(always)]
+ pub(crate) fn range_tag(&self, lo: u8, span: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ le(
+ _mm_sub_epi8(
+ _mm_and_si128(v, _mm_set1_epi8(0x0F)),
+ _mm_set1_epi8(lo as i8),
+ ),
+ span,
+ )
+ })
+ }
+
+ #[inline(always)]
+ pub(crate) fn range_full(&self, lo: u8, span: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe { le(_mm_sub_epi8(v, _mm_set1_epi8(lo as i8)), span) })
+ }
+
+ #[inline(always)]
+ pub(crate) fn eq_full(&self, k: u8) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe { _mm_cmpeq_epi8(v, _mm_set1_epi8(k as i8)) })
+ }
+
+ #[inline(always)]
+ pub(crate) fn any_range_tag(&self, lo: u8, span: u8) -> bool {
+ self.range_tag(lo, span) != 0
+ }
+
+ #[inline(always)]
+ pub(crate) fn any_eq_full(&self, k: u8) -> bool {
+ self.eq_full(k) != 0
+ }
+
+ #[inline(always)]
+ pub(crate) fn ascii_alpha(&self) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ le(
+ _mm_sub_epi8(
+ _mm_or_si128(v, _mm_set1_epi8(0x20)),
+ _mm_set1_epi8(b'a' as i8),
+ ),
+ 25,
+ )
+ })
+ }
+
+ #[inline(always)]
+ pub(crate) fn ascii_punct(&self) -> u64 {
+ // SAFETY: register arithmetic only.
+ self.mask(|v| unsafe {
+ let r = |v: __m128i, lo: u8, span: u8| {
+ le(_mm_sub_epi8(v, _mm_set1_epi8(lo as i8)), span)
+ };
+ _mm_or_si128(
+ _mm_or_si128(r(v, 0x21, 0x0E), r(v, 0x3A, 0x06)),
+ _mm_or_si128(r(v, 0x5B, 0x05), r(v, 0x7B, 0x03)),
+ )
+ })
+ }
+ }
+}
+
+// ── wasm32 / SIMD128 ────────────────────────────────────────────────────────────────────────────
+#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
+pub(crate) use wasm::Block;
+
+#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
+mod wasm {
+ use core::arch::wasm32::*;
+
+ /// 64 bytes in four v128 registers; `u8x16_bitmask` is the native movemask.
+ pub(crate) struct Block {
+ v: [v128; 4],
+ }
+
+ #[inline(always)]
+ fn movemask64(v0: v128, v1: v128, v2: v128, v3: v128) -> u64 {
+ (u8x16_bitmask(v0) as u64)
+ | ((u8x16_bitmask(v1) as u64) << 16)
+ | ((u8x16_bitmask(v2) as u64) << 32)
+ | ((u8x16_bitmask(v3) as u64) << 48)
+ }
+
+ impl Block {
+ /// Load `bytes[at..at + 64]`.
+ ///
+ /// # Safety
+ ///
+ /// `at + 64 <= bytes.len()`. `v128_load` is alignment-free.
+ #[inline(always)]
+ pub(crate) unsafe fn load(bytes: &[u8], at: usize) -> Self {
+ debug_assert!(at + 64 <= bytes.len());
+ // SAFETY: the fn contract puts all four 16-byte loads in bounds.
+ unsafe {
+ let p = bytes.as_ptr().add(at);
+ Self {
+ v: [
+ v128_load(p.cast()),
+ v128_load(p.add(16).cast()),
+ v128_load(p.add(32).cast()),
+ v128_load(p.add(48).cast()),
+ ],
+ }
+ }
+ }
+
+ #[inline(always)]
+ fn mask(&self, f: impl Fn(v128) -> v128) -> u64 {
+ movemask64(f(self.v[0]), f(self.v[1]), f(self.v[2]), f(self.v[3]))
+ }
+
+ #[inline(always)]
+ pub(crate) fn eq_tag(&self, k: u8) -> u64 {
+ self.mask(|v| u8x16_eq(v128_and(v, u8x16_splat(0x0F)), u8x16_splat(k)))
+ }
+
+ #[inline(always)]
+ pub(crate) fn range_tag(&self, lo: u8, span: u8) -> u64 {
+ self.mask(|v| {
+ u8x16_le(
+ u8x16_sub(v128_and(v, u8x16_splat(0x0F)), u8x16_splat(lo)),
+ u8x16_splat(span),
+ )
+ })
+ }
+
+ #[inline(always)]
+ pub(crate) fn range_full(&self, lo: u8, span: u8) -> u64 {
+ self.mask(|v| u8x16_le(u8x16_sub(v, u8x16_splat(lo)), u8x16_splat(span)))
+ }
+
+ #[inline(always)]
+ pub(crate) fn eq_full(&self, k: u8) -> u64 {
+ self.mask(|v| u8x16_eq(v, u8x16_splat(k)))
+ }
+
+ #[inline(always)]
+ pub(crate) fn any_range_tag(&self, lo: u8, span: u8) -> bool {
+ self.range_tag(lo, span) != 0
+ }
+
+ #[inline(always)]
+ pub(crate) fn any_eq_full(&self, k: u8) -> bool {
+ self.eq_full(k) != 0
+ }
+
+ #[inline(always)]
+ pub(crate) fn ascii_alpha(&self) -> u64 {
+ self.mask(|v| {
+ u8x16_le(
+ u8x16_sub(v128_or(v, u8x16_splat(0x20)), u8x16_splat(b'a')),
+ u8x16_splat(25),
+ )
+ })
+ }
+
+ #[inline(always)]
+ pub(crate) fn ascii_punct(&self) -> u64 {
+ self.mask(|v| {
+ let r = |v: v128, lo: u8, span: u8| {
+ u8x16_le(u8x16_sub(v, u8x16_splat(lo)), u8x16_splat(span))
+ };
+ v128_or(
+ v128_or(r(v, 0x21, 0x0E), r(v, 0x3A, 0x06)),
+ v128_or(r(v, 0x5B, 0x05), r(v, 0x7B, 0x03)),
+ )
+ })
+ }
+ }
+}
diff --git a/tokenizers/atomsplit/src/fsm/masked/byte_level.rs b/tokenizers/atomsplit/src/fsm/masked/byte_level.rs
new file mode 100644
index 000000000..900885434
--- /dev/null
+++ b/tokenizers/atomsplit/src/fsm/masked/byte_level.rs
@@ -0,0 +1,146 @@
+//! Masked scheme for the byte-level (GPT-2) regex — the r50k boundary algebra (via gigatoken's
+//! `r50k.rs`, MIT) over tag-fed class masks.
+//!
+//! Every byte-level rule is local: a token starts exactly at a class change that is not a space
+//! prefix, at the first byte of a whitespace run, at the last whitespace byte before a
+//! non-whitespace (the `\s+(?!\S)` give-back), or at a contraction edge. Bad zones: a
+//! whitespace char straddling the batch edge (its give-back is per char, not per byte), an
+//! apostrophe too close to the edge for the contraction peek, a multi-byte whitespace char, or
+//! a `Sentinel`/`MultiByte` tag.
+
+use super::super::byte_level::advance_byte_level;
+use super::block::Block;
+use super::{MaskedFsm, char_lead, cont_runs, fill};
+use crate::fsm::{APO, CONT, LET, NLN, NO, NW, SPC, WSO, in_mask, mask};
+
+pub(super) struct ByteLevelMasked;
+
+impl MaskedFsm for ByteLevelMasked {
+ #[inline(always)]
+ fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) {
+ batch_masks(text, tags, scan)
+ }
+
+ #[inline(always)]
+ fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
+ advance_byte_level(text, tags, i, end)
+ }
+}
+
+#[inline(always)]
+fn batch_masks(text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) {
+ debug_assert!(scan + 64 < tags.len() && tags.len() == text.len());
+ // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract.
+ let b = unsafe { Block::load(tags, scan) };
+ if b.any_range_tag(13, 1) {
+ // Sentinel / MultiByte: the scalar dispatch has a defensive arm for these; keep its
+ // behavior by refusing the whole batch.
+ return (0, u64::MAX);
+ }
+ let l0 = b.eq_tag(LET);
+ let d0 = b.range_tag(NW, 1);
+ let ws0 = b.range_tag(NLN, 2);
+ let s = b.eq_tag(SPC);
+ // Apostrophes and continuations only matter to the fixups below; skip their movemask when
+ // an any-test says the batch has none.
+ let ap = if b.any_eq_full(APO) {
+ b.eq_full(APO)
+ } else {
+ 0
+ };
+ let c = if b.any_eq_full(CONT) {
+ b.eq_full(CONT)
+ } else {
+ 0
+ };
+ let (mut l, mut d, mut ws) = (l0, d0, ws0);
+
+ // Carries: the classes of the byte just before the batch, which is the class of the char
+ // containing it.
+ let (pl, pd, pws, ps, po) = if scan == 0 {
+ (0u64, 0u64, 0u64, 0u64, 0u64)
+ } else {
+ match tags[char_lead(tags, scan - 1)] & 0x0F {
+ LET => (1, 0, 0, 0, 0),
+ NW | NO => (0, 1, 0, 0, 0),
+ SPC => (0, 0, 1, 1, 0),
+ NLN | WSO => (0, 0, 1, 0, 0),
+ _ => (0, 0, 0, 0, 1),
+ }
+ };
+
+ if c != 0 {
+ // A char straddling into the batch has its leading continuation bytes take the carry
+ // class ("other" needs no action: it is derived as the complement).
+ if c & 1 != 0 {
+ let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1);
+ l |= lead_in * pl;
+ d |= lead_in * pd;
+ ws |= lead_in * pws;
+ }
+ let (c2, c3) = cont_runs(c);
+ l = fill(l, c, c2, c3);
+ d = fill(d, c, c2, c3);
+ ws = fill(ws, c, c2, c3);
+ if ws & c != 0 {
+ // Multi-byte whitespace: the `\s+(?!\S)` give-back is one char, not one byte, and
+ // the algebra below works in bytes. Rare (NBSP and friends); scalar batch.
+ return (0, u64::MAX);
+ }
+ }
+ let o = !(l | d | ws);
+
+ // The r50k boundary algebra (gigatoken, MIT). A byte starts a token when it is not
+ // whitespace, does not continue a same-class run, and does not follow a space (the ` ?`
+ // prefix glues it to the space instead).
+ let cont_same = (l & ((l << 1) | pl)) | (d & ((d << 1) | pd)) | (o & ((o << 1) | po));
+ let after_sp = (s << 1) | ps;
+ let nb = !ws & !cont_same & !after_sp;
+
+ let mut bad = 0u64;
+
+ // Whitespace-run splits. `split_ok` = the last whitespace byte before a non-whitespace (the
+ // `\s+(?!\S)` give-back starts a token there); bit 63 needs the lookahead tag.
+ let mut split_ok = ws & (!ws >> 1);
+ let la = tags[scan + 64] & 0x0F;
+ if la == CONT {
+ // The char at the batch edge straddles out. Whitespace is the only class whose rules
+ // look at char ends, so only a whitespace lead poisons the tail.
+ let p = char_lead(tags, scan + 63);
+ if in_mask(tags[p], mask::WS) {
+ bad |= u64::MAX << (p - scan);
+ }
+ } else if !in_mask(la, mask::WS) {
+ split_ok |= ws & (1u64 << 63);
+ }
+ let pwsb = (ws << 1) | pws;
+ let wsboundary = ws & (!pwsb | split_ok);
+ let mut boundary = nb | wsboundary;
+
+ // Contraction fixup, only when the batch has an apostrophe that starts a token: a match
+ // (case-sensitive, as in the scalar dispatch) absorbs the next 1-2 letters and re-opens a
+ // token after them. Too close to the edge (i >= 61) the peek and the re-opened bit can
+ // cross the batch; refuse the tail instead.
+ if ap != 0 {
+ let mut cand = ap & boundary;
+ while cand != 0 {
+ let i = cand.trailing_zeros() as usize;
+ cand &= cand - 1;
+ if i >= 61 {
+ bad |= u64::MAX << i;
+ break;
+ }
+ let k = match text[scan + i + 1] {
+ b's' | b't' | b'm' | b'd' => 2,
+ b'r' | b'v' if text[scan + i + 2] == b'e' => 3,
+ b'l' if text[scan + i + 2] == b'l' => 3,
+ _ => 0,
+ };
+ if k != 0 {
+ boundary &= !(1u64 << (i + 1));
+ boundary |= 1u64 << (i + k);
+ }
+ }
+ }
+ (boundary & !bad, bad)
+}
diff --git a/tokenizers/atomsplit/src/fsm/masked/cl100k.rs b/tokenizers/atomsplit/src/fsm/masked/cl100k.rs
new file mode 100644
index 000000000..438835073
--- /dev/null
+++ b/tokenizers/atomsplit/src/fsm/masked/cl100k.rs
@@ -0,0 +1,313 @@
+//! Masked scheme for the cl100k regex family (cl100k / Llama-3 / GLM at digit cap 3, Qwen2 at
+//! cap 1, unbounded `\p{N}+` at `usize::MAX`) — the boundary algebra via gigatoken's
+//! `cl100k_family.rs` (MIT) over tag-fed class masks.
+//!
+//! Boundary rules (the regex is `'(?i:contractions)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,cap}|
+//! ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+`):
+//! - A letter starts a token unless it continues a letter run, follows a space or a non-newline
+//! whitespace char (those always sit at a boundary before a non-ws char and absorb one letter
+//! run via the `[^\r\n\p{L}\p{N}]?` prefix), or follows a punct char that is itself at a
+//! boundary, i.e. whose own predecessor is neither punct nor space (a two-chars-back test,
+//! made char-aware by shifting per the previous char's byte length).
+//! - Digits split every `cap` chars from each run start and never absorb a preceding space.
+//! - A punct char starts a token unless it continues a punct run or follows a space.
+//! - Newlines directly after a punct run are absorbed (`[\r\n]*`).
+//! - A whitespace run containing newlines emits one token through its LAST newline
+//! (`\s*[\r\n]`), then the give-back rules; NL-free runs split before their last char when
+//! followed by non-ws (`\s+(?!\S)`).
+//!
+//! Bad zones: multi-byte `\p{N}` chars under cap 3 (the grouping is char-counted, byte hops
+//! would misphase it) and any digit run whose phase starts before the batch; whitespace runs
+//! touching the batch end while the next char is whitespace (their last newline may lie
+//! beyond); apostrophes near the batch edge or before a non-ASCII char (`(?i:'s)` also matches
+//! `'ſ`); `Sentinel`/`MultiByte` tags.
+
+use super::super::cl100k::advance_cl100k_cap;
+use super::block::Block;
+use super::{MaskedFsm, char_lead, cont_runs, digit_run_splits3, fill, shl_sat, smear_up};
+use crate::fsm::{APO, CONT, LET, NLN, NO, NW, SPC, WSO, in_mask, mask};
+
+pub(super) struct Cl100kMasked {
+ /// 1, 3 or `usize::MAX`; the entry point routes any other cap to the scalar scan.
+ pub(super) digit_cap: usize,
+}
+
+impl MaskedFsm for Cl100kMasked {
+ #[inline(always)]
+ fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) {
+ batch_masks(text, tags, scan, self.digit_cap)
+ }
+
+ #[inline(always)]
+ fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
+ advance_cl100k_cap(text, tags, i, end, self.digit_cap)
+ }
+}
+
+/// Boundary carries from the two chars before the batch: P1 is the char containing byte
+/// `scan - 1`, P2 the one before it (the two-chars-back absorb test).
+#[derive(Default)]
+struct Carries {
+ pl: u64,
+ ps: u64,
+ pwt: u64,
+ po: u64,
+ pws: u64,
+ pd: u64,
+ /// P2 is punct-or-space, for a char lead at bit 0 (P1 entirely before the batch).
+ c2_os: u64,
+ /// The same test positioned at the first lead after a P1 that straddles into the batch
+ /// (P1's own predecessor is then P2).
+ b2b_in: u64,
+}
+
+fn carries(tags: &[u8], scan: usize, c: u64) -> Carries {
+ let mut cr = Carries::default();
+ if scan == 0 {
+ return cr;
+ }
+ let p1 = char_lead(tags, scan - 1);
+ let c2v = if p1 == 0 {
+ 0
+ } else {
+ let t2 = tags[char_lead(tags, p1 - 1)] & 0x0F;
+ u64::from(t2 == SPC || in_mask(t2, mask::NOT_WS_L_N))
+ };
+ if c & 1 != 0 {
+ cr.b2b_in = c2v << (!c).trailing_zeros();
+ } else {
+ cr.c2_os = c2v;
+ }
+ match tags[p1] & 0x0F {
+ LET => cr.pl = 1,
+ NW | NO => cr.pd = 1,
+ SPC => {
+ cr.pws = 1;
+ cr.ps = 1;
+ }
+ WSO => {
+ cr.pws = 1;
+ cr.pwt = 1;
+ }
+ // A newline is whitespace but never a `[^\r\n\p{L}\p{N}]?` prefix, so pwt stays 0.
+ NLN => cr.pws = 1,
+ _ => cr.po = 1,
+ }
+ cr
+}
+
+fn batch_masks(text: &[u8], tags: &[u8], scan: usize, digit_cap: usize) -> (u64, u64) {
+ debug_assert!(scan + 64 < tags.len() && tags.len() == text.len());
+ // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract.
+ let blk = unsafe { Block::load(tags, scan) };
+ if blk.any_range_tag(13, 1) {
+ // Sentinel / MultiByte: the scalar dispatch has a defensive arm for these.
+ return (0, u64::MAX);
+ }
+ let l0 = blk.eq_tag(LET);
+ let d0 = blk.range_tag(NW, 1);
+ let nl = blk.eq_tag(NLN);
+ let s = blk.eq_tag(SPC);
+ let wt0 = blk.eq_tag(WSO);
+ let ap = if blk.any_eq_full(APO) {
+ blk.eq_full(APO)
+ } else {
+ 0
+ };
+ let c = if blk.any_eq_full(CONT) {
+ blk.eq_full(CONT)
+ } else {
+ 0
+ };
+
+ let cr = carries(tags, scan, c);
+
+ // Fill: continuation bytes join their char's class; a char straddling into the batch has
+ // its leading continuation bytes take P1's class (punct needs no action: it is derived as
+ // the complement).
+ let (mut l, mut d, mut wt) = (l0, d0, wt0);
+ let (c2, c3) = cont_runs(c);
+ if c != 0 {
+ if c & 1 != 0 {
+ let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1);
+ l |= lead_in * cr.pl;
+ d |= lead_in * cr.pd;
+ wt |= lead_in * cr.pwt;
+ }
+ l = fill(l, c, c2, c3);
+ d = fill(d, c, c2, c3);
+ wt = fill(wt, c, c2, c3);
+ }
+ // Per-length char leads (`\s` chars are at most 3 bytes; letters up to 4).
+ let lead = !c;
+ let len1 = lead & !(c >> 1);
+ let len2 = lead & (c >> 1) & !(c >> 2);
+ let len3 = lead & (c >> 1) & (c >> 2) & !(c >> 3);
+ let len4 = lead & (c >> 1) & (c >> 2) & (c >> 3);
+ let w2 = wt0 & len2;
+ let w3 = wt0 & len3;
+
+ let ws_f = s | nl | wt;
+ let o = !(l | d | ws_f);
+
+ // --- Letters: `[^\r\n\p{L}\p{N}]?\p{L}+` --------------------------------------------------
+ // b2back: "the char two back is punct or space", evaluated at each char's lead by shifting
+ // the prev-byte test by the PREVIOUS char's byte length.
+ let c_test = ((o | s) << 1) | cr.po | cr.ps;
+ let b2back = ((c_test & len1) << 1)
+ | ((c_test & len2) << 2)
+ | ((c_test & len3) << 3)
+ | ((c_test & len4) << 4)
+ | cr.c2_os
+ | cr.b2b_in;
+ let p_l = (l << 1) | cr.pl;
+ let p_s = (s << 1) | cr.ps;
+ let p_wt = (wt << 1) | cr.pwt;
+ let p_o = (o << 1) | cr.po;
+ let absorb = p_o & !b2back;
+ let b_letters = l0 & !p_l & !p_s & !p_wt & !absorb;
+
+ // --- Digits: `\p{N}{1,cap}` ----------------------------------------------------------------
+ // Only 1-byte digit chars can be split by byte hops; multi-byte `\p{N}` runs go to the
+ // scalar path under cap 3 (bad below). Cap 1 tokens are single chars (any width), and the
+ // unbounded cap only needs run starts.
+ let d_ascii = d0 & len1;
+ let dmb = (d & c) | (d0 & !len1);
+ let b_digits = match digit_cap {
+ 3 => {
+ if d_ascii & (d_ascii >> 1) != 0 {
+ digit_run_splits3(d_ascii)
+ } else {
+ d_ascii
+ }
+ }
+ 1 => d0,
+ _ => d0 & !((d << 1) | cr.pd),
+ };
+
+ // --- Punct: ` ?[^\s\p{L}\p{N}]+[\r\n]*` ----------------------------------------------------
+ let b_punct = (o & lead) & !p_o & !p_s;
+
+ // Newlines directly after a punct run are absorbed (`[\r\n]*`).
+ let abs_seed = nl & ((o << 1) | cr.po);
+ let abs_n = if abs_seed == 0 {
+ 0
+ } else {
+ smear_up(abs_seed, nl)
+ };
+ let ws_eff = ws_f & !abs_n;
+
+ let mut bad = if digit_cap == 3 {
+ dmb | dmb << 1 | dmb >> 1
+ } else {
+ 0
+ };
+
+ // Byte-64 lookahead: is the char at the next batch's first byte non-ws? Decides whether
+ // ws-like runs touching bit 63 resolve in-batch. A CONT lookahead means a char straddles
+ // out; treating that as "ws" is safe: a live give-back at bit 63 would need the next char's
+ // lead at byte 64, which contradicts the straddle.
+ let la = tags[scan + 64] & 0x0F;
+ let nn64 = la != CONT && !in_mask(la, mask::WS);
+ let nn64m = u64::from(nn64).wrapping_neg();
+
+ // An absorbed newline touching the batch end: if byte 64 is ws the token may continue with
+ // another newline, and the next batch cannot tell an absorbed `\n` before its bit 0 from a
+ // ws-run `\n` — defer.
+ if abs_n >> 63 != 0 && !nn64 {
+ bad |= 1u64 << 63;
+ }
+
+ // A ws run touching the batch end resolves in-batch only when byte 64's char is non-ws
+ // (its last newline and `(?!\S)` split are then all visible); otherwise defer it.
+ let nonws = !ws_eff;
+ if ws_eff >> 63 != 0 && !nn64 {
+ if nonws == 0 {
+ return (0, u64::MAX); // whole batch one ws run
+ }
+ let h = 63 - nonws.leading_zeros();
+ bad |= u64::MAX << (h + 1);
+ }
+
+ // A digit run whose grouping phase did not start inside this batch (a continuation from
+ // before it, or following a bad zone that may hold digit chars) defers too.
+ if digit_cap == 3 {
+ let seed = (d_ascii & (bad << 1)) | (d_ascii & cr.pd);
+ if seed != 0 {
+ bad |= smear_up(seed, d_ascii);
+ }
+ }
+
+ // --- Whitespace ---------------------------------------------------------------------------
+ // Base rule (correct for NL-free runs; NL runs are overridden below): run start, or split
+ // before the last char when followed by non-ws.
+ let ws_leads1 = (s | nl | (wt0 & len1)) & ws_eff;
+ let ws_leads = (ws_leads1 | w2 | w3) & !abs_n;
+ let p_ws = (ws_eff << 1) | cr.pws;
+ let edge_last = (ws_leads1 & (1 << 63)) | (w2 & (1 << 62)) | (w3 & (1 << 61));
+ let split_ok = (ws_leads1 & (nonws >> 1))
+ | (w2 & (nonws >> 2))
+ | (w3 & (nonws >> 3))
+ | (edge_last & nn64m);
+ let mut b_ws = ws_leads & (!p_ws | split_ok);
+
+ // Override every run containing a (non-absorbed) newline: one token through the run's last
+ // newline, then the give-back rules on the remainder.
+ let mut runs_n = nl & ws_eff & !bad;
+ while runs_n != 0 {
+ let f = runs_n.trailing_zeros();
+ let below_gap = nonws & ((1u64 << f) - 1);
+ let a = if below_gap == 0 {
+ 0
+ } else {
+ 64 - below_gap.leading_zeros()
+ };
+ let e = (nonws & (u64::MAX << f)).trailing_zeros();
+ let run_mask = (u64::MAX << a) & !shl_sat(u64::MAX, e);
+ b_ws &= !run_mask;
+ b_ws |= 1u64 << a;
+ let q = 63 - (nl & run_mask).leading_zeros(); // last newline in the run
+ if q + 1 < e {
+ // Tail after the last newline: starts a token, and its last char splits off before
+ // the following non-ws char.
+ b_ws |= 1u64 << (q + 1);
+ let tail = run_mask & (u64::MAX << (q + 1));
+ let tail_leads = ws_leads & tail;
+ b_ws |= 1u64 << (63 - tail_leads.leading_zeros());
+ }
+ runs_n &= !run_mask;
+ }
+
+ let mut boundary = b_letters | b_digits | b_punct | b_ws;
+
+ // --- Contractions: `'(?i:[sdmt]|ll|ve|re)` -------------------------------------------------
+ let mut cand = ap & boundary & !bad;
+ while cand != 0 {
+ let i = cand.trailing_zeros() as usize;
+ cand &= cand - 1;
+ if i >= 61 {
+ bad |= u64::MAX << i;
+ break;
+ }
+ let b1 = text[scan + i + 1];
+ if b1 >= 0x80 {
+ // `(?i:'s)` also matches 'ſ (U+017F): an apostrophe before any non-ASCII char
+ // defers to the scalar path.
+ bad |= 0b111u64 << i;
+ continue;
+ }
+ let k = match b1 | 0x20 {
+ b's' | b'd' | b'm' | b't' => 2,
+ b'l' if text[scan + i + 2] | 0x20 == b'l' => 3,
+ b'v' if text[scan + i + 2] | 0x20 == b'e' => 3,
+ b'r' if text[scan + i + 2] | 0x20 == b'e' => 3,
+ _ => 0,
+ };
+ if k != 0 {
+ boundary &= !(1u64 << (i + 1));
+ boundary |= 1u64 << (i + k);
+ }
+ }
+
+ (boundary & !bad, bad)
+}
diff --git a/tokenizers/atomsplit/src/fsm/masked/deepseek.rs b/tokenizers/atomsplit/src/fsm/masked/deepseek.rs
new file mode 100644
index 000000000..4598d4ae7
--- /dev/null
+++ b/tokenizers/atomsplit/src/fsm/masked/deepseek.rs
@@ -0,0 +1,306 @@
+//! Masked scheme for the deepseek-v3 Sequence (digits `\p{N}{1,3}` → CJK-range runs → the big
+//! regex) — boundary rules derived from [`scan_deepseek`]'s dispatch, following gigatoken's
+//! `deepseek_v3.rs` scoping (MIT).
+//!
+//! Deepseek has no case rules and no contractions, but three shapes of its own:
+//! - The alt-2 prefix class is `[^\r\n\p{L}\p{P}\p{S}]?`: letters absorb a preceding space,
+//! non-newline whitespace char, or the LAST char of a gap run (Control / NumericOther /
+//! ZWJ — chars matching no alternative), and never a punct char or a digit.
+//! - alt-1 `[ascii-punct][A-Za-z]+`: an ASCII punct char at a token start absorbs a following
+//! ASCII-letter run. Where that run collides with a non-ASCII letter or mark, the two rules
+//! diverge (`[A-Za-z]+` stops, `[\p{L}\p{M}]+` would continue); those collision bits defer.
+//! - A whitespace run followed by a digit or CJK char is one whole token (Split-1/2 isolated
+//! the follower), so the `\s+(?!\S)` give-back is gated on the follower's class.
+//!
+//! CJK-range chars are closed units re-split into same-kind sub-runs, and every other rule
+//! stops at them; a batch containing any byte with a lead in `0xE3..=0xE9` (a superset of the
+//! CJK ranges) defers whole, before the tag masks are built. On CJK-dominated text the scanner
+//! is then the scalar scan plus one 64-byte test per batch, which is the accepted trade: the
+//! win is on latin/code text, and CJK batches resolve through the same scalar rules as before.
+
+use super::super::deepseek::advance_deepseek;
+use super::block::Block;
+use super::{MaskedFsm, char_lead, cont_runs, digit_run_splits3, fill, shl_sat, smear_up};
+use crate::fsm::{ASM, CON, CONT, CTL, LET, MRK, NLN, NMO, NO, NW, SPC, WSO, ZWJ, in_mask, mask};
+
+pub(super) struct DeepSeekMasked;
+
+impl MaskedFsm for DeepSeekMasked {
+ #[inline(always)]
+ fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) {
+ batch_masks(text, tags, scan)
+ }
+
+ #[inline(always)]
+ fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
+ advance_deepseek(text, tags, i, end)
+ }
+}
+
+/// A deepseek letter-run member tag: `[\p{L}\p{M}]` minus the refined Marks (ASM/ZWJ). The
+/// CJK-range exclusion is handled by the bad cover, not here.
+#[inline(always)]
+fn is_member(t: u8) -> bool {
+ in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ
+}
+
+#[derive(Default)]
+struct Carries {
+ pl: u64,
+ ps: u64,
+ pwt: u64,
+ po: u64,
+ pws: u64,
+ pd: u64,
+ pgap: u64,
+ /// Byte `scan - 1` is an ASCII punct char at a token start: an alt-1 absorb may reach into
+ /// this batch.
+ alt1: u64,
+ /// P1 is a CJK-range char. Its all-zero carries say "closed unit", which is right for the
+ /// char AFTER it; bytes of the char itself straddling into the batch stay unclassified and
+ /// must defer (see `batch_masks`).
+ cjk: bool,
+}
+
+fn carries(text: &[u8], tags: &[u8], scan: usize) -> Carries {
+ let mut cr = Carries::default();
+ if scan == 0 {
+ return cr;
+ }
+ let p1 = char_lead(tags, scan - 1);
+ if (0xE3..=0xE9).contains(&text[p1]) {
+ // A CJK-range char is a closed unit: everything after it starts fresh, which is what
+ // all-zero carries say.
+ cr.cjk = true;
+ return cr;
+ }
+ let t1 = tags[p1];
+ match t1 & 0x0F {
+ LET | MRK if is_member(t1) => cr.pl = 1,
+ NW | NO => cr.pd = 1,
+ SPC => {
+ cr.pws = 1;
+ cr.ps = 1;
+ }
+ WSO => {
+ cr.pws = 1;
+ cr.pwt = 1;
+ }
+ NLN => cr.pws = 1,
+ NMO | CTL => cr.pgap = 1,
+ _ if t1 == ZWJ => cr.pgap = 1,
+ _ => cr.po = 1,
+ }
+ if text[scan - 1].is_ascii_punctuation() {
+ // Was P1 (one byte) at a token start? Not when it continues a punct run or follows a
+ // space; a CJK-range P2 is a closed unit, so P1 starts fresh after it.
+ let at_start = if p1 == 0 {
+ true
+ } else {
+ let p2 = char_lead(tags, p1 - 1);
+ let t2 = tags[p2];
+ (0xE3..=0xE9).contains(&text[p2])
+ || !(t2 & 0x0F == SPC || in_mask(t2, mask::PUNCT_SYM) || t2 == ASM)
+ };
+ cr.alt1 = u64::from(at_start);
+ }
+ cr
+}
+
+fn batch_masks(text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) {
+ debug_assert!(scan + 64 < tags.len() && tags.len() == text.len());
+ // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the blocks' load contract.
+ let tb = unsafe { Block::load(text, scan) };
+ // CJK-range chars (lead 0xE3..=0xE9, a superset of the Split-2 ranges) are closed units
+ // that every other rule stops at; a batch containing any defers whole, BEFORE the tag
+ // masks are built. On CJK-dominated text the scanner degenerates to the scalar scan plus
+ // this one test (see the module doc).
+ if tb.range_full(0xE3, 6) != 0 {
+ return (0, u64::MAX);
+ }
+ // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract.
+ let blk = unsafe { Block::load(tags, scan) };
+ if blk.any_range_tag(13, 1) {
+ return (0, u64::MAX);
+ }
+ let l0 = blk.eq_tag(LET);
+ let mk0 = blk.eq_full(MRK); // true marks join deepseek letter runs
+ let o0 = blk.range_tag(CON, 3) | blk.eq_full(ASM); // `[\p{P}\p{S}]` = Con|Pun|Apo|Sym (+ASM)
+ let gap0 = blk.range_tag(NMO, 1) | blk.eq_full(ZWJ);
+ let d0 = blk.range_tag(NW, 1);
+ let nl = blk.eq_tag(NLN);
+ let s = blk.eq_tag(SPC);
+ let wt0 = blk.eq_tag(WSO);
+ let c = if blk.any_eq_full(CONT) {
+ blk.eq_full(CONT)
+ } else {
+ 0
+ };
+ let pa = tb.ascii_punct();
+ let alpha = tb.ascii_alpha();
+
+ let cr = carries(text, tags, scan);
+
+ let lml0 = l0 | mk0;
+ let (mut lm, mut d, mut wt, mut g) = (lml0, d0, wt0, gap0);
+ let (c2, c3) = cont_runs(c);
+ if c != 0 {
+ if c & 1 != 0 {
+ let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1);
+ lm |= lead_in * cr.pl;
+ d |= lead_in * cr.pd;
+ wt |= lead_in * cr.pwt;
+ g |= lead_in * cr.pgap;
+ }
+ lm = fill(lm, c, c2, c3);
+ d = fill(d, c, c2, c3);
+ wt = fill(wt, c, c2, c3);
+ g = fill(g, c, c2, c3);
+ }
+ let lead = !c;
+ let len1 = lead & !(c >> 1);
+ let len2 = lead & (c >> 1) & !(c >> 2);
+ let len3 = lead & (c >> 1) & (c >> 2) & !(c >> 3);
+ let len4 = lead & (c >> 1) & (c >> 2) & (c >> 3);
+ let w2 = wt0 & len2;
+ let w3 = wt0 & len3;
+
+ let ws_f = s | nl | wt;
+ // The complement is the punct-run class here too (gap and CJK bytes land in it; both are
+ // covered by their own masks or the bad cover below).
+ let o = !(lm | d | ws_f | g);
+
+ let mut bad = 0u64;
+ if cr.cjk && c & 1 != 0 {
+ // A CJK char straddling into the batch: its continuation bytes carry no class (the
+ // all-zero carries only cover the char after it), so they and the char right after
+ // them defer.
+ let e1 = (!c).trailing_zeros();
+ bad |= (c & ((1u64 << e1) - 1)) | (1u64 << e1);
+ }
+
+ // --- Letters: `[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+` and alt-1 `[ascii-punct][A-Za-z]+` ---
+ let p_lm = (lm << 1) | cr.pl;
+ let p_s = (s << 1) | cr.ps;
+ let p_wt = (wt << 1) | cr.pwt;
+ let p_g = (g << 1) | cr.pgap;
+ let p_o = (o << 1) | cr.po;
+
+ // --- Punct: ` ?[\p{P}\p{S}]+[\r\n]*` -------------------------------------------------------
+ let b_punct = o0 & !p_o & !p_s;
+ let abs_seed = nl & ((o << 1) | cr.po);
+ let abs_n = if abs_seed == 0 {
+ 0
+ } else {
+ smear_up(abs_seed, nl)
+ };
+ let ws_eff = ws_f & !abs_n;
+
+ // alt-1 absorbs: an ASCII-letter byte right after a token-starting ASCII punct char, and
+ // the rest of that `[A-Za-z]+` run. Where the run's end meets a letter-run member the two
+ // letter rules diverge: defer that bit. A run touching the batch end defers too (the next
+ // batch cannot tell an alt-1 run from a plain letter run).
+ let absorb_a = alpha & (((pa & b_punct) << 1) | cr.alt1);
+ let b_letters = lml0 & !p_lm & !p_s & !p_wt & !p_g & !absorb_a;
+ if absorb_a != 0 {
+ let zone = smear_up(absorb_a, alpha);
+ bad |= (zone << 1) & !alpha & lml0;
+ if zone >> 63 != 0 {
+ bad |= 1u64 << 63;
+ }
+ }
+
+ // --- Gap runs: boundary at the run start, and at the last char before an absorbed letter
+ // run (the prefix split; one and the same bit for a single-char gap). The prefix split
+ // reads the NEXT char's lead, so a gap char whose follower sits past the batch edge
+ // defers.
+ let b_gap = (gap0 & !p_g)
+ | (gap0
+ & ((len1 & (lml0 >> 1))
+ | (len2 & (lml0 >> 2))
+ | (len3 & (lml0 >> 3))
+ | (len4 & (lml0 >> 4))));
+ bad |= (gap0 & len1 & (u64::MAX << 63))
+ | (gap0 & len2 & (u64::MAX << 62))
+ | (gap0 & len3 & (u64::MAX << 61))
+ | (gap0 & len4 & (u64::MAX << 60));
+
+ // --- Digits: `\p{N}{1,3}` (the cl100k cap-3 machinery) -------------------------------------
+ let d_ascii = d0 & len1;
+ let dmb = (d & c) | (d0 & !len1);
+ let b_digits = if d_ascii & (d_ascii >> 1) != 0 {
+ digit_run_splits3(d_ascii)
+ } else {
+ d_ascii
+ };
+ bad |= dmb | dmb << 1 | dmb >> 1;
+
+ // --- Whitespace -----------------------------------------------------------------------------
+ // The give-back is gated on the follower: a digit or CJK char after the run means the whole
+ // run is one token (`iso` covers the follower's bytes; CJK is inside `bad` anyway, but the
+ // gate keeps the algebra honest about why).
+ let la = tags[scan + 64] & 0x0F;
+ let la_iso = matches!(la, NW | NO) || (0xE3..=0xE9).contains(&text[scan + 64]);
+ let nn64 = la != CONT && !in_mask(la, mask::WS);
+ let nn64m = u64::from(nn64 && !la_iso).wrapping_neg();
+ if abs_n >> 63 != 0 && !nn64 {
+ bad |= 1u64 << 63;
+ }
+ let nonws = !ws_eff;
+ if ws_eff >> 63 != 0 && !nn64 {
+ if nonws == 0 {
+ return (0, u64::MAX);
+ }
+ let h = 63 - nonws.leading_zeros();
+ bad |= u64::MAX << (h + 1);
+ }
+ let seed = (d_ascii & (bad << 1)) | (d_ascii & cr.pd);
+ if seed != 0 {
+ bad |= smear_up(seed, d_ascii);
+ }
+
+ // A whitespace run followed by a digit keeps its last char (no give-back); the CJK case
+ // never reaches this path (any in-batch CJK deferred above), leaving only the lookahead.
+ let iso = d;
+ let nonws_ni = nonws & !iso;
+ let p_ws = (ws_eff << 1) | cr.pws;
+ let ws_leads1 = (s | nl | (wt0 & len1)) & ws_eff;
+ let ws_leads = (ws_leads1 | w2 | w3) & !abs_n;
+ let edge_last = (ws_leads1 & (1 << 63)) | (w2 & (1 << 62)) | (w3 & (1 << 61));
+ let split_ok = (ws_leads1 & (nonws_ni >> 1))
+ | (w2 & (nonws_ni >> 2))
+ | (w3 & (nonws_ni >> 3))
+ | (edge_last & nn64m);
+ let mut b_ws = ws_leads & (!p_ws | split_ok);
+
+ let mut runs_n = nl & ws_eff & !bad;
+ while runs_n != 0 {
+ let f = runs_n.trailing_zeros();
+ let below_gap = nonws & ((1u64 << f) - 1);
+ let a = if below_gap == 0 {
+ 0
+ } else {
+ 64 - below_gap.leading_zeros()
+ };
+ let e = (nonws & (u64::MAX << f)).trailing_zeros();
+ let run_mask = (u64::MAX << a) & !shl_sat(u64::MAX, e);
+ b_ws &= !run_mask;
+ b_ws |= 1u64 << a;
+ let q = 63 - (nl & run_mask).leading_zeros();
+ // The post-newline tail: no give-back when the run's follower is a digit or CJK char
+ // (the whole tail is then one token).
+ let follower_iso = if e >= 64 { la_iso } else { iso >> e & 1 != 0 };
+ if q + 1 < e {
+ b_ws |= 1u64 << (q + 1);
+ if !follower_iso {
+ let tail = run_mask & (u64::MAX << (q + 1));
+ let tail_leads = ws_leads & tail;
+ b_ws |= 1u64 << (63 - tail_leads.leading_zeros());
+ }
+ }
+ runs_n &= !run_mask;
+ }
+
+ let boundary = b_letters | b_digits | b_punct | b_gap | b_ws;
+ (boundary & !bad, bad)
+}
diff --git a/tokenizers/atomsplit/src/fsm/masked/o200k.rs b/tokenizers/atomsplit/src/fsm/masked/o200k.rs
new file mode 100644
index 000000000..dbc77c477
--- /dev/null
+++ b/tokenizers/atomsplit/src/fsm/masked/o200k.rs
@@ -0,0 +1,473 @@
+//! Masked scheme for the o200k regex family (o200k / GPT-4o / gpt-oss with contractions and
+//! digit cap 3, Mistral tekken without contractions at cap 1) — the boundary algebra via
+//! gigatoken's `o200k_family.rs` (MIT) over tag-fed class masks.
+//!
+//! Differences from the cl100k family:
+//! - Letter runs are case-structured. Under leftmost-greedy backtracking the two letter
+//! alternatives reduce to a phase automaton: a strict-upper char (`\p{Lu}\p{Lt}`) ends a
+//! token exactly when the previous char is strict-lower (`\p{Ll}`) — "camelCase" splits
+//! `camel|Case`, "HTTPResponse" stays one token. A strict-upper after a CASELESS letter
+//! needs the phase and lookahead (the backtrack to the last caseless char), so those chars
+//! defer to the scalar path.
+//! - Contractions are attached suffixes of letter tokens, not a standalone alternative:
+//! "don't" is ONE token and the char after a consumed suffix always starts a new one
+//! ("can'ts" is `can't|s`). A contraction applies only when the apostrophe directly follows
+//! a letter-run char; elsewhere `'` is ordinary punctuation.
+//! - Punct runs absorb a `[\r\n/]*` tail. `/` is itself punct, so an absorbed tail always
+//! begins with a newline; whether a batch-leading `[\r\n/]` run continues such a tail is
+//! resolved by a bounded walkback over the preceding text.
+//! - Marks (`\p{M}`) are dual-class: they join letter runs AND continue punct runs, so their
+//! effective class is run-contextual. Mark chars (rare) defer to the scalar path with a bad
+//! smear wide enough (8 bytes forward) to cover every boundary their class can influence
+//! (two chars of multi-byte followers).
+
+use super::super::o200k::advance_o200k;
+use super::block::Block;
+use super::{MaskedFsm, char_lead, cont_runs, digit_run_splits3, fill, shl_sat, smear_up};
+use crate::fsm::{APO, ASM, Atom, CONT, LET, MRK, NLN, NO, NW, SPC, WSO, ZWJ, in_mask, mask};
+
+pub(super) struct O200kMasked;
+
+impl MaskedFsm
+ for O200kMasked
+{
+ #[inline(always)]
+ fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) {
+ batch_masks::(text, tags, scan)
+ }
+
+ #[inline(always)]
+ fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize {
+ advance_o200k::(text, tags, i, end)
+ }
+}
+
+#[inline(always)]
+fn is_tail_byte(b: u8) -> bool {
+ matches!(b, b'\r' | b'\n' | b'/')
+}
+
+#[inline(always)]
+fn is_member_mark(t: u8) -> bool {
+ t & 0x0F == MRK && t != ASM && t != ZWJ
+}
+
+/// Was the tail-class byte at `scan - 1` absorbed by a punct run's `[\r\n/]*` tail (as opposed
+/// to being a fresh punct-run `/` or a ws-run newline)? Walks the tail-class run back (bounded)
+/// and classifies the char before it. `None`: unresolved (over-long run, or a preceding mark
+/// whose own class is run-contextual).
+fn prev_tail_absorbed(text: &[u8], tags: &[u8], scan: usize) -> Option {
+ debug_assert!(scan >= 1 && is_tail_byte(text[scan - 1]));
+ let mut r = scan - 1;
+ let mut steps = 0;
+ while r > 0 && is_tail_byte(text[r - 1]) {
+ r -= 1;
+ steps += 1;
+ if steps > 8 {
+ return None;
+ }
+ }
+ // T-run = text[r..scan]. The `[\r\n/]*` tail is greedy, so once absorption triggers — at
+ // the first newline that directly follows a punct-run char (an in-run slash, or the
+ // pre-run char for a run-leading newline) — everything to the run's end is absorbed.
+ // Before the trigger, newlines are ws-run members and slashes ordinary punct-run bytes.
+ let run = &text[r..scan];
+ let mut trigger = usize::MAX;
+ let mut seen_slash = false;
+ for (j, &b) in run.iter().enumerate() {
+ if b == b'/' {
+ seen_slash = true;
+ continue;
+ }
+ if seen_slash {
+ trigger = j;
+ break;
+ }
+ if j == 0 {
+ if r == 0 {
+ continue;
+ }
+ let t = tags[char_lead(tags, r - 1)];
+ if is_member_mark(t) || t & 0x0F >= 13 {
+ // A mark continues whatever run precedes it; Sentinel/MultiByte are opaque.
+ return None;
+ }
+ if !in_mask(t, mask::LETTER | mask::NUMBER | mask::WS) {
+ trigger = 0;
+ break;
+ }
+ }
+ }
+ Some(scan - 1 - r >= trigger)
+}
+
+/// Two-back "punct or space" test for the char whose lead is `p2`. A slash may be an absorbed
+/// tail byte (a token end, neither punct-run member nor space), so it resolves through the
+/// walkback. `None`: unresolved (the caller sets `force_bad_lead`). A mark P2 answers 0: the
+/// bits that read it wrongly sit inside the previous batch's mark smear, and the scalar
+/// overrun from there covers them (the walker's resume masking).
+fn c2_os_at(text: &[u8], tags: &[u8], p2: usize) -> Option {
+ if text[p2] == b'/' {
+ return prev_tail_absorbed(text, tags, p2 + 1).map(|abs| u64::from(!abs));
+ }
+ let t2 = tags[p2];
+ Some(u64::from(
+ t2 & 0x0F == SPC || (!is_member_mark(t2) && in_mask(t2, mask::NOT_WS_L_N)),
+ ))
+}
+
+/// Boundary carries from the two chars before the batch (the cl100k set, plus the case classes
+/// and the absorbed-tail resolution).
+#[derive(Default)]
+struct Carries {
+ pl: u64,
+ pu: u64,
+ pcl: u64,
+ ps: u64,
+ pwt: u64,
+ po: u64,
+ pws: u64,
+ pd: u64,
+ /// P1 is a member-mark (seeds the mark smear at bit 0).
+ pmk: u64,
+ c2_os: u64,
+ b2b_in: u64,
+ /// P1 is an absorbed `[\r\n/]*` tail byte whose token may continue into this batch.
+ p_abs: bool,
+ /// The tail walkback could not resolve: the batch's leading tail-class run (plus the byte
+ /// after it) can't be trusted.
+ force_bad_lead: bool,
+}
+
+fn carries(text: &[u8], tags: &[u8], scan: usize, c: u64) -> Carries {
+ let mut cr = Carries::default();
+ if scan == 0 {
+ return cr;
+ }
+ if is_tail_byte(text[scan - 1]) {
+ // An absorbed tail ended the previous token, so every "P1 is X" carry is zero and only
+ // the tail-continuation seed survives. A fresh `/` is an ordinary punct byte; fresh
+ // `\r\n` are ws-run newlines.
+ match prev_tail_absorbed(text, tags, scan) {
+ None => cr.force_bad_lead = true,
+ Some(true) => cr.p_abs = true,
+ Some(false) => {
+ if text[scan - 1] == b'/' {
+ cr.po = 1;
+ } else {
+ cr.pws = 1;
+ }
+ match c2_os_at(text, tags, char_lead(tags, scan - 2)) {
+ Some(v) => cr.c2_os = v,
+ None => cr.force_bad_lead = true,
+ }
+ }
+ }
+ return cr;
+ }
+ let p1 = char_lead(tags, scan - 1);
+ let c2v = if p1 == 0 {
+ Some(0)
+ } else {
+ c2_os_at(text, tags, char_lead(tags, p1 - 1))
+ };
+ match c2v {
+ Some(v) => {
+ if c & 1 != 0 {
+ cr.b2b_in = v << (!c).trailing_zeros();
+ } else {
+ cr.c2_os = v;
+ }
+ }
+ None => cr.force_bad_lead = true,
+ }
+ let t1 = tags[p1];
+ match t1 & 0x0F {
+ LET => {
+ cr.pl = 1;
+ cr.pu = u64::from(t1 == Atom::UpperLetter as u8);
+ cr.pcl = u64::from(t1 == Atom::Letter as u8);
+ }
+ NW | NO => cr.pd = 1,
+ SPC => {
+ cr.pws = 1;
+ cr.ps = 1;
+ }
+ WSO => {
+ cr.pws = 1;
+ cr.pwt = 1;
+ }
+ MRK if is_member_mark(t1) => cr.pmk = 1,
+ _ => cr.po = 1,
+ }
+ cr
+}
+
+fn batch_masks(
+ text: &[u8],
+ tags: &[u8],
+ scan: usize,
+) -> (u64, u64) {
+ debug_assert!(scan + 64 < tags.len() && tags.len() == text.len());
+ // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract.
+ let blk = unsafe { Block::load(tags, scan) };
+ if blk.any_range_tag(13, 1) {
+ return (0, u64::MAX);
+ }
+ let l0 = blk.eq_tag(LET);
+ let ub0 = blk.eq_full(Atom::UpperLetter as u8);
+ let cl0 = blk.eq_full(Atom::Letter as u8);
+ let mk0 = blk.eq_full(MRK); // true marks: refined Mark tags (ASM/ZWJ) are punct-class
+ let d0 = blk.range_tag(NW, 1);
+ let nl = blk.eq_tag(NLN);
+ let s = blk.eq_tag(SPC);
+ let wt0 = blk.eq_tag(WSO);
+ let ap = if blk.any_eq_full(APO) {
+ blk.eq_full(APO)
+ } else {
+ 0
+ };
+ let c = if blk.any_eq_full(CONT) {
+ blk.eq_full(CONT)
+ } else {
+ 0
+ };
+
+ let cr = carries(text, tags, scan, c);
+
+ let (mut l, mut u, mut clb, mut d, mut wt, mut mk) = (l0, ub0, cl0, d0, wt0, mk0);
+ let (c2, c3) = cont_runs(c);
+ if c != 0 {
+ if c & 1 != 0 {
+ let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1);
+ l |= lead_in * cr.pl;
+ u |= lead_in * cr.pu;
+ clb |= lead_in * cr.pcl;
+ d |= lead_in * cr.pd;
+ wt |= lead_in * cr.pwt;
+ mk |= lead_in * cr.pmk;
+ }
+ l = fill(l, c, c2, c3);
+ u = fill(u, c, c2, c3);
+ clb = fill(clb, c, c2, c3);
+ d = fill(d, c, c2, c3);
+ wt = fill(wt, c, c2, c3);
+ mk = fill(mk, c, c2, c3);
+ }
+ mk |= cr.pmk; // a mark P1 poisons bit 0 even when it ends exactly at the batch edge
+ let lead = !c;
+ let len1 = lead & !(c >> 1);
+ let len2 = lead & (c >> 1) & !(c >> 2);
+ let len3 = lead & (c >> 1) & (c >> 2) & !(c >> 3);
+ let len4 = lead & (c >> 1) & (c >> 2) & (c >> 3);
+ let w2 = wt0 & len2;
+ let w3 = wt0 & len3;
+
+ let ws_f = s | nl | wt;
+ // Marks land in the complement (dual-class); every bit that can read them is inside the
+ // mark bad smear below, so their punct-class reading is never trusted.
+ let o = !(l | d | ws_f);
+
+ // --- Absorbed `[\r\n/]*` tails ---------------------------------------------------------
+ // The tail class needs the slash mask from the TEXT block; skip that load when the batch
+ // has no newline and no tail context carried in.
+ let (tcls, abs_t) = if nl != 0 || cr.p_abs || cr.force_bad_lead {
+ // SAFETY: `scan + 64 < text.len()` (walker guarantee), the block's load contract.
+ let sl = unsafe { Block::load(text, scan) }.eq_full(b'/');
+ let tcls = nl | sl;
+ let abs_seed = (nl & ((o << 1) | cr.po)) | (u64::from(cr.p_abs) & tcls);
+ let abs_t = if abs_seed == 0 {
+ 0
+ } else {
+ smear_up(abs_seed, tcls)
+ };
+ (tcls, abs_t)
+ } else {
+ (nl, 0)
+ };
+ let ob_eff = o & !abs_t;
+
+ // --- Letters (see the cl100k scheme for the base rules) ---------------------------------
+ let c_test = ((ob_eff | s) << 1) | cr.po | cr.ps;
+ let b2back = ((c_test & len1) << 1)
+ | ((c_test & len2) << 2)
+ | ((c_test & len3) << 3)
+ | ((c_test & len4) << 4)
+ | cr.c2_os
+ | cr.b2b_in;
+ let p_l = (l << 1) | cr.pl;
+ let p_u = (u << 1) | cr.pu;
+ let p_cl = (clb << 1) | cr.pcl;
+ let p_s = (s << 1) | cr.ps;
+ let p_wt = (wt << 1) | cr.pwt;
+ let p_o = (ob_eff << 1) | cr.po;
+ let absorb = p_o & !b2back;
+ // Casing boundary: a strict-upper char after a strict-lower one. (For ASCII text this is
+ // the whole rule; upper-after-caseless defers below.)
+ let p_sl = p_l & !p_u & !p_cl;
+ let b_letters = (l0 & !p_l & !p_s & !p_wt & !absorb) | (ub0 & p_sl);
+
+ // --- Digits ------------------------------------------------------------------------------
+ let d_ascii = d0 & len1;
+ let dmb = (d & c) | (d0 & !len1);
+ let b_digits = if DIGIT_CAP == 3 {
+ if d_ascii & (d_ascii >> 1) != 0 {
+ digit_run_splits3(d_ascii)
+ } else {
+ d_ascii
+ }
+ } else {
+ d0 // cap 1: every digit char its own token
+ };
+
+ // --- Punct: ` ?[^\s\p{L}\p{N}]+[\r\n/]*` --------------------------------------------------
+ let b_punct = (ob_eff & lead) & !p_o & !p_s;
+
+ // --- Bad zones ----------------------------------------------------------------------------
+ let mut bad = if DIGIT_CAP == 3 {
+ dmb | dmb << 1 | dmb >> 1
+ } else {
+ 0
+ };
+ if mk != 0 {
+ // A mark's run-contextual class can affect boundaries up to two chars after it (8
+ // bytes of multi-byte followers) and the byte before.
+ bad |= mk
+ | (mk << 1)
+ | (mk << 2)
+ | (mk << 3)
+ | (mk << 4)
+ | (mk << 5)
+ | (mk << 6)
+ | (mk << 7)
+ | (mk << 8)
+ | (mk >> 1);
+ }
+ // A strict-upper char after a caseless letter: phase- and lookahead-dependent.
+ bad |= ub0 & ((clb << 1) | cr.pcl);
+ if cr.force_bad_lead {
+ bad |= (smear_up(tcls & 1, tcls) << 1) | 0b11;
+ }
+
+ // --- Whitespace ---------------------------------------------------------------------------
+ let ws_eff = ws_f & !abs_t;
+ let la = tags[scan + 64] & 0x0F;
+ let nn64 = la != CONT && !in_mask(la, mask::WS);
+ let nn64m = u64::from(nn64).wrapping_neg();
+
+ // An absorbed tail touching the batch end continues iff byte 64 is tail-class; the next
+ // batch's tail walkback re-derives the context either way, so nothing defers here. A ws
+ // run touching the batch end still defers when byte 64 is ws.
+ let nonws = !ws_eff;
+ if ws_eff >> 63 != 0 && !nn64 {
+ if nonws == 0 {
+ return (0, u64::MAX);
+ }
+ let h = 63 - nonws.leading_zeros();
+ bad |= u64::MAX << (h + 1);
+ }
+ if DIGIT_CAP == 3 {
+ let seed = (d_ascii & (bad << 1)) | (d_ascii & cr.pd);
+ if seed != 0 {
+ bad |= smear_up(seed, d_ascii);
+ }
+ }
+
+ let ws_leads1 = (s | nl | (wt0 & len1)) & ws_eff;
+ let ws_leads = (ws_leads1 | w2 | w3) & !abs_t;
+ let p_ws = (ws_eff << 1) | cr.pws;
+ let edge_last = (ws_leads1 & (1 << 63)) | (w2 & (1 << 62)) | (w3 & (1 << 61));
+ let split_ok = (ws_leads1 & (nonws >> 1))
+ | (w2 & (nonws >> 2))
+ | (w3 & (nonws >> 3))
+ | (edge_last & nn64m);
+ let mut b_ws = ws_leads & (!p_ws | split_ok);
+
+ let mut runs_n = nl & ws_eff & !bad;
+ while runs_n != 0 {
+ let f = runs_n.trailing_zeros();
+ let below_gap = nonws & ((1u64 << f) - 1);
+ let a = if below_gap == 0 {
+ 0
+ } else {
+ 64 - below_gap.leading_zeros()
+ };
+ let e = (nonws & (u64::MAX << f)).trailing_zeros();
+ let run_mask = (u64::MAX << a) & !shl_sat(u64::MAX, e);
+ b_ws &= !run_mask;
+ b_ws |= 1u64 << a;
+ let q = 63 - (nl & run_mask).leading_zeros();
+ if q + 1 < e {
+ b_ws |= 1u64 << (q + 1);
+ let tail = run_mask & (u64::MAX << (q + 1));
+ let tail_leads = ws_leads & tail;
+ b_ws |= 1u64 << (63 - tail_leads.leading_zeros());
+ }
+ runs_n &= !run_mask;
+ }
+
+ let mut boundary = b_letters | b_digits | b_punct | b_ws;
+
+ // --- Contractions: suffix `(?i:'s|'t|'re|'ve|'m|'ll|'d)?` ---------------------------------
+ // An apostrophe at a boundary right after a letter-run char merges the suffix into that
+ // token and forces a boundary right after it.
+ if CONTRACTION {
+ let mut cand = ap & boundary & p_l & !bad;
+ let mut last_forced = usize::MAX;
+ while cand != 0 {
+ let i = cand.trailing_zeros() as usize;
+ cand &= cand - 1;
+ if i <= 2 {
+ // The preceding letter could itself end an earlier contraction that started
+ // before the batch: scalar.
+ bad |= 0b111u64 << i;
+ continue;
+ }
+ if i >= 61 {
+ bad |= u64::MAX << i;
+ break;
+ }
+ if i == last_forced {
+ // "x'll'd": the letter before this apostrophe is a consumed suffix's last
+ // char; a new (prefix) match starts here instead.
+ continue;
+ }
+ // The letter before this apostrophe may itself be a consumed suffix's last char
+ // resolved where `last_forced` can't see it (a scalar-walked zone, or a fixup
+ // before the batch): locally ambiguous, defer.
+ let p = scan + i;
+ let prev_suffix_possible = (text[p - 2] == b'\''
+ && matches!(text[p - 1] | 0x20, b's' | b'd' | b'm' | b't'))
+ || (text[p - 3] == b'\''
+ && (matches!(
+ (text[p - 2] | 0x20, text[p - 1] | 0x20),
+ (b'l', b'l') | (b'v', b'e') | (b'r', b'e')
+ ) || (text[p - 2] == 0xC5 && text[p - 1] == 0xBF)));
+ if prev_suffix_possible {
+ bad |= 0b111u64 << i;
+ continue;
+ }
+ let b1 = text[p + 1];
+ if b1 >= 0x80 {
+ // `(?i:'s)` also matches 'ſ (U+017F): defer.
+ bad |= 0b111u64 << i;
+ continue;
+ }
+ let k = match b1 | 0x20 {
+ b's' | b'd' | b'm' | b't' => 2,
+ b'l' if text[p + 2] | 0x20 == b'l' => 3,
+ b'v' if text[p + 2] | 0x20 == b'e' => 3,
+ b'r' if text[p + 2] | 0x20 == b'e' => 3,
+ _ => 0,
+ };
+ if k != 0 {
+ boundary &= !(1u64 << i);
+ boundary &= !(((1u64 << (k - 1)) - 1) << (i + 1));
+ boundary |= 1u64 << (i + k);
+ last_forced = i + k;
+ }
+ }
+ }
+
+ (boundary & !bad, bad)
+}
diff --git a/tokenizers/atomsplit/src/fsm/o200k.rs b/tokenizers/atomsplit/src/fsm/o200k.rs
index 28e15640a..6914e02c7 100644
--- a/tokenizers/atomsplit/src/fsm/o200k.rs
+++ b/tokenizers/atomsplit/src/fsm/o200k.rs
@@ -50,19 +50,18 @@ fn o200k_letter_match(tags: &[u8], p: usize, re: usize) -> usize {
e
}
-/// Emit the o200k case-split of the letter run `[ls, re)` into `out[*w..]`: the first sub-token starts at
+/// Emit the o200k case-split of the letter run `[ls, re)`: the first sub-token starts at
/// `pfx` (the optional `[^\r\n\p{L}\p{N}]?` prefix; `pfx == ls` when none), the last absorbs a trailing
/// contraction (`CONTRACTION` — off for tekken). Returns the new cursor (past the contraction).
/// `ls < re` (caller-guaranteed).
#[inline(always)]
-fn emit_o200k_letters(
+fn emit_o200k_letters(
text: &[u8],
tags: &[u8],
pfx: usize,
ls: usize,
re: usize,
- out: &mut [Span],
- w: &mut usize,
+ emit: &mut E,
) -> usize {
let (mut p, mut first, mut cursor) = (ls, true, re);
while p < re {
@@ -73,13 +72,10 @@ fn emit_o200k_letters(
} else {
e
};
- unsafe {
- *out.get_unchecked_mut(*w) = Span {
- start: start as u32,
- end: tok_end as u32,
- }
- };
- *w += 1;
+ emit(Span {
+ start: start as u32,
+ end: tok_end as u32,
+ });
first = false;
cursor = tok_end;
p = e;
@@ -94,86 +90,199 @@ fn emit_o200k_letters(
/// Unlike deepseek there are no gaps: rule 4's `[^\s\p{L}\p{N}]+` is a catch-all. Scalar; ┌ OWNER: shared ┐
#[must_use]
pub fn fsm_o200k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
- o200k::(text, tags, out)
+ debug_assert!(out.len() >= text.len());
+ let mut w = 0usize;
+ scan_o200k(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
}
/// Mistral tekken ([`crate::regexes::TEKKEN`]) — the o200k FSM with the contraction suffix off and one
/// token per digit. Every other rule is shared, so both are the same code monomorphized twice.
#[must_use]
pub fn fsm_tekken(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
- o200k::(text, tags, out)
+ debug_assert!(out.len() >= text.len());
+ let mut w = 0usize;
+ scan_tekken(text, tags, |span| {
+ // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`.
+ unsafe { *out.get_unchecked_mut(w) = span };
+ w += 1;
+ });
+ w
+}
+
+/// The scan under [`fsm_o200k`]: hands each token to `emit` the moment it is cut,
+/// so a caller can consume tokens in place instead of collecting a span buffer first.
+pub fn scan_o200k(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) {
+ o200k::(text, tags, emit);
+}
+
+/// The scan under [`fsm_tekken`]; see [`scan_o200k`].
+pub fn scan_tekken(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) {
+ o200k::(text, tags, emit);
+}
+
+/// Is tag `t` a real `[\p{L}\p{M}]` member? Coarse `Letter` (any case) is always in; coarse
+/// `Mark` is in only as a true `\p{M}` — ALPHA_SYM (`\p{S}`) and ZWJ/ZWNJ (`\p{Cf}`) are `\w`
+/// but not `[\p{L}\p{M}]`.
+#[inline(always)]
+fn member(t: u8) -> bool {
+ let c = t & 0x0F;
+ c == LET || (c == MRK && t != ASM && t != ZWJ)
+}
+
+#[inline(always)]
+fn is_lm(tags: &[u8], a: usize, end: usize) -> bool {
+ a < end && member(tags[a])
}
-fn o200k(
+/// Maximal `[\p{L}\p{M}]+` run from `a` (byte-wise; continuation bytes ride along — see `run_end`).
+#[inline(always)]
+fn letter_end(tags: &[u8], a: usize, end: usize) -> usize {
+ let mut p = a;
+ // logos-style fast loop: 16 tags/chunk, one bounds check, unchecked reads. A plain `Letter`
+ // (low nibble 0, incl Han — o200k keeps all letters) or a `Cont` byte stays in-run; only a
+ // coarse `Mark` lane pays the refinement test. Byte-exact with the plain scan below.
+ // SAFETY: `p + 16 <= end <= tags.len()` in the body.
+ while p + 16 <= end {
+ let mut brk = 16;
+ for k in 0..16 {
+ let t = unsafe { *tags.get_unchecked(p + k) };
+ if t == CONT || t & 0x0F == LET {
+ continue;
+ }
+ if t & 0x0F == MRK && t != ASM && t != ZWJ {
+ continue;
+ }
+ brk = k;
+ break;
+ }
+ if brk < 16 {
+ return p + brk;
+ }
+ p += 16;
+ }
+ while p < end && (tags[p] == CONT || member(tags[p])) {
+ p += 1;
+ }
+ p
+}
+
+/// rule 4 `[^\s\p{L}\p{N}]+[\r\n/]*` from `sp0` (any leading space already consumed); `sp0` if
+/// none. `/` is in the `+` body too — the trailing class only matters after the `+` stops at a
+/// `\r\n`.
+#[inline(always)]
+fn other(text: &[u8], tags: &[u8], sp0: usize, end: usize) -> usize {
+ let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N);
+ if p > sp0 {
+ while p < end && (tags[p] == NLN || text[p] == b'/') {
+ p += char_len(text[p]);
+ }
+ }
+ p
+}
+
+/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of
+/// the o200k-family regex. A letter-path token is the FIRST case sub-run from its start (the
+/// case split restarts at every sub-token, so one dispatch per sub-token equals the run loop in
+/// [`emit_o200k_letters`]); the masked scanner re-derives tokens with this where its batch
+/// masks are not trustworthy.
+#[inline(always)]
+pub(super) fn advance_o200k(
text: &[u8],
tags: &[u8],
- out: &mut [Span],
+ i: usize,
+ end: usize,
) -> usize {
- debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
- let end = text.len();
- // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
- // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
- let tags = &tags[..end];
-
- // Is tag `t` (at byte `p`) a real `[\p{L}\p{M}]` member? Coarse `Letter` (any case) is always in;
- // coarse `Mark` is in only as a true `\p{M}` — ALPHA_SYM (`\p{S}`) and ZWJ/ZWNJ (`\p{Cf}`) are `\w`
- // but not `[\p{L}\p{M}]`. The `ds_is_zwj` byte-peek is thus paid ONLY for Marks, never for letters.
- let member = |t: u8, p: usize| -> bool {
- let c = t & 0x0F;
- let _ = p;
- c == LET || (c == MRK && t != ASM && t != ZWJ)
+ let one_letter = |ls: usize| -> usize {
+ let re = letter_end(tags, ls, end);
+ let e = o200k_letter_match(tags, ls, re);
+ if CONTRACTION && e == re {
+ e + contraction(text, e)
+ } else {
+ e
+ }
};
- let is_lm = |a: usize| a < end && member(tags[a], a);
- // maximal `[\p{L}\p{M}]+` run from `a` (byte-wise; continuation bytes ride along — see `run_end`).
- let letter_end = |a: usize| -> usize {
- let mut p = a;
- // logos-style fast loop: 16 tags/chunk, one bounds check, unchecked reads. A plain `Letter`
- // (low nibble 0, incl Han — o200k keeps all letters) or a `Cont` byte stays in-run with no
- // `ds_is_zwj` peek; only a coarse `Mark` lane pays the peek. Byte-exact with the scalar scan.
- // SAFETY: `p + 16 <= end <= tags.len()`/`text.len()` in the body.
- while p + 16 <= end {
- let mut brk = 16;
- for k in 0..16 {
- let t = unsafe { *tags.get_unchecked(p + k) };
- if t == CONT || t & 0x0F == LET {
- continue;
- }
- if t & 0x0F == MRK && t != ASM && t != ZWJ {
- continue;
+ let b = text[i];
+ match tags[i] & 0x0F {
+ NW | NO => {
+ let (mut p, mut cnt) = (i, 0);
+ while p < end && cnt < DIGIT_CAP && in_mask(tags[p], mask::NUMBER) {
+ p += char_len(text[p]);
+ cnt += 1;
+ }
+ p
+ }
+ LET | MRK => {
+ if tags[i] != ASM && tags[i] != ZWJ {
+ one_letter(i)
+ } else {
+ let a = i + char_len(b);
+ if is_lm(tags, a, end) {
+ one_letter(a)
+ } else {
+ other(text, tags, i, end)
}
- brk = k;
- break;
}
- if brk < 16 {
- return p + brk;
+ }
+ SPC => {
+ let a = i + 1; // Space is ASCII (0x20)
+ if is_lm(tags, a, end) {
+ one_letter(a)
+ } else {
+ let p = other(text, tags, a, end);
+ if p > a {
+ p
+ } else {
+ ws_tail(text, tags, i, end)
+ }
}
- p += 16;
}
- while p < end && (tags[p] == CONT || member(tags[p], p)) {
- p += 1;
+ WSO => {
+ let a = i + char_len(b);
+ if is_lm(tags, a, end) {
+ one_letter(a)
+ } else {
+ ws_tail(text, tags, i, end)
+ }
}
- p
- };
- // rule 4 `[^\s\p{L}\p{N}]+[\r\n/]*` from `sp0` (any leading space already consumed); `sp0` if none.
- // `/` is in the `+` body too — the trailing class only matters after the `+` stops at a `\r\n`.
- let other = |sp0: usize| -> usize {
- let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N);
- if p > sp0 {
- while p < end && (tags[p] == NLN || text[p] == b'/') {
- p += char_len(text[p]);
+ NLN => ws_tail(text, tags, i, end),
+ CON | PUN | APO | SYM | NMO | CTL => {
+ let a = i + char_len(b);
+ if is_lm(tags, a, end) {
+ one_letter(a)
+ } else {
+ other(text, tags, i, end)
}
}
- p
- };
+ _ => i + char_len(b),
+ }
+}
+
+fn o200k(
+ text: &[u8],
+ tags: &[u8],
+ mut emit: E,
+) {
+ debug_assert!(tags.len() >= text.len());
+ let end = text.len();
+ // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior
+ // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.)
+ let tags = &tags[..end];
+
+ let is_lm = |a: usize| is_lm(tags, a, end);
// rules 5-7 (`\s*[\r\n]+ | \s+(?!\S) | \s+`) → the shared `ws_tail` (identical to cl100k).
let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) };
+ let other = |sp0: usize| -> usize { other(text, tags, sp0, end) };
// The letter rules: case-split the run starting at `ls`, first sub-token starting at the prefix `pfx`.
- let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize {
- emit_o200k_letters::(text, tags, pfx, ls, letter_end(ls), out, w)
+ let letters = |pfx: usize, ls: usize, emit: &mut E| -> usize {
+ emit_o200k_letters::(text, tags, pfx, ls, letter_end(tags, ls, end), emit)
};
let mut i = 0;
- let mut w = 0usize;
while i < end {
let start = i;
let b = text[i];
@@ -191,12 +300,12 @@ fn o200k(
// take the `[^\r\n\p{L}\p{N}]?` prefix / rule-4 path instead.
LET | MRK => {
if tags[i] != ASM && tags[i] != ZWJ {
- i = letters(i, i, out, &mut w);
+ i = letters(i, i, &mut emit);
continue;
}
let a = i + char_len(b);
if is_lm(a) {
- i = letters(i, a, out, &mut w);
+ i = letters(i, a, &mut emit);
continue;
}
i = other(i); // ∈ NOT_WS_L_N ⇒ > i
@@ -205,7 +314,7 @@ fn o200k(
SPC => {
let a = i + 1; // Space is ASCII (0x20)
if is_lm(a) {
- i = letters(i, a, out, &mut w);
+ i = letters(i, a, &mut emit);
continue;
}
let p = other(a);
@@ -215,7 +324,7 @@ fn o200k(
WSO => {
let a = i + char_len(b);
if is_lm(a) {
- i = letters(i, a, out, &mut w);
+ i = letters(i, a, &mut emit);
continue;
}
i = ws(i);
@@ -225,7 +334,7 @@ fn o200k(
CON | PUN | APO | SYM | NMO | CTL => {
let a = i + char_len(b);
if is_lm(a) {
- i = letters(i, a, out, &mut w);
+ i = letters(i, a, &mut emit);
continue;
}
i = other(i); // ∈ NOT_WS_L_N ⇒ > i
@@ -233,14 +342,9 @@ fn o200k(
// Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively.
_ => i += char_len(b),
}
- // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1).
- unsafe {
- *out.get_unchecked_mut(w) = Span {
- start: start as u32,
- end: i as u32,
- }
- };
- w += 1;
+ emit(Span {
+ start: start as u32,
+ end: i as u32,
+ });
}
- w
}
diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs
index 5f22f52d0..888263190 100644
--- a/tokenizers/atomsplit/src/lib.rs
+++ b/tokenizers/atomsplit/src/lib.rs
@@ -3,7 +3,7 @@
//! One SIMD pass ([`classify`]) maps every codepoint to a tiny "atom" alphabet; a family of no-push
//! FSMs ([`fsm`]) turn that atom stream into token spans (byte ranges) — the pre-tokenizer stage that
//! runs before a BPE/WordPiece model. Pre-tokenizers implemented: `WhitespaceSplit`, `Punctuation`,
-//! `Digits`, `Whitespace`, `Bert`, `Cl100k`, `DeepSeek`, `ByteLevel`, `CharDelimiterSplit` (o200k and
+//! `Digits`, `Whitespace`, `Bert`, `Cl100k`, `DeepSeek`, `ByteLevel` (o200k and
//! Mistral's tekken are exposed as the [`fsm::fsm_o200k`] / [`fsm::fsm_tekken`] functions rather than
//! recipe structs).
//!
@@ -28,6 +28,12 @@ mod simd_avx_classify;
#[cfg(target_arch = "aarch64")]
mod simd_classify;
mod simd_fsm;
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+))]
+mod simd_literal;
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
mod simd_wasm_classify;
pub mod tables;
diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs
index c826b2dda..9e7106c1a 100644
--- a/tokenizers/atomsplit/src/literal.rs
+++ b/tokenizers/atomsplit/src/literal.rs
@@ -4,7 +4,15 @@
//! FSMs cut where the class changes ([`crate::fsm`]).
//!
//! The atom classification is unnecessary for pre-tokenizers splitting on an exact character or literal string:
-//! We can use a simpler byte search looking at 16 bytes at a time.
+//! a plain byte search finds the same cuts. Two searches cover the two match densities:
+//!
+//! - [`Literal::matches`] iterates with [`memmem`], which is built for needles that are rare in
+//! the haystack: it skips ahead fast and restarts after every match.
+//! - [`Literal::matches_into`] scans the whole text once and writes every offset into a caller
+//! buffer. Pre-tokenizer delimiters are the opposite of rare (running English text has a space
+//! about every six bytes), and at that density the per-match restarts dominate the iterator.
+//! The scan answers a whole block of text at once and turns the answers into offsets;
+//! `simd_literal.rs` explains how, step by step.
use memchr::memmem;
use std::fmt;
@@ -53,4 +61,135 @@ impl Literal {
pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator- + 't {
self.finder.find_iter(text)
}
+
+ /// The same offsets as [`Literal::matches`], written into `out[..count]` in one scan of
+ /// `text`; returns the count. Use this when matches are frequent: the iterator restarts its
+ /// search machinery after every match, the scan never stops.
+ ///
+ /// # Preconditions
+ /// - `out.len() >= text.len() / pattern.len() + 4` (asserted). The division is the most
+ /// matches the text can hold; the `+ 4` is slack the SIMD path writes past the last match.
+ /// - `text.len()` must fit in `u32` (asserted), so that every offset does too.
+ pub fn matches_into(&self, text: &[u8], out: &mut [u32]) -> usize {
+ assert!(
+ out.len() >= text.len() / self.pattern().len() + 4,
+ "matches_into needs out.len() >= text.len() / pattern.len() + 4"
+ );
+ assert!(
+ u32::try_from(text.len()).is_ok(),
+ "matches_into writes u32 offsets"
+ );
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ if self.scannable() {
+ match *self.pattern() {
+ [a] => return crate::simd_literal::matches_into([a], text, out),
+ [a, b] => return crate::simd_literal::matches_into([a, b], text, out),
+ [a, b, c] => return crate::simd_literal::matches_into([a, b, c], text, out),
+ _ => unreachable!("scannable patterns are one to three bytes"),
+ }
+ }
+ let mut count = 0;
+ for position in self.finder.find_iter(text) {
+ out[count] = position as u32;
+ count += 1;
+ }
+ count
+ }
+
+ /// How many matches [`Literal::matches`] would report. Counting runs the compare step
+ /// without producing any offsets, so it is several times faster than a full scan; count
+ /// first to size an output exactly, then fill it with a [`Literal::for_each_match`] pass.
+ #[must_use]
+ pub fn count_matches(&self, text: &[u8]) -> usize {
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ if self.scannable() {
+ match *self.pattern() {
+ [a] => return crate::simd_literal::count_matches([a], text),
+ [a, b] => return crate::simd_literal::count_matches([a, b], text),
+ [a, b, c] => return crate::simd_literal::count_matches([a, b, c], text),
+ _ => unreachable!("scannable patterns are one to three bytes"),
+ }
+ }
+ self.finder.find_iter(text).count()
+ }
+
+ /// Calls `on_match` with the byte offset of every match, left to right: the same offsets
+ /// as [`Literal::matches`] at batch-scan speed, with no buffer for the caller to provide.
+ /// The scan streams through a small stack window, so its footprint stays flat however
+ /// long the text is.
+ pub fn for_each_match(&self, text: &[u8], mut on_match: impl FnMut(usize)) {
+ if !self.scannable() {
+ for position in self.finder.find_iter(text) {
+ on_match(position);
+ }
+ return;
+ }
+ self.scan_windows(text, |base, matches| {
+ for &position in matches {
+ on_match(base + position as usize);
+ }
+ });
+ }
+
+ /// Whether the batch scan covers this pattern. The scan emits every position where the
+ /// pattern matches, with no ordering check between them; that is only the non-overlapping
+ /// match list when the pattern cannot overlap itself, which one byte comparison per
+ /// length decides: a two-byte pattern overlaps itself when both bytes are equal, a
+ /// three-byte one when last equals first. Longer or self-overlapping patterns search
+ /// through the [`memmem`] engine instead.
+ #[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ ))]
+ fn scannable(&self) -> bool {
+ match *self.pattern() {
+ [_] => true,
+ [a, b] => a != b,
+ [a, _, c] => a != c,
+ _ => false,
+ }
+ }
+
+ /// Without a SIMD kernel there is no batch scan; everything searches through [`memmem`].
+ #[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "x86_64",
+ all(target_arch = "wasm32", target_feature = "simd128")
+ )))]
+ fn scannable(&self) -> bool {
+ false
+ }
+
+ /// Streams the batch scan window by window: [`Literal::matches_into`] fills a stack
+ /// buffer for one stretch of text, `on_window` consumes it (offsets are relative to the
+ /// window's `base`), and the next window starts where a match could no longer fit in the
+ /// previous one, so nothing at a window edge is missed or reported twice.
+ fn scan_windows(&self, text: &[u8], mut on_window: impl FnMut(usize, &[u32])) {
+ // 4KB of stack; a window of (1024 - 4) * pattern.len() bytes fills it exactly.
+ let mut buffer = [0u32; 1024];
+ let width = self.pattern().len();
+ let window = (buffer.len() - 4) * width;
+ let mut base = 0;
+ loop {
+ let end = usize::min(base + window, text.len());
+ let count = self.matches_into(&text[base..end], &mut buffer);
+ on_window(base, &buffer[..count]);
+ if end == text.len() {
+ return;
+ }
+ // A match crossing `end` was not reported; the next window re-covers its
+ // possible starts. Matches of a scannable pattern never overlap, so the reports
+ // stay disjoint.
+ base = end - (width - 1);
+ }
+ }
}
diff --git a/tokenizers/atomsplit/src/simd_literal.rs b/tokenizers/atomsplit/src/simd_literal.rs
new file mode 100644
index 000000000..78e6bfa58
--- /dev/null
+++ b/tokenizers/atomsplit/src/simd_literal.rs
@@ -0,0 +1,783 @@
+//! Finds every place a short pattern (up to three bytes) occurs in a text, in one pass.
+//! This is the scan behind [`Literal::matches_into`](crate::literal::Literal::matches_into);
+//! [`crate::literal`] explains when it is picked over the plain iterator.
+//!
+//! # Why a dedicated scan
+//!
+//! A pre-tokenizer splits on a delimiter that shows up every few bytes: running English text
+//! has a space about every six bytes, SentencePiece text has a `▁` at every word. A
+//! next-match engine like `memmem` is built for the opposite case, a needle that is rare, so
+//! at this density it spends most of its time stopping and restarting: every match returns to
+//! the caller, and every call re-enters the search machinery. This scan never stops. It
+//! answers "where are ALL the matches in these 16 bytes?" with a handful of instructions,
+//! block after block.
+//!
+//! # How: compare, mask, decode
+//!
+//! Two words come up in everything below. A SIMD **register** is an extra-wide CPU value
+//! holding 16 bytes side by side; one instruction operates on all 16 at once. Each of those
+//! 16 byte slots is a **lane**. A SIMD compare answers lane by lane: a lane becomes all ones
+//! where the answer is yes, all zeros where it is no.
+//!
+//! **Step 1, compare.** Compare the text against the pattern's first byte (copied across all
+//! 16 lanes), the text shifted by one against the second byte, and so on, then AND the
+//! answers together: a lane that passed every compare is the start of a whole match, and
+//! nothing needs a second look. Searching for `▁` (three bytes, `E2 96 81`) in
+//! `"a…b▁c"`, where the `…` (`E2 80 A6`) shares a first byte with `▁` and acts as the decoy:
+//!
+//! ```text
+//! offset 0 1 2 3 4 5 6 7 8
+//! text a E2 80 A6 b E2 96 81 c
+//! └── … ──────┘ └── ▁ ──────┘
+//! text[j] == E2 ? · ✓ · · · ✓ · · ·
+//! text[j+1] == 96 ? · · · · · ✓ · · ·
+//! text[j+2] == 81 ? · · · · · ✓ · · ·
+//! AND of the three · · · · · ✓ · · ·
+//! ```
+//!
+//! Only offset 5 passes all three rows: a match starts at 5, and the decoy at 1 is out after
+//! the second row.
+//!
+//! **Step 2, mask.** The 16 lane answers are squeezed into one ordinary integer so plain
+//! arithmetic can take over: bit `j` set means "a match starts at offset `j`". For the
+//! example above the mask is `0b100000`, bit 5. Each target has a one-or-two instruction way
+//! to build this integer; on NEON the bits come out spaced four positions apart instead of
+//! one, a spacing the shared code carries along as `SHIFT` (the per-target layer below).
+//!
+//! **Step 3, decode.** Count-trailing-zeros on the mask gives the lowest set bit, which is the
+//! next match offset; clear that bit and repeat. The one twist is that the code decodes four
+//! offsets *without checking whether any bits are left*. Say a block starting at text offset
+//! `base` had matches at offsets 2 and 5 (shown with the one-bit spacing):
+//!
+//! ```text
+//! mask 0b100100 trailing zeros = 2 write base + 2 clear bit → 0b100000
+//! mask 0b100000 trailing zeros = 5 write base + 5 clear bit → 0
+//! mask 0 trailing zeros = 64 write garbage (slot 3)
+//! mask 0 trailing zeros = 64 write garbage (slot 4)
+//! cursor += 2 (how many bits were set)
+//! ```
+//!
+//! The two garbage slots sit past the cursor, so the next block's writes (or nothing) land on
+//! them and they never reach the caller. The honest alternative, checking "still something
+//! left?" before each decode, puts a branch in the hottest loop; with a match every four to
+//! eight bytes that branch guesses wrong about once per block, which costs more than all the
+//! compares together.
+//!
+//! Sparse text takes none of this: one instruction per 64-byte block answers "nothing here"
+//! (see `any`), so a pattern that is rare or absent stays close to `memmem` speed.
+//!
+//! The compare trick is Wojciech Muła's ("SIMD-friendly algorithms for substring searching");
+//! the branch-free decode is how simdjson reads its masks.
+//!
+//! # The per-target layer
+//!
+//! Everything from the mask on is plain integer arithmetic and is shared. A target provides
+//! four primitives in its `arch` module:
+//!
+//! - `splat`: one byte copied into every lane, the shape a compare wants.
+//! - `match_starts`: step 1, the shifted compares ANDed together.
+//! - `any`: "did anything in these 64 bytes match?", the cheap exit for sparse text.
+//! - `bits`: step 2, lane answers to integer mask.
+//!
+//! x86 (`movemask`) and wasm (`bitmask`) each have an instruction that grabs one bit from
+//! every lane, so offset `j` lands on bit `j`. NEON has no such instruction; its standard
+//! substitute (Danila Kutenin's) leaves the bits four positions apart, so offset `j` lands
+//! on bit `4 * j`. `SHIFT` is that spacing (`bit position >> SHIFT` recovers the offset),
+//! and it is the only difference the shared code ever sees.
+//!
+//! On x86 the 16-byte baseline is SSE2, which every x86_64 CPU has; when the CPU reports
+//! wider vectors, [`matches_into`] and [`count_matches`] pick a faster front end per call
+//! (the `wide` module: AVX2 and AVX-512).
+//!
+//! The wasm layer is compile-checked but not run in CI; the other targets run the full
+//! `tests/literal.rs` parity suite.
+
+#[cfg(target_arch = "aarch64")]
+mod arch {
+ use core::arch::aarch64::*;
+
+ /// 16 bytes of text in one register.
+ pub(super) type Chunk = uint8x16_t;
+ /// Offsets sit four bit positions apart in a mask from [`bits`]: offset `j` is bit `4 * j`.
+ pub(super) const SHIFT: u32 = 2;
+
+ /// One byte copied into every lane.
+ #[inline(always)]
+ pub(super) fn splat(byte: u8) -> Chunk {
+ unsafe { vdupq_n_u8(byte) }
+ }
+
+ /// All-ones in every lane where a whole `K`-byte match starts.
+ ///
+ /// # Safety
+ /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each
+ /// of `p`, `p + 1`, .., `p + K - 1`.
+ #[inline(always)]
+ pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk {
+ // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`,
+ // stays inside that readable range.
+ let mut starts = unsafe { vceqq_u8(vld1q_u8(p), pattern[0]) };
+ for (k, &byte) in pattern.iter().enumerate().skip(1) {
+ starts = unsafe { vandq_u8(starts, vceqq_u8(vld1q_u8(p.add(k)), byte)) };
+ }
+ starts
+ }
+
+ /// `true` when anything matched in the four chunks.
+ #[inline(always)]
+ pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool {
+ unsafe { vmaxvq_u8(vorrq_u8(vorrq_u8(s0, s1), vorrq_u8(s2, s3))) != 0 }
+ }
+
+ /// The lane answers as one integer: one bit per lane, four bit positions apart.
+ ///
+ /// NEON cannot grab one bit from every lane in a single instruction. Instead `vshrn`
+ /// ("shift right and narrow") halves the register to 64 bits, shrinking each lane's
+ /// answer from eight bits to four (still all ones or all zeros); the `0x1111..` mask
+ /// then keeps one bit of each four, so clearing one bit in the decode drops exactly
+ /// one match.
+ #[inline(always)]
+ pub(super) fn bits(starts: Chunk) -> u64 {
+ unsafe {
+ let halved = vshrn_n_u16::<4>(vreinterpretq_u16_u8(starts));
+ vget_lane_u64::<0>(vreinterpret_u64_u8(halved)) & 0x1111_1111_1111_1111
+ }
+ }
+}
+
+#[cfg(target_arch = "x86_64")]
+mod arch {
+ use core::arch::x86_64::*;
+
+ /// 16 bytes of text in one register.
+ pub(super) type Chunk = __m128i;
+ /// One bit per lane in a mask from [`bits`]: offset `j` is bit `j`, nothing to shift.
+ pub(super) const SHIFT: u32 = 0;
+
+ /// One byte copied into every lane.
+ #[inline(always)]
+ pub(super) fn splat(byte: u8) -> Chunk {
+ unsafe { _mm_set1_epi8(byte as i8) }
+ }
+
+ /// All-ones in every lane where a whole `K`-byte match starts.
+ ///
+ /// # Safety
+ /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each
+ /// of `p`, `p + 1`, .., `p + K - 1`.
+ #[inline(always)]
+ pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk {
+ // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`,
+ // stays inside that readable range.
+ let mut starts = unsafe { _mm_cmpeq_epi8(_mm_loadu_si128(p.cast()), pattern[0]) };
+ for (k, &byte) in pattern.iter().enumerate().skip(1) {
+ starts = unsafe {
+ _mm_and_si128(
+ starts,
+ _mm_cmpeq_epi8(_mm_loadu_si128(p.add(k).cast()), byte),
+ )
+ };
+ }
+ starts
+ }
+
+ /// `true` when anything matched in the four chunks.
+ #[inline(always)]
+ pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool {
+ unsafe { _mm_movemask_epi8(_mm_or_si128(_mm_or_si128(s0, s1), _mm_or_si128(s2, s3))) != 0 }
+ }
+
+ /// The lane answers as one integer: `movemask` grabs each lane's top bit, so offset `j`
+ /// is bit `j`.
+ #[inline(always)]
+ pub(super) fn bits(starts: Chunk) -> u64 {
+ unsafe { _mm_movemask_epi8(starts) as u32 as u64 }
+ }
+}
+
+#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
+mod arch {
+ use core::arch::wasm32::*;
+
+ /// 16 bytes of text in one register.
+ pub(super) type Chunk = v128;
+ /// One bit per lane in a mask from [`bits`]: offset `j` is bit `j`, nothing to shift.
+ pub(super) const SHIFT: u32 = 0;
+
+ /// One byte copied into every lane.
+ #[inline(always)]
+ pub(super) fn splat(byte: u8) -> Chunk {
+ u8x16_splat(byte)
+ }
+
+ /// All-ones in every lane where a whole `K`-byte match starts.
+ ///
+ /// # Safety
+ /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each
+ /// of `p`, `p + 1`, .., `p + K - 1`.
+ #[inline(always)]
+ pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk {
+ // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`,
+ // stays inside that readable range.
+ let mut starts = unsafe { u8x16_eq(v128_load(p.cast()), pattern[0]) };
+ for (k, &byte) in pattern.iter().enumerate().skip(1) {
+ starts = unsafe { v128_and(starts, u8x16_eq(v128_load(p.add(k).cast()), byte)) };
+ }
+ starts
+ }
+
+ /// `true` when anything matched in the four chunks.
+ #[inline(always)]
+ pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool {
+ v128_any_true(v128_or(v128_or(s0, s1), v128_or(s2, s3)))
+ }
+
+ /// The lane answers as one integer: `bitmask` grabs each lane's top bit, so offset `j`
+ /// is bit `j`.
+ #[inline(always)]
+ pub(super) fn bits(starts: Chunk) -> u64 {
+ u8x16_bitmask(starts) as u64
+ }
+}
+
+/// Decodes every set bit of `bits` into a `base + offset` value at `cursor`, and returns the
+/// cursor advanced by the true match count.
+///
+/// The first four decodes run without checking `bits` (the module docs say why): once the mask
+/// is used up, `trailing_zeros` is 64 and the slot gets a garbage offset, scratch that the
+/// count never covers and a later call overwrites. Blocks with more than four matches finish
+/// in the loop.
+///
+/// # Safety
+/// There must be room for `max(4, count_ones(bits))` writes between `cursor` and `end`: four
+/// slots for the unconditional decodes, one per match when a block has more. Debug builds
+/// check this bound before writing anything.
+#[inline(always)]
+unsafe fn decode(mut bits: u64, base: u32, cursor: *mut u32, end: *const u32) -> *mut u32 {
+ let count = bits.count_ones() as usize;
+ debug_assert!(
+ end as usize - cursor as usize >= count.max(4) * size_of::(),
+ "decode would write past the match buffer"
+ );
+ // SAFETY: the caller's contract above gives these four slots.
+ for slot in 0..4 {
+ unsafe { *cursor.add(slot) = base + (bits.trailing_zeros() >> arch::SHIFT) };
+ bits &= bits.wrapping_sub(1);
+ }
+ if bits != 0 {
+ let mut slot = 4;
+ // SAFETY: one slot per remaining match, still within the caller's `count` slots.
+ while bits != 0 {
+ unsafe { *cursor.add(slot) = base + (bits.trailing_zeros() >> arch::SHIFT) };
+ bits &= bits.wrapping_sub(1);
+ slot += 1;
+ }
+ }
+ // SAFETY: `count` slots were just written, so one past the last of them is in bounds.
+ unsafe { cursor.add(count) }
+}
+
+/// How many matches of `pattern` are in `text`: the compare and mask steps only, counting set
+/// bits instead of decoding them, which makes counting several times faster than the full
+/// scan. Same pattern rules as [`matches_into`].
+pub(crate) fn count_matches(pattern: [u8; K], text: &[u8]) -> usize {
+ const {
+ assert!(
+ K >= 1 && K <= 3,
+ "the scan is built for patterns of 1 to 3 bytes"
+ );
+ }
+ let len = text.len();
+ if len < K {
+ return 0;
+ }
+ #[cfg(target_arch = "x86_64")]
+ if let Some(count) = wide::count_matches_wide::(pattern, text) {
+ return count;
+ }
+ let p = text.as_ptr();
+ let mut pattern_chunks = [arch::splat(0); K];
+ for (chunk, &byte) in pattern_chunks.iter_mut().zip(&pattern) {
+ *chunk = arch::splat(byte);
+ }
+ let mut count = 0usize;
+ let mut i = 0usize;
+
+ // The same block layout as `scan`, and the same SAFETY story for the reads: each loop
+ // condition bounds its loads inside `text`. Nothing here writes.
+ while i + 64 + (K - 1) <= len {
+ let (s0, s1, s2, s3) = unsafe {
+ (
+ arch::match_starts::(p.add(i), &pattern_chunks),
+ arch::match_starts::(p.add(i + 16), &pattern_chunks),
+ arch::match_starts::(p.add(i + 32), &pattern_chunks),
+ arch::match_starts::(p.add(i + 48), &pattern_chunks),
+ )
+ };
+ if arch::any(s0, s1, s2, s3) {
+ count += (arch::bits(s0).count_ones()
+ + arch::bits(s1).count_ones()
+ + arch::bits(s2).count_ones()
+ + arch::bits(s3).count_ones()) as usize;
+ }
+ i += 64;
+ }
+
+ count + unsafe { count_tail::(pattern, &pattern_chunks, text, i) }
+}
+
+/// The counting twin of [`scan_tail`]: 16-byte rounds from `i`, the overlapped final block,
+/// or the byte-by-byte walk for texts shorter than one block.
+///
+/// # Safety
+/// Reads are bounded by the loop conditions, exactly as in [`scan_tail`]. Nothing writes.
+unsafe fn count_tail(
+ pattern: [u8; K],
+ pattern_chunks: &[arch::Chunk; K],
+ text: &[u8],
+ mut i: usize,
+) -> usize {
+ let len = text.len();
+ let p = text.as_ptr();
+ let mut count = 0usize;
+
+ // SAFETY: reads end at `i + 16 + (K - 1) <= len`.
+ while i + 16 + (K - 1) <= len {
+ let starts = unsafe { arch::match_starts::(p.add(i), pattern_chunks) };
+ count += arch::bits(starts).count_ones() as usize;
+ i += 16;
+ }
+
+ if i + K <= len {
+ if len >= 16 + (K - 1) {
+ // The overlapped final block: bits below `i` were counted above and are masked off.
+ // SAFETY: the block reads `base .. base + 16 + K - 1`, which ends exactly at `len`.
+ let base = len - 16 - (K - 1);
+ let starts = unsafe { arch::match_starts::(p.add(base), pattern_chunks) };
+ let bits = arch::bits(starts) & (!0u64 << ((i - base) << arch::SHIFT));
+ count += bits.count_ones() as usize;
+ } else {
+ while i + K <= len {
+ count += (text[i..i + K] == pattern[..]) as usize;
+ i += 1;
+ }
+ }
+ }
+ count
+}
+
+/// The wider x86 front ends. Every x86_64 CPU has the SSE2 baseline above; these two rungs
+/// are picked per call when the CPU reports the feature:
+///
+/// - **AVX2** compares 32 bytes per instruction, so one round covers 64 bytes with two
+/// compares and one combined mask.
+/// - **AVX-512** compares 64 bytes straight into a bit mask (`_mm512_cmpeq_epi8_mask`, no
+/// separate mask step) and decodes with `vpcompressd`: the CPU packs the matching offsets
+/// itself, 16 at a time, replacing the whole count-trailing-zeros loop.
+///
+/// Each rung is its own function because `#[target_feature]` applies per function; the bodies
+/// mirror [`scan`] and [`count_matches`] round for round and share [`decode`], [`scan_tail`]
+/// and [`count_tail`]. The AVX2 rung is exercised by the `tests/literal.rs` suite on any
+/// AVX2 machine or emulator; the AVX-512 rung runs the same suite only on hardware that has
+/// it, so run the tests on such a machine before trusting changes to it.
+#[cfg(target_arch = "x86_64")]
+mod wide {
+ use super::{arch, count_tail, decode, scan_tail};
+ use core::arch::x86_64::*;
+
+ /// Runs the widest scan this CPU supports; `None` means only the baseline exists.
+ #[inline]
+ pub(super) fn matches_into_wide(
+ pattern: [u8; K],
+ text: &[u8],
+ out: &mut [u32],
+ ) -> Option {
+ if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
+ // SAFETY: the features were just detected; the buffer contract is the caller's.
+ return Some(unsafe { scan_avx512::(pattern, text, out) });
+ }
+ if is_x86_feature_detected!("avx2") {
+ // SAFETY: the feature was just detected; the buffer contract is the caller's.
+ return Some(unsafe { scan_avx2::(pattern, text, out) });
+ }
+ None
+ }
+
+ /// The counting twin of [`matches_into_wide`].
+ #[inline]
+ pub(super) fn count_matches_wide(
+ pattern: [u8; K],
+ text: &[u8],
+ ) -> Option {
+ if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
+ // SAFETY: the features were just detected. Counting never writes.
+ return Some(unsafe { count_avx512::(pattern, text) });
+ }
+ if is_x86_feature_detected!("avx2") {
+ // SAFETY: the feature was just detected. Counting never writes.
+ return Some(unsafe { count_avx2::(pattern, text) });
+ }
+ None
+ }
+
+ /// All-ones per lane where a whole `K`-byte match starts, for 32 positions at `p`.
+ ///
+ /// # Safety
+ /// AVX2 must be available, and `p .. p + 32 + K - 1` must be readable.
+ #[inline]
+ #[target_feature(enable = "avx2")]
+ unsafe fn match_starts32(p: *const u8, pattern: &[__m256i; K]) -> __m256i {
+ // SAFETY: the caller's contract above; each load of 32 bytes from `p + k`, `k < K`,
+ // stays inside that readable range.
+ let mut starts = _mm256_cmpeq_epi8(unsafe { _mm256_loadu_si256(p.cast()) }, pattern[0]);
+ for (k, &byte) in pattern.iter().enumerate().skip(1) {
+ starts = _mm256_and_si256(
+ starts,
+ _mm256_cmpeq_epi8(unsafe { _mm256_loadu_si256(p.add(k).cast()) }, byte),
+ );
+ }
+ starts
+ }
+
+ /// One byte copied into every lane, and the baseline chunks [`scan_tail`] wants.
+ #[target_feature(enable = "avx2")]
+ fn splat32(pattern: [u8; K]) -> ([__m256i; K], [arch::Chunk; K]) {
+ let mut wide = [_mm256_set1_epi8(0); K];
+ let mut baseline = [arch::splat(0); K];
+ for k in 0..K {
+ wide[k] = _mm256_set1_epi8(pattern[k] as i8);
+ baseline[k] = arch::splat(pattern[k]);
+ }
+ (wide, baseline)
+ }
+
+ /// [`scan`](super::scan) with the AVX2 front end: 64 bytes per round as two 32-byte
+ /// compares, one combined mask, decoded in the same 16-bit quarters as the baseline so
+ /// [`decode`]'s four-slot unroll stays matched to typical match density.
+ ///
+ /// # Safety
+ /// AVX2 must be available, and `out` must satisfy
+ /// [`matches_into`](super::matches_into)'s buffer bound, which it asserts.
+ #[target_feature(enable = "avx2")]
+ unsafe fn scan_avx2(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize {
+ let len = text.len();
+ if len < K {
+ return 0;
+ }
+ let p = text.as_ptr();
+ let (pattern_chunks, baseline) = splat32::(pattern);
+ let start = out.as_mut_ptr();
+ // SAFETY: one past the buffer's last slot is a valid address to form (never read).
+ let end = unsafe { start.add(out.len()) } as *const u32;
+ let mut cursor = start;
+ let mut i = 0usize;
+
+ // SAFETY: reads end at `i + 32 + 32 + (K - 1) <= len`; writes are the buffer
+ // contract, handed to `decode` as `end`.
+ while i + 64 + (K - 1) <= len {
+ let (s0, s1) = unsafe {
+ (
+ match_starts32::(p.add(i), &pattern_chunks),
+ match_starts32::(p.add(i + 32), &pattern_chunks),
+ )
+ };
+ let bits = (_mm256_movemask_epi8(s0) as u32 as u64)
+ | ((_mm256_movemask_epi8(s1) as u32 as u64) << 32);
+ if bits != 0 {
+ unsafe {
+ cursor = decode(bits & 0xFFFF, i as u32, cursor, end);
+ cursor = decode((bits >> 16) & 0xFFFF, i as u32 + 16, cursor, end);
+ cursor = decode((bits >> 32) & 0xFFFF, i as u32 + 32, cursor, end);
+ cursor = decode(bits >> 48, i as u32 + 48, cursor, end);
+ }
+ }
+ i += 64;
+ }
+ // SAFETY: `scan_tail` continues under the same contract.
+ cursor = unsafe { scan_tail::(pattern, &baseline, text, i, cursor, end) };
+ // SAFETY: `cursor` and `start` both point into `out`'s buffer.
+ unsafe { cursor.offset_from(start) as usize }
+ }
+
+ /// The counting twin of [`scan_avx2`].
+ ///
+ /// # Safety
+ /// AVX2 must be available.
+ #[target_feature(enable = "avx2")]
+ unsafe fn count_avx2(pattern: [u8; K], text: &[u8]) -> usize {
+ let len = text.len();
+ if len < K {
+ return 0;
+ }
+ let p = text.as_ptr();
+ let (pattern_chunks, baseline) = splat32::(pattern);
+ let mut count = 0usize;
+ let mut i = 0usize;
+
+ // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Nothing writes.
+ while i + 64 + (K - 1) <= len {
+ let (s0, s1) = unsafe {
+ (
+ match_starts32::(p.add(i), &pattern_chunks),
+ match_starts32::(p.add(i + 32), &pattern_chunks),
+ )
+ };
+ count += (_mm256_movemask_epi8(s0).count_ones() + _mm256_movemask_epi8(s1).count_ones())
+ as usize;
+ i += 64;
+ }
+ // SAFETY: `count_tail` continues under the same read bounds.
+ count + unsafe { count_tail::(pattern, &baseline, text, i) }
+ }
+
+ /// The 64-byte match mask, straight from the compares: AVX-512 byte compares produce a
+ /// bit per lane natively, so there is no separate mask step to pay for.
+ ///
+ /// # Safety
+ /// AVX-512 F and BW must be available, and `p .. p + 64 + K - 1` must be readable.
+ #[inline]
+ #[target_feature(enable = "avx512f,avx512bw")]
+ unsafe fn match_bits64(p: *const u8, pattern: &[__m512i; K]) -> u64 {
+ // SAFETY: the caller's contract above; each load of 64 bytes from `p + k`, `k < K`,
+ // stays inside that readable range.
+ let mut bits = _mm512_cmpeq_epi8_mask(unsafe { _mm512_loadu_si512(p.cast()) }, pattern[0]);
+ for (k, &byte) in pattern.iter().enumerate().skip(1) {
+ bits &= _mm512_cmpeq_epi8_mask(unsafe { _mm512_loadu_si512(p.add(k).cast()) }, byte);
+ }
+ bits
+ }
+
+ /// One byte copied into every lane, and the baseline chunks the shared tails want.
+ #[target_feature(enable = "avx512f")]
+ fn splat64(pattern: [u8; K]) -> ([__m512i; K], [arch::Chunk; K]) {
+ let mut wide = [_mm512_set1_epi8(0); K];
+ let mut baseline = [arch::splat(0); K];
+ for k in 0..K {
+ wide[k] = _mm512_set1_epi8(pattern[k] as i8);
+ baseline[k] = arch::splat(pattern[k]);
+ }
+ (wide, baseline)
+ }
+
+ /// [`scan`](super::scan) with the AVX-512 front end. The decode is `vpcompressd`
+ /// (`_mm512_mask_compressstoreu_epi32`): give the CPU sixteen candidate offsets and a
+ /// mask, and it stores exactly the matching ones, packed. The count-trailing-zeros loop
+ /// and its density concerns disappear; four compress-stores cover a 64-byte round.
+ ///
+ /// # Safety
+ /// AVX-512 F and BW must be available, and `out` must satisfy
+ /// [`matches_into`](super::matches_into)'s buffer bound, which it asserts.
+ #[target_feature(enable = "avx512f,avx512bw")]
+ unsafe fn scan_avx512(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize {
+ let len = text.len();
+ if len < K {
+ return 0;
+ }
+ let p = text.as_ptr();
+ let (pattern_chunks, baseline) = splat64::(pattern);
+ let start = out.as_mut_ptr();
+ // SAFETY: one past the buffer's last slot is a valid address to form (never read).
+ let end = unsafe { start.add(out.len()) } as *const u32;
+ let mut cursor = start;
+ let mut i = 0usize;
+ // Lane numbers 0..16, the base of every offset vector.
+ let lanes = _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0);
+
+ // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Each compress-store writes exactly
+ // the matches of its quarter, within the buffer contract.
+ while i + 64 + (K - 1) <= len {
+ let bits = unsafe { match_bits64::(p.add(i), &pattern_chunks) };
+ if bits != 0 {
+ let base = _mm512_add_epi32(lanes, _mm512_set1_epi32(i as i32));
+ for quarter in 0..4 {
+ let mask = ((bits >> (16 * quarter)) & 0xFFFF) as __mmask16;
+ let offsets = _mm512_add_epi32(base, _mm512_set1_epi32(16 * quarter));
+ debug_assert!(
+ end as usize - cursor as usize
+ >= mask.count_ones() as usize * size_of::(),
+ "compress-store would write past the match buffer"
+ );
+ unsafe {
+ _mm512_mask_compressstoreu_epi32(cursor.cast(), mask, offsets);
+ cursor = cursor.add(mask.count_ones() as usize);
+ }
+ }
+ }
+ i += 64;
+ }
+ // SAFETY: `scan_tail` continues under the same contract.
+ cursor = unsafe { scan_tail::(pattern, &baseline, text, i, cursor, end) };
+ // SAFETY: `cursor` and `start` both point into `out`'s buffer.
+ unsafe { cursor.offset_from(start) as usize }
+ }
+
+ /// The counting twin of [`scan_avx512`].
+ ///
+ /// # Safety
+ /// AVX-512 F and BW must be available.
+ #[target_feature(enable = "avx512f,avx512bw")]
+ unsafe fn count_avx512(pattern: [u8; K], text: &[u8]) -> usize {
+ let len = text.len();
+ if len < K {
+ return 0;
+ }
+ let p = text.as_ptr();
+ let (pattern_chunks, baseline) = splat64::(pattern);
+ let mut count = 0usize;
+ let mut i = 0usize;
+
+ // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Nothing writes.
+ while i + 64 + (K - 1) <= len {
+ count += unsafe { match_bits64::(p.add(i), &pattern_chunks) }.count_ones() as usize;
+ i += 64;
+ }
+ // SAFETY: `count_tail` continues under the same read bounds.
+ count + unsafe { count_tail::(pattern, &baseline, text, i) }
+ }
+}
+
+/// One-pass scan: writes the start offset of every match of `pattern` in `text` into `out` and
+/// returns how many. The pattern must not be able to overlap itself, so that every matching
+/// position belongs to the non-overlapping match list;
+/// [`Literal::matches_into`](crate::literal::Literal::matches_into) routes self-overlapping
+/// patterns away.
+///
+/// The buffer bound is asserted here, not trusted from the caller: everything unsafe in this
+/// module leans on it, so the check lives next to what it protects.
+pub(crate) fn matches_into(
+ pattern: [u8; K],
+ text: &[u8],
+ out: &mut [u32],
+) -> usize {
+ const {
+ assert!(
+ K >= 1 && K <= 3,
+ "the scan is built for patterns of 1 to 3 bytes"
+ );
+ }
+ assert!(
+ out.len() >= text.len() / K + 4,
+ "the match buffer must hold text.len() / pattern.len() + 4 offsets"
+ );
+ #[cfg(target_arch = "x86_64")]
+ if let Some(count) = wide::matches_into_wide::(pattern, text, out) {
+ return count;
+ }
+ // SAFETY: the assert above is exactly `scan`'s buffer contract.
+ unsafe { scan::(pattern, text, out) }
+}
+
+/// # Safety
+/// `out.len() >= text.len() / K + 4` must hold; [`matches_into`] asserts it. Every match found
+/// counts one slot and matches cannot overlap, so at most `text.len() / K` slots are counted
+/// and the four extra keep `decode`'s unconditional writes inside the buffer. Reads from
+/// `text` are justified region by region below.
+unsafe fn scan(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize {
+ let len = text.len();
+ if len < K {
+ return 0;
+ }
+ let p = text.as_ptr();
+ let mut pattern_chunks = [arch::splat(0); K];
+ for (chunk, &byte) in pattern_chunks.iter_mut().zip(&pattern) {
+ *chunk = arch::splat(byte);
+ }
+ let start = out.as_mut_ptr();
+ // SAFETY: one past the buffer's last slot is a valid address to form (never read);
+ // only `decode`'s debug bound check uses it.
+ let end = unsafe { start.add(out.len()) } as *const u32;
+ let mut cursor = start;
+ let mut i = 0usize;
+
+ // 64 bytes per round; `any` lets a match-free round skip the masks and decodes.
+ // SAFETY: the loop condition keeps every read inside `text`: the furthest `match_starts`
+ // begins at `p + i + 48` and reads `16 + K - 1` bytes, ending at `i + 64 + (K - 1) <= len`.
+ // Writes are `scan`'s buffer contract, handed to `decode` as `end`.
+ while i + 64 + (K - 1) <= len {
+ let (s0, s1, s2, s3) = unsafe {
+ (
+ arch::match_starts::(p.add(i), &pattern_chunks),
+ arch::match_starts::(p.add(i + 16), &pattern_chunks),
+ arch::match_starts::(p.add(i + 32), &pattern_chunks),
+ arch::match_starts::(p.add(i + 48), &pattern_chunks),
+ )
+ };
+ if arch::any(s0, s1, s2, s3) {
+ unsafe {
+ cursor = decode(arch::bits(s0), i as u32, cursor, end);
+ cursor = decode(arch::bits(s1), i as u32 + 16, cursor, end);
+ cursor = decode(arch::bits(s2), i as u32 + 32, cursor, end);
+ cursor = decode(arch::bits(s3), i as u32 + 48, cursor, end);
+ }
+ }
+ i += 64;
+ }
+
+ cursor = unsafe { scan_tail::(pattern, &pattern_chunks, text, i, cursor, end) };
+ // SAFETY: `cursor` and `start` both point into `out`'s buffer.
+ unsafe { cursor.offset_from(start) as usize }
+}
+
+/// The remainder every scan front end shares, once its 64-byte rounds got to `i`: 16-byte
+/// rounds, then one overlapped final block, or a byte-by-byte walk when the whole text is
+/// shorter than a block. Returns the advanced cursor.
+///
+/// # Safety
+/// Same contract as [`scan`]: reads are bounded by the loop conditions, and `cursor .. end`
+/// must have room for every remaining match plus `decode`'s four-slot slack.
+unsafe fn scan_tail(
+ pattern: [u8; K],
+ pattern_chunks: &[arch::Chunk; K],
+ text: &[u8],
+ mut i: usize,
+ mut cursor: *mut u32,
+ end: *const u32,
+) -> *mut u32 {
+ let len = text.len();
+ let p = text.as_ptr();
+
+ // 16 bytes per round over what the 64-byte rounds left.
+ // SAFETY: reads end at `i + 16 + (K - 1) <= len`; writes as in [`scan`].
+ while i + 16 + (K - 1) <= len {
+ cursor = unsafe {
+ decode(
+ arch::bits(arch::match_starts::(p.add(i), pattern_chunks)),
+ i as u32,
+ cursor,
+ end,
+ )
+ };
+ i += 16;
+ }
+
+ if i + K <= len {
+ if len >= 16 + (K - 1) {
+ // One last 16-byte block, moved back to end flush with the last position a match
+ // can still start at. Its low bits re-cover positions the rounds above already
+ // decoded, so they are masked off.
+ // SAFETY: the block reads `base .. base + 16 + K - 1`, which ends exactly at `len`;
+ // writes as in [`scan`].
+ let base = len - 16 - (K - 1);
+ unsafe {
+ let mut bits = arch::bits(arch::match_starts::(p.add(base), pattern_chunks));
+ bits &= !0u64 << ((i - base) << arch::SHIFT);
+ cursor = decode(bits, base as u32, cursor, end);
+ }
+ } else {
+ // The whole text is shorter than one block: byte-by-byte costs nothing here.
+ while i + K <= len {
+ if text[i..i + K] == pattern[..] {
+ debug_assert!(end as usize - cursor as usize >= size_of::());
+ // SAFETY: one slot per match; the buffer holds a slot for every possible
+ // match plus four ([`matches_into`]'s assert).
+ unsafe {
+ *cursor = i as u32;
+ cursor = cursor.add(1);
+ }
+ }
+ i += 1;
+ }
+ }
+ }
+ cursor
+}
diff --git a/tokenizers/atomsplit/tests/fsm.rs b/tokenizers/atomsplit/tests/fsm.rs
index 1cad95e1b..3aa75ab7f 100644
--- a/tokenizers/atomsplit/tests/fsm.rs
+++ b/tokenizers/atomsplit/tests/fsm.rs
@@ -1,10 +1,48 @@
//! Integration tests for the FSM pre-tokenizers. Kept out of `src/` so the core stays production-only.
use atomsplit::classify::{classify, mask};
use atomsplit::fsm::{
- CharDelimiterSplit, Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k,
- fsm_deepseek, fsm_o200k, fsm_tekken,
+ Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k,
+ fsm_tekken, scan_byte_level, scan_byte_level_masked, scan_cl100k_cap, scan_cl100k_cap_masked,
+ scan_deepseek, scan_deepseek_masked, scan_o200k, scan_o200k_masked, scan_tekken,
+ scan_tekken_masked,
};
+/// The shared sweep for a masked/scalar scanner pair: compare spans on the corpus behind
+/// every 64-byte-edge offset (leading padding 0..=70), and with truncations exercising the
+/// scalar tail at every remaining length.
+type ScanFn = dyn Fn(&[u8], &[u8], &mut dyn FnMut(Span));
+
+fn masked_matches_scalar(corpus: &str, scalar: &ScanFn, masked: &ScanFn) {
+ fn spans_of(scan: &ScanFn, s: &[u8]) -> Vec {
+ let mut tags = vec![0u8; s.len()];
+ classify(s, &mut tags);
+ let mut v = Vec::new();
+ scan(s, &tags, &mut |sp| v.push(sp));
+ v
+ }
+ let check = |s: &str| {
+ let a = spans_of(scalar, s.as_bytes());
+ let b = spans_of(masked, s.as_bytes());
+ assert_eq!(b, a, "input len {}: {:?}", s.len(), s);
+ };
+ for pad in 0..=70 {
+ check(&format!("{}{}", "x".repeat(pad), corpus));
+ }
+ for pad in [0, 37] {
+ let padded = format!("{}{}", "x".repeat(pad), corpus);
+ for len in padded.len().saturating_sub(140)..=padded.len() {
+ if padded.is_char_boundary(len) {
+ check(&padded[..len]);
+ }
+ }
+ }
+ for len in 0..=70 {
+ if corpus.is_char_boundary(len) {
+ check(&corpus[..len]);
+ }
+ }
+}
+
/// Run a no-push fsm into a fresh buffer and return the emitted spans.
fn spans(f: impl Fn(&[u8], &[u8], &mut [Span]) -> usize, s: &str) -> Vec {
let mut tags = vec![0u8; s.len()];
@@ -48,6 +86,111 @@ fn deepseek_rules() {
assert_eq!(ds("!!!"), vec![(0, 3)]); // \p{P}∪\p{S} run
}
+/// Byte-exactness gate for the masked byte-level scanner: `scan_byte_level_masked` must emit the
+/// spans `scan_byte_level` emits, on every input. The corpus stresses every rule shape the batch
+/// algebra rewrites: contractions in both cases and at rejection edges, prefix spaces before all
+/// three run classes, unbounded numbers including non-ASCII digits, multi-byte whitespace (the
+/// bad-zone route), CJK/emoji/ZWJ runs, CRLF and tab runs, and runs longer than a batch. The
+/// padding sweep (0..=70 leading bytes) moves every shape across every 64-byte batch edge, so
+/// batch-edge carries, the bit-63 lookahead and every bad-zone route are hit at every offset; the
+/// truncation sweep exercises the scalar tail at every remaining length.
+#[test]
+fn masked_scan_matches_scan_byte_level() {
+ let corpus = concat!(
+ "I'm 12345 ok, don't they'll 've 'lx x's ''s IT'S 's mid'dle end' ",
+ "hello world\r\nmore\ttabs\t\t spaced end ",
+ "no.1 中文漢字テスト مرحبا १२३४५६७८९० ¹²³ ½¾ ",
+ "emoji 😀😀 zwj 👩\u{200d}🔬 nbsp\u{a0}x thin\u{2009}y wide\u{3000}z ",
+ "((()))!!!??? #$%&' “curly” apostrophe’d ",
+ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa ",
+ "9999999999999999999999999999999999999999999999999999999999999999999999999999999 ",
+ "中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中 end",
+ );
+ masked_matches_scalar(
+ corpus,
+ &|t, tg, e| scan_byte_level(t, tg, e),
+ &|t, tg, e| scan_byte_level_masked(t, tg, e),
+ );
+}
+
+/// Byte-exactness gate for the masked cl100k-family scanner, at every shipped digit cap. On top
+/// of the byte-level shapes: capped digit runs (pure ASCII, pure Devanagari, and mixed, so the
+/// char-counted bad route is hit mid-run), letters after punct at run starts vs mid-run (the
+/// two-chars-back absorb test), newlines absorbed after punct runs (`[\r\n]*`), whitespace runs
+/// with interior newlines (`\s*[\r\n]`), tab prefixes, and an apostrophe before a non-ASCII
+/// char (the `'ſ` defer).
+#[test]
+fn masked_scan_matches_scan_cl100k_cap() {
+ let corpus = concat!(
+ "I'm 12345 ok, don't they'LL 've 'lx x's ''s IT'S 's 3'ts end' ",
+ "a1234 12٣45 ١٢٣٤٥٦ १२३४५६७८९० 999999999999999999999999999999999999999999999999 ",
+ "!!a x!a ?!x. ;\n\n};\r\n\r\n foo();\nbar() //x\n\n\n end\n",
+ "hello world\r\nmore\ttabs\t\t spaced \ta \r\n \n\t\r\n\t x ",
+ "no.1 中文漢字テスト مرحبا nbsp\u{a0}x thin\u{2009}y wide\u{3000}z 'ſok “curly”’d ",
+ "((()))!!!??? #$%&' 😀😀 👩\u{200d}🔬 ",
+ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa ",
+ "中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中 end",
+ );
+ for cap in [3, 1, usize::MAX] {
+ masked_matches_scalar(
+ corpus,
+ &move |t, tg, e| scan_cl100k_cap(t, tg, cap, e),
+ &move |t, tg, e| scan_cl100k_cap_masked(t, tg, cap, e),
+ );
+ }
+}
+
+/// Byte-exactness gate for the masked o200k/tekken scanner. On top of the cl100k shapes: case
+/// splits (camelCase, all-upper runs, upper after caseless — the deferred backtrack), suffix
+/// contractions incl. chains ("can'ts", "x'll'd") and prefix apostrophes after digits, `[\r\n/]*`
+/// tails with slash runs before and after newlines (the walkback shapes), and combining marks
+/// (run-contextual class, the wide bad smear).
+#[test]
+fn masked_scan_matches_scan_o200k_and_tekken() {
+ let corpus = concat!(
+ "camelCase HTTPResponse XMLHttpRequest AAAA aaaa aA Aa 中B B中b ʰupper Xʰa 中中中中 ",
+ "don't they'LL CAN'TS x'll'd 3'ts I'm'll a'sit 's'' end' 'ſok ",
+ "a1234 12٣45 १२३४५६ 99999999999999999999999999999999999999999999999999999999999999999 ",
+ "foo();\nbar() .\n// ///x\r\n/// a/b//c\n\n//d e\u{301}f g\u{5bf}h zwj\u{200c}x ",
+ "hello world\r\nmore\ttabs\t\t spaced \ta \r\n \n\t\r\n\t x nbsp\u{a0}x wide\u{3000}z ",
+ "((()))!!!??? #$%&' “curly”’d 😀😀 مرحبا מבחן ",
+ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa end"
+ );
+ masked_matches_scalar(
+ corpus,
+ &|t, tg, e| scan_o200k(t, tg, e),
+ &|t, tg, e| scan_o200k_masked(t, tg, e),
+ );
+ masked_matches_scalar(
+ corpus,
+ &|t, tg, e| scan_tekken(t, tg, e),
+ &|t, tg, e| scan_tekken_masked(t, tg, e),
+ );
+}
+
+/// Byte-exactness gate for the masked deepseek scanner. On top of the shared shapes: CJK
+/// letter/punct runs and their neighborhoods (the closed-unit rule and the bad cover), gap runs
+/// (controls / NumericOther / ZWJ) with and without a following letter run (the prefix split),
+/// alt-1 `[ascii-punct][A-Za-z]+` including collisions with non-ASCII letters ("_naïve"), and
+/// whitespace runs followed by digits or CJK (no give-back).
+#[test]
+fn masked_scan_matches_scan_deepseek() {
+ let corpus = concat!(
+ "abc中def 中文漢字テスト!ひらがな・カタカナ 中中中中中中中 mixed中123中ok 拼音列表(e.g. 表!x ",
+ "_abc (foo) .py x!a _naïve _né !!a a/b.c 3'ts don't ½x ¼¼y\u{7f}z zwj\u{200c}gap\u{200c}\u{200c}word ",
+ "a1234 12٣45 999999999999999999999999999999999999999999999999999999999999999999 12 中 ",
+ "x 123 y\t\t\t45 z \n 99 w 中 Model):\n money = f(x):\n\n\t indent ",
+ "hello world\r\nmore\ttabs\t\t spaced \ta \r\n \n\t\r\n\t x wide\u{3000}z 中 12\n\n34 ",
+ "((()))!!!??? #$%&' “curly”’d 😀😀 مرحبا ",
+ "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa end"
+ );
+ masked_matches_scalar(
+ corpus,
+ &|t, tg, e| scan_deepseek(t, tg, e),
+ &|t, tg, e| scan_deepseek_masked(t, tg, e),
+ );
+}
+
#[test]
fn byte_level_rules() {
let bl = |s| spans(fsm_byte_level, s);
@@ -57,14 +200,6 @@ fn byte_level_rules() {
assert_eq!(bl("hi ok"), vec![(0, 2), (2, 4), (4, 7)]); // \s+(?!\S) leaves one space
}
-#[test]
-fn char_delimiter_split() {
- let mut out = vec![Span::default(); 8];
- // split on '/', Removed → drop delimiters, drop the empty gap between "//"
- let k = CharDelimiterSplit('/').pre_tokenize(b"a/bc//d", &mut [], &mut out);
- assert_eq!(&out[..k], &[(0, 1), (2, 4), (6, 7)]);
-}
-
/// Byte-exactness gate for the class family: the NEON boundary extractor (`class_runs_into`) must equal
/// the scalar run-end core (`emit_class_spans`) for every recipe, at every char-aligned truncation length so
/// the < 16-byte NEON tail starts at every offset — including mid-char (chunk loop steps by 16). Corpus
diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs
index 14ceed9a0..894dde46a 100644
--- a/tokenizers/atomsplit/tests/literal.rs
+++ b/tokenizers/atomsplit/tests/literal.rs
@@ -48,3 +48,203 @@ fn an_empty_pattern_is_rejected() {
"an empty pattern matches everywhere"
);
}
+
+// ---- the batch scan, `matches_into` ----
+
+/// Run `matches_into` with a buffer of exactly the documented size and return what it wrote.
+fn batch(literal: &Literal, text: &[u8]) -> Vec {
+ let mut out = vec![0u32; text.len() / literal.pattern().len() + 4];
+ let count = literal.matches_into(text, &mut out);
+ out.truncate(count);
+ out
+}
+
+fn iterated(literal: &Literal, text: &[u8]) -> Vec {
+ literal.matches(text).map(|p| p as u32).collect()
+}
+
+#[test]
+fn matches_into_reports_the_iterator_positions() {
+ let dash = Literal::new(b"-").unwrap();
+ assert_eq!(batch(&dash, b"a-b--c"), [1, 3, 4]);
+ assert_eq!(batch(&dash, b"-starts and ends-"), [0, 16]);
+ assert_eq!(batch(&dash, b"none here"), []);
+ assert_eq!(batch(&dash, b""), []);
+
+ let space = Literal::new(b" ").unwrap();
+ let text = b"the quick brown fox jumps over the lazy dog".repeat(40);
+ assert_eq!(batch(&space, &text), iterated(&space, &text));
+
+ let metaspace = Literal::new("\u{2581}".as_bytes()).unwrap();
+ let text = "\u{2581}the\u{2581}quick\u{2581}brown\u{2581}fox".repeat(40);
+ assert_eq!(
+ batch(&metaspace, text.as_bytes()),
+ iterated(&metaspace, text.as_bytes())
+ );
+}
+
+/// The scan works in blocks; a match starting on either side of an internal block edge, or on
+/// the last position where the pattern still fits, must not be missed or doubled.
+#[test]
+fn matches_into_is_seamless_across_block_boundaries() {
+ for pattern in [&b"-"[..], "\u{2581}".as_bytes()] {
+ let literal = Literal::new(pattern).unwrap();
+ for start in [0, 13, 14, 15, 16, 17, 61, 62, 63, 64, 65, 127, 128] {
+ for len in [start + pattern.len(), 90, 129, 200] {
+ if start + pattern.len() > len {
+ continue;
+ }
+ let mut text = vec![b'x'; len];
+ text[start..start + pattern.len()].copy_from_slice(pattern);
+ assert_eq!(
+ batch(&literal, &text),
+ [start as u32],
+ "pattern {pattern:?} at {start} in {len} bytes"
+ );
+ }
+ }
+ }
+}
+
+/// Every position matches: the largest count a text can produce, written into a buffer of
+/// exactly the documented size.
+#[test]
+fn matches_into_handles_the_densest_text() {
+ let a = Literal::new(b"a").unwrap();
+ let text = vec![b'a'; 300];
+ assert_eq!(batch(&a, &text), (0..300).collect::>());
+
+ let metaspace = Literal::new("\u{2581}".as_bytes()).unwrap();
+ let text = "\u{2581}".repeat(100);
+ assert_eq!(
+ batch(&metaspace, text.as_bytes()),
+ (0..100).map(|i| i * 3).collect::>()
+ );
+}
+
+#[test]
+fn a_self_overlapping_pattern_keeps_non_overlapping_matches() {
+ let aa = Literal::new(b"aa").unwrap();
+ assert_eq!(batch(&aa, b"aaa"), [0]);
+ let run = [b'a'; 100];
+ assert_eq!(batch(&aa, &run), iterated(&aa, &run));
+
+ let aba = Literal::new(b"aba").unwrap();
+ assert_eq!(batch(&aba, b"ababa"), [0]);
+}
+
+#[test]
+fn a_pattern_longer_than_three_bytes_still_matches() {
+ let literal = Literal::new(b"abcd").unwrap();
+ let text = b"abcd xabcdx abcabcd".repeat(20);
+ assert_eq!(batch(&literal, &text), iterated(&literal, &text));
+}
+
+#[test]
+#[should_panic(expected = "matches_into")]
+fn an_undersized_buffer_is_rejected() {
+ let a = Literal::new(b"a").unwrap();
+ let mut out = vec![0u32; 7]; // "aaaa" needs 4 / 1 + 4 = 8
+ a.matches_into(b"aaaa", &mut out);
+}
+
+// ---- the streaming pair, `count_matches` + `for_each_match` ----
+
+fn streamed(literal: &Literal, text: &[u8]) -> Vec {
+ let mut out = Vec::new();
+ literal.for_each_match(text, |start| out.push(start as u32));
+ out
+}
+
+/// The streaming scan works through a fixed window; a text much longer than any plausible
+/// window must report the same matches as the iterator, including around every window edge.
+#[test]
+fn for_each_match_streams_the_iterator_offsets_across_windows() {
+ for pattern in [&b" "[..], "\u{2581}".as_bytes(), b"ab"] {
+ let literal = Literal::new(pattern).unwrap();
+ let mut text = b"word ".repeat(4_000); // ~20KB, matches every 5 bytes
+ text.extend_from_slice("\u{2581}ab".repeat(2_000).as_bytes());
+ assert_eq!(streamed(&literal, &text), iterated(&literal, &text));
+ assert_eq!(
+ literal.count_matches(&text),
+ literal.matches(&text).count(),
+ "count for {pattern:?}"
+ );
+ }
+}
+
+/// The densest possible long text: every byte starts a match, in every window.
+#[test]
+fn for_each_match_handles_dense_windows() {
+ let a = Literal::new(b"a").unwrap();
+ let text = vec![b'a'; 10_000];
+ assert_eq!(streamed(&a, &text), (0..10_000).collect::>());
+ assert_eq!(a.count_matches(&text), 10_000);
+}
+
+#[test]
+fn for_each_match_takes_the_iterator_route_for_uncovered_patterns() {
+ // Self-overlapping and longer-than-three patterns are exactly the ones the batch scan
+ // refuses; the streaming pair must still report the iterator's non-overlapping matches.
+ for pattern in [&b"aa"[..], b"aba", b"abcd"] {
+ let literal = Literal::new(pattern).unwrap();
+ let text = b"aabaabcdaaaba".repeat(2_000);
+ assert_eq!(streamed(&literal, &text), iterated(&literal, &text));
+ assert_eq!(literal.count_matches(&text), literal.matches(&text).count());
+ }
+}
+
+#[test]
+fn streaming_agrees_with_the_iterator_on_random_lengths() {
+ let mut state = 0x1234_5678_9ABC_DEF0u64;
+ let mut next = move || {
+ state ^= state << 13;
+ state ^= state >> 7;
+ state ^= state << 17;
+ state
+ };
+ let patterns: &[&[u8]] = &[b"a", b"ab", b"aa", b"aba", b"abc", b"abca"];
+ for _ in 0..500 {
+ // lengths spread around plausible window sizes, so edges get hit both exactly and off by one
+ let len = (next() % 8_000) as usize;
+ let text: Vec = (0..len).map(|_| b'a' + (next() % 3) as u8).collect();
+ let literal = Literal::new(patterns[(next() % patterns.len() as u64) as usize]).unwrap();
+ assert_eq!(
+ streamed(&literal, &text),
+ iterated(&literal, &text),
+ "pattern {:?} len {len}",
+ literal.pattern()
+ );
+ assert_eq!(literal.count_matches(&text), literal.matches(&text).count());
+ }
+}
+
+#[test]
+fn matches_into_agrees_with_the_iterator_on_random_inputs() {
+ // xorshift64: deterministic, no dev-dependency needed.
+ let mut state = 0x9E37_79B9_7F4A_7C15u64;
+ let mut next = move || {
+ state ^= state << 13;
+ state ^= state >> 7;
+ state ^= state << 17;
+ state
+ };
+ let patterns: &[&[u8]] = &[
+ b"a", b"b", b"ab", b"aa", b"ba", b"aab", b"aba", b"abc", b"abca",
+ ];
+ for _ in 0..20_000 {
+ let len = (next() % 300) as usize;
+ let alphabet = 2 + (next() % 2) as u8;
+ let text: Vec = (0..len)
+ .map(|_| b'a' + (next() % alphabet as u64) as u8)
+ .collect();
+ let literal = Literal::new(patterns[(next() % patterns.len() as u64) as usize]).unwrap();
+ assert_eq!(
+ batch(&literal, &text),
+ iterated(&literal, &text),
+ "pattern {:?} text {:?}",
+ literal.pattern(),
+ text
+ );
+ }
+}
diff --git a/tokenizers/tk-encode/examples/fixture_bench.rs b/tokenizers/tk-encode/examples/fixture_bench.rs
index 9ac6d89bf..947e566e5 100644
--- a/tokenizers/tk-encode/examples/fixture_bench.rs
+++ b/tokenizers/tk-encode/examples/fixture_bench.rs
@@ -20,7 +20,10 @@
//! on, so the headline includes the post-process stage the ladder charges.
//! 2. **Stage breakdown** — the `encode_generic::` ablation ladder plus
//! the pre-tokenize-vs-regex-engine references, on fresh caller-owned
-//! scratches, fully separate from the phase-1 timings.
+//! scratches, fully separate from the phase-1 timings. Every rung drives the
+//! same route a full `encode` takes (fused scan, zero-copy or staged, reported
+//! as `route`), so the split bar is the splitter that actually ships and the
+//! model bar is a difference between two runs of it. See `stage_secs`.
//! 3. **Scaling & memory** — a multi-thread throughput sweep (1/2/4/8/max) over
//! the whole corpus on fresh instances, and resident-set deltas measured by
//! re-spawning this binary as `--memory ` children — one
@@ -299,6 +302,23 @@ fn bench_throughput(
/// successive levels and subtracting gives each stage's marginal cost (the ablation
/// ladder), no profiler and no per-segment instrumentation.
///
+/// Every rung runs the **route a full encode runs**, which is what makes subtracting
+/// them meaningful. A fused pipeline's split rung drives the same masked FSM scan the
+/// model rung does, emitting each span into `pre_tokens` instead of the model; a
+/// zero-copy pipeline's split rung is its proven-cut pass, with neither the rewrite
+/// write nor the added-token scan over rewritten text that the route never performs.
+/// Two consequences to read the numbers with:
+///
+/// - The split rung pays one span write per pre-token that a full fused encode does not
+/// (it hands the span straight to the model), so the split bar reads a little high and
+/// the model bar a little low.
+/// - The `&str` slice the model call makes is charged to the model, not to the split. On
+/// these corpora that is UTF-8 boundary checking, and it is not a small share.
+///
+/// Before this was route-aware the split rung ran the *staged* array FSM, a splitter no
+/// shipping encode uses. That put the split bar about twice too high and made the model
+/// bar a subtraction across two different pipelines.
+///
/// The scratch is created fresh here (never taken from the pipeline's pool), so the
/// stage numbers are warmed on this fixture alone and can't perturb — or be flattered
/// by — the phase-1 cache state. Both caller-owned buffers are reused across chunks
@@ -333,6 +353,46 @@ fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String])
median_secs(samples)
}
+/// The top of the ladder with the output and span buffers allocated **per chunk**, the way
+/// `PipelineTokenizer::encode` does for every call. The pipeline, the fresh scratch and
+/// the warm-up are all identical to `stage_secs`, so subtracting `t_post` isolates
+/// allocation with no other variable in play.
+///
+/// It comes out near zero on these corpora: growing and freeing the output `Vec` is not a
+/// cost worth chasing. It is reported anyway, because it looks like one.
+///
+/// Timing the real `encode` here instead was tried and dropped. Phase 2 shares one pipeline
+/// across every fixture, so its scratch pool holds the words of all previously benched
+/// corpora, and against 64k slots that saturates. The delta then came out larger than the
+/// whole encode: it measures cache eviction, not call overhead. A number that size in the
+/// report would have read as a per-call cost.
+fn fresh_buffer_secs(pipeline: &PipelineTokenizer, chunks: &[String]) -> f64 {
+ let mut scratch = pipeline.get_model().init_scratch();
+ let mut run = || {
+ for chunk in chunks {
+ let mut out = Vec::new();
+ let mut pre_tokens = Vec::new();
+ let _ = pipeline.encode_generic::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(
+ chunk,
+ true,
+ &mut pre_tokens,
+ &mut scratch,
+ &mut out,
+ );
+ black_box(&out);
+ black_box(&pre_tokens);
+ }
+ };
+ run(); // warm-up
+ let mut samples = Vec::with_capacity(REPS);
+ for _ in 0..REPS {
+ let start = Instant::now();
+ run();
+ samples.push(start.elapsed().as_secs_f64());
+ }
+ median_secs(samples)
+}
+
/// Stage decomposition + regex-engine references for one fixture: the
/// `stage_ns_per_byte` and `pretok_vs_regex` objects of its report row.
fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -> (Value, Value) {
@@ -341,10 +401,23 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -
let t_split = stage_secs::<{ PipelineTokenizer::STAGE_SPLIT }>(pipeline, &f.chunks);
let t_model = stage_secs::<{ PipelineTokenizer::STAGE_MODEL }>(pipeline, &f.chunks);
let t_post = stage_secs::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(pipeline, &f.chunks);
+ let t_fresh = fresh_buffer_secs(pipeline, &f.chunks);
// Two distinct "split" costs: `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, `post` the special-token id-frame
- // splice. All five stages sum exactly to `total`.
+ // `pre_tokenize` is the pre-tokenizer split *as this model's route performs it*,
+ // `post` the special-token id-frame splice. The five stages sum to `total`, which is
+ // the ladder's own full encode.
+ //
+ // `clamped` is what the five do NOT account for: the rungs are independent medians, so
+ // a stage at or below the noise floor goes negative and clamps to zero, after which the
+ // parts sum slightly above `total`. A `clamped` that is not near zero means this
+ // fixture's rungs are noise-dominated and its breakdown should not be trusted.
+ //
+ // `alloc` is the one cost outside the decomposition that can be isolated cleanly; see
+ // `fresh_buffer_secs`. `total` is still below what phase 1 measures through the public
+ // `encode`, which also acquires a pooled scratch and hands back an owned `Vec`; that
+ // remainder is NOT reported here, because no measurement in this bench separates it
+ // from the pool's cache state. Better an acknowledged gap than a mislabelled bar.
let nspb = |secs: f64| secs * 1e9 / f.bytes as f64;
let (ns_added, ns_norm, ns_split, ns_model, ns_post) = (
nspb(t_frame.max(0.0)),
@@ -353,15 +426,23 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -
nspb((t_model - t_split).max(0.0)),
nspb((t_post - t_model).max(0.0)),
);
+ let ns_total = nspb(t_post);
+ let clamped = ns_total - (ns_added + ns_norm + ns_split + ns_model + ns_post);
+ let ns_alloc = nspb(t_fresh - t_post);
eprintln!(
- " {} stages ns/byte: added-split {ns_added:.2}, norm {ns_norm:.2}, pre-split {ns_split:.2}, model {ns_model:.2}, post {ns_post:.2}",
+ " {} stages ns/byte: added-split {ns_added:.2}, norm {ns_norm:.2}, pre-split {ns_split:.2}, model {ns_model:.2}, post {ns_post:.2} → total {ns_total:.2} (clamped {clamped:+.2}) · buffer alloc {ns_alloc:+.2}",
f.name
);
- // pre_tokenize (= classify SIMD + fsm) vs classify-scalar and vs real regex engines
- // over the same corpus, so the report shows the split beating a regex engine both
- // WITH and WITHOUT SIMD. `scalar_pipe` = pre_tokenize + (cls_scalar − cls_simd):
- // fsm is the scalar jump-table in both pipes, SIMD/scalar is the classify pass only.
+ // The shipped split vs real regex engines over the same corpus. `ns_split` is now the
+ // route's own splitter, so these ratios describe the code that runs. Before the ladder
+ // was route-aware they compared the engines against the staged array FSM and undersold
+ // the split by roughly 2x.
+ //
+ // `scalar_pipe` = ns_split + (cls_scalar − cls_simd) swaps *only* the classify pass
+ // for its scalar version. It is not "the split without SIMD": for a regex-shaped FSM
+ // the scan is a boundary-mask scanner, so SIMD lives in the scan too and cannot be
+ // subtracted out this way. Read it as the split with a scalar classify pass.
let corpus: String = f.chunks.concat();
let cls_simd = classify_ns(corpus.as_bytes(), false);
let cls_scalar = classify_ns(corpus.as_bytes(), true);
@@ -384,7 +465,7 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -
})
};
eprintln!(
- " {} pre-tok: SIMD-cls {ns_split:.2} / scalar-cls {scalar_pipe:.2} ns/B · vs onig {} · vs fancy {} · vs pcre2 {} · vs logos {}",
+ " {} split: {ns_split:.2} / scalar-classify {scalar_pipe:.2} ns/B · vs onig {} · vs fancy {} · vs pcre2 {} · vs logos {}",
f.name,
vs(onig_ns),
vs(fancy_ns),
@@ -400,7 +481,9 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -
"pre_tokenize": ns_split,
"model": ns_model,
"post": ns_post,
- "total": nspb(t_post),
+ "total": ns_total,
+ "clamped": clamped,
+ "alloc": ns_alloc,
}),
json!({
"cls_simd": cls_simd,
@@ -1188,6 +1271,9 @@ fn main() {
models.push(json!({
"model": name, "desc": desc, "shape": shape,
+ // Which encode route the stage numbers decompose: the three split the work
+ // differently, so a stage chart is only readable against the route it timed.
+ "route": pipeline.encode_route().as_str(),
"results": rows, "memory": memory, "threads": threads,
"decode_threads": decode_threads, "decode_reason": decode_reason,
}));
diff --git a/tokenizers/tk-encode/examples/fold_ab.rs b/tokenizers/tk-encode/examples/fold_ab.rs
new file mode 100644
index 000000000..65e57aa27
--- /dev/null
+++ b/tokenizers/tk-encode/examples/fold_ab.rs
@@ -0,0 +1,80 @@
+//! One end-to-end throughput number for one model on one corpus, for A/B runs
+//! between two builds of this crate (fold on vs fold off, run alternately at
+//! process level so binary-layout noise averages out).
+//!
+//! Mirrors `fixture_bench`'s throughput phase: single thread, ~10 kB chunks,
+//! `add_special_tokens` on, one warm-up pass to fill the caches, then the
+//! median of the timed passes.
+//!
+//! cargo run --release --example fold_ab -- [passes]
+
+use std::convert::TryFrom;
+use std::time::Instant;
+
+use tk_encode::Tokenizer;
+use tk_encode::pipeline::PipelineTokenizer;
+
+const CHUNK: usize = 10_000;
+const CORPUS_CAP: usize = 4 << 20;
+
+fn main() {
+ let args: Vec = std::env::args().collect();
+ let [_, model, corpus] = &args[..3] else {
+ eprintln!("usage: fold_ab [passes]");
+ std::process::exit(2);
+ };
+ let passes: usize = args.get(3).map_or(5, |p| p.parse().unwrap());
+
+ let tok = Tokenizer::from_file(model).unwrap();
+ let pipeline = PipelineTokenizer::try_from(&tok).unwrap();
+
+ let mut text = std::fs::read_to_string(corpus).unwrap();
+ if text.len() > CORPUS_CAP {
+ let mut cap = CORPUS_CAP;
+ while !text.is_char_boundary(cap) {
+ cap -= 1;
+ }
+ text.truncate(cap);
+ }
+ let chunks: Vec<&str> = {
+ let mut chunks = Vec::new();
+ let mut rest = text.as_str();
+ while !rest.is_empty() {
+ let mut cut = CHUNK.min(rest.len());
+ while !rest.is_char_boundary(cut) {
+ cut -= 1;
+ }
+ let (head, tail) = rest.split_at(cut);
+ chunks.push(head);
+ rest = tail;
+ }
+ chunks
+ };
+ let bytes: usize = chunks.iter().map(|c| c.len()).sum();
+
+ let encode_all = || {
+ let mut ids = 0usize;
+ for chunk in &chunks {
+ ids += pipeline.encode(chunk, true).unwrap().len();
+ }
+ ids
+ };
+
+ let ids = encode_all(); // warm-up: fills the scratch pool and the word cache
+ let mut timings: Vec = (0..passes)
+ .map(|_| {
+ let start = Instant::now();
+ encode_all();
+ start.elapsed().as_secs_f64()
+ })
+ .collect();
+ timings.sort_by(f64::total_cmp);
+ let median = timings[timings.len() / 2];
+
+ println!(
+ "{:.3} MB/s ({} B, {} ids, {passes} passes)",
+ bytes as f64 / median / 1e6,
+ bytes,
+ ids,
+ );
+}
diff --git a/tokenizers/tk-encode/examples/normalize_ab.rs b/tokenizers/tk-encode/examples/normalize_ab.rs
new file mode 100644
index 000000000..2c8a8304a
--- /dev/null
+++ b/tokenizers/tk-encode/examples/normalize_ab.rs
@@ -0,0 +1,98 @@
+//! A/B harness for changes to the pipeline's normalize stage: for each tokenizer config given on
+//! the command line, times the `encode_generic` ablation ladder over `data/big.txt` and prints the
+//! frame and normalize levels plus the full encode throughput.
+//!
+//! cargo run --release -p tk-encode --example normalize_ab -- data/llama-2.json data/gemma-4.json
+//!
+//! Binary layout alone moves numbers by a few percent, so never compare two builds from one
+//! process each: build one binary per side, keep both, and alternate runs.
+
+use std::convert::TryFrom;
+use std::hint::black_box;
+use std::path::Path;
+use std::time::Instant;
+
+use tk_encode::Tokenizer;
+use tk_encode::pipeline::{Model, PipelineTokenizer};
+
+const DATA_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data");
+/// Per-input-overhead is amortized at this size; same regime as `fixture_bench`.
+const CHUNK_BYTES: usize = 10 * 1024;
+const TOTAL_BYTES: usize = 4 * 1024 * 1024;
+const REPS: usize = 9;
+
+fn median(mut samples: Vec) -> f64 {
+ samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
+ samples[samples.len() / 2]
+}
+
+/// Median seconds of one warm pass over `chunks` at ladder level `STAGE`;
+/// same shape as `fixture_bench::stage_secs`.
+fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[&str]) -> f64 {
+ let mut out = Vec::new();
+ let mut pre_tokens = Vec::new();
+ let mut scratch = pipeline.get_model().init_scratch();
+ let mut run = || {
+ for chunk in chunks {
+ out.clear();
+ let _ = pipeline.encode_generic::(
+ chunk,
+ true,
+ &mut pre_tokens,
+ &mut scratch,
+ &mut out,
+ );
+ black_box(&out);
+ black_box(&pre_tokens);
+ }
+ };
+ run(); // warm-up
+ let mut samples = Vec::with_capacity(REPS);
+ for _ in 0..REPS {
+ let start = Instant::now();
+ run();
+ samples.push(start.elapsed().as_secs_f64());
+ }
+ median(samples)
+}
+
+/// The first `TOTAL_BYTES` of big.txt in `CHUNK_BYTES` pieces, cut on char boundaries.
+fn chunks_of(text: &str) -> Vec<&str> {
+ let mut chunks = Vec::new();
+ let mut start = 0;
+ while start < text.len().min(TOTAL_BYTES) {
+ let mut end = (start + CHUNK_BYTES).min(text.len());
+ while !text.is_char_boundary(end) {
+ end += 1;
+ }
+ chunks.push(&text[start..end]);
+ start = end;
+ }
+ chunks
+}
+
+fn main() {
+ let text = std::fs::read_to_string(Path::new(DATA_DIR).join("big.txt"))
+ .expect("data/big.txt (fetch with `make data/big.txt`)");
+ let chunks = chunks_of(&text);
+ let bytes: usize = chunks.iter().map(|c| c.len()).sum();
+
+ for config in std::env::args().skip(1) {
+ let tok = Tokenizer::from_file(&config).expect("tokenizer config");
+ let pipeline = PipelineTokenizer::try_from(&tok).expect("pipeline builds");
+
+ let t_frame = stage_secs::<{ PipelineTokenizer::STAGE_FRAME }>(&pipeline, &chunks);
+ let t_norm = stage_secs::<{ PipelineTokenizer::STAGE_NORMALIZE }>(&pipeline, &chunks);
+ let t_full = stage_secs::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(&pipeline, &chunks);
+
+ let nspb = |secs: f64| secs * 1e9 / bytes as f64;
+ println!(
+ "{config}: frame {:.3} ns/B, normalize {:.3} ns/B (marginal {:.3}), full {:.3} ns/B = {:.1} MB/s",
+ nspb(t_frame),
+ nspb(t_norm),
+ nspb(t_norm - t_frame),
+ nspb(t_full),
+ bytes as f64 / t_full / 1e6,
+ );
+ }
+}
diff --git a/tokenizers/tk-encode/examples/normalize_claims.rs b/tokenizers/tk-encode/examples/normalize_claims.rs
new file mode 100644
index 000000000..75dd97cf9
--- /dev/null
+++ b/tokenizers/tk-encode/examples/normalize_claims.rs
@@ -0,0 +1,337 @@
+//! What the space rewrite costs when written, and what the zero-copy path wins back.
+//!
+//! For each SentencePiece-shaped model this runs the `encode_generic::` ablation
+//! ladder (the same methodology as `fixture_bench`) over long and short inputs, then
+//! decomposes the fused normalizer's pass into its three parts: counting the spaces,
+//! allocating the rewrite `String`, and writing it. Models taking the zero-copy path
+//! (see `ZeroCopyMetaspace`) get an interleaved A/B against the written rewrite, with
+//! ids compared over the whole corpus first. A last probe measures what the second
+//! special-token scan adds once a tokenizer holds `normalized` added tokens.
+
+use std::hint::black_box;
+use std::path::Path;
+use std::time::Instant;
+
+use atomsplit::literal::Literal;
+use tk_encode::pipeline::{Model, PipelineTokenizer};
+use tk_encode::{AddedToken, NormalizerWrapper, Tokenizer};
+
+const DATA: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data");
+const REPS: usize = 9;
+const CHUNK_BYTES: usize = 10 * 1024;
+const MAX_BYTES: usize = 512 * 1024;
+
+/// (config file, prepend for the decomposition replica; `None` skips it because the
+/// model runs the `drop_whitespace` path, which is a different rewrite)
+const MODELS: &[(&str, Option)] = &[
+ ("llama-2.json", Some(true)),
+ ("gemma-4.json", Some(false)),
+ ("t5-base.json", None),
+ ("albert-base-v1-tokenizer.json", None),
+];
+
+const FIXTURES: &[&str] = &[
+ "fixtures/lang/eng_Latn.txt",
+ "fixtures/lang/cmn_Hani.txt",
+ "fixtures/modalities/code_mixed.txt",
+];
+
+/// The words `fixture_bench` injects as `normalized:true` added tokens, so the probe
+/// exercises the same second-scan state the comparative benchmark runs under.
+const NORMALIZED_WORDS: &[&str] = &["widgetron", "flibberjast", "zorptastic", "quibblenaut"];
+
+fn median(mut samples: Vec) -> f64 {
+ samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
+ samples[samples.len() / 2]
+}
+
+fn timed(mut run: impl FnMut()) -> f64 {
+ run(); // warm-up
+ let mut samples = Vec::with_capacity(REPS);
+ for _ in 0..REPS {
+ let start = Instant::now();
+ run();
+ samples.push(start.elapsed().as_secs_f64());
+ }
+ median(samples)
+}
+
+fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String]) -> f64 {
+ let mut out = Vec::new();
+ let mut pre_tokens = Vec::new();
+ let mut scratch = pipeline.get_model().init_scratch();
+ timed(|| {
+ for chunk in chunks {
+ out.clear();
+ let _ = pipeline.encode_generic::(
+ chunk,
+ true,
+ &mut pre_tokens,
+ &mut scratch,
+ &mut out,
+ );
+ black_box(&out);
+ black_box(&pre_tokens);
+ }
+ })
+}
+
+/// Whole lines accumulated up to `chunk_bytes`, `MAX_BYTES` in total. `chunk_bytes = 0`
+/// keeps each line its own chunk (the short-input regime).
+fn chunks_of(text: &str, chunk_bytes: usize) -> Vec {
+ let mut chunks = Vec::new();
+ let mut current = String::new();
+ let mut total = 0usize;
+ for line in text.lines().filter(|l| !l.trim().is_empty()) {
+ current.push_str(line);
+ if current.len() >= chunk_bytes {
+ total += current.len();
+ chunks.push(std::mem::take(&mut current));
+ if total >= MAX_BYTES {
+ return chunks;
+ }
+ } else {
+ current.push(' ');
+ }
+ }
+ if !current.is_empty() {
+ chunks.push(current);
+ }
+ chunks
+}
+
+struct Ladder {
+ added: f64,
+ norm: f64,
+ split: f64,
+ model: f64,
+ post: f64,
+ total_mbs: f64,
+}
+
+fn ladder(pipeline: &PipelineTokenizer, chunks: &[String], bytes: usize) -> Ladder {
+ 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);
+ let t_post = stage_secs::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(pipeline, chunks);
+ let nspb = |secs: f64| secs * 1e9 / bytes as f64;
+ Ladder {
+ added: nspb(t_frame.max(0.0)),
+ norm: nspb((t_norm - t_frame).max(0.0)),
+ split: nspb((t_split - t_norm).max(0.0)),
+ model: nspb((t_model - t_split).max(0.0)),
+ post: nspb((t_post - t_model).max(0.0)),
+ total_mbs: bytes as f64 / t_post / 1e6,
+ }
+}
+
+/// The fused normalizer's pass, split into its three costs over the same chunks:
+/// the space count, the `String` allocation, and the full rewrite (count + alloc +
+/// write). All three follow `MetaspaceNormalizer::normalize`'s non-`drop_whitespace`
+/// arm byte for byte, so `rewrite` should land on the ladder's norm marginal.
+fn decompose(chunks: &[String], bytes: usize, prepend: bool) -> (f64, f64, f64) {
+ let space = Literal::new(b" ").unwrap();
+ let delimiter = "\u{2581}";
+ let counts: Vec = chunks
+ .iter()
+ .map(|c| space.count_matches(c.as_bytes()))
+ .collect();
+ let nspb = |secs: f64| secs * 1e9 / bytes as f64;
+
+ let count_only = timed(|| {
+ for chunk in chunks {
+ black_box(space.count_matches(chunk.as_bytes()));
+ }
+ });
+ let alloc_only = timed(|| {
+ for (chunk, &count) in chunks.iter().zip(&counts) {
+ let s = String::with_capacity(chunk.len() + 2 * count + if prepend { 3 } else { 0 });
+ black_box(&s);
+ }
+ });
+ let rewrite = timed(|| {
+ for chunk in chunks {
+ let count = space.count_matches(chunk.as_bytes());
+ if !prepend && count == 0 {
+ black_box(chunk.as_str());
+ continue;
+ }
+ let mut rewritten =
+ String::with_capacity(chunk.len() + 2 * count + if prepend { 3 } else { 0 });
+ if prepend {
+ rewritten.push_str(delimiter);
+ }
+ let mut prev = 0;
+ space.for_each_match(chunk.as_bytes(), |start| {
+ rewritten.push_str(&chunk[prev..start]);
+ rewritten.push_str(delimiter);
+ prev = start + 1;
+ });
+ rewritten.push_str(&chunk[prev..]);
+ black_box(&rewritten);
+ }
+ });
+ (nspb(count_only), nspb(alloc_only), nspb(rewrite))
+}
+
+fn main() {
+ let fixtures: Vec<(String, String)> = FIXTURES
+ .iter()
+ .map(|rel| {
+ let path = Path::new(DATA).join(rel);
+ let name = path.file_stem().unwrap().to_str().unwrap().to_string();
+ (name, std::fs::read_to_string(&path).unwrap())
+ })
+ .collect();
+
+ for &(file, replica_prepend) in MODELS {
+ let path = Path::new(DATA).join(file);
+ let mut tok = match Tokenizer::from_file(&path) {
+ Ok(t) => t,
+ Err(e) => {
+ println!("== {file}: load failed: {e}");
+ continue;
+ }
+ };
+ let mut pipeline = match PipelineTokenizer::try_from(&tok) {
+ Ok(p) => p,
+ Err(e) => {
+ println!("== {file}: no pipeline: {e}");
+ continue;
+ }
+ };
+ // The ladder times the written stages; the zero-copy path gets its own A/B below.
+ pipeline.disable_zero_copy();
+ let pretok = format!("{:?}", pipeline.get_pre_tokenizer());
+ let pretok = pretok.split(['(', ' ']).next().unwrap_or("?");
+ println!("== {file} (pre-tokenizer: {pretok})");
+
+ for (name, text) in &fixtures {
+ for (regime, chunk_bytes) in [("10kB", CHUNK_BYTES), ("line", 0)] {
+ let chunks = chunks_of(text, chunk_bytes);
+ let bytes: usize = chunks.iter().map(String::len).sum();
+ let l = ladder(&pipeline, &chunks, bytes);
+ let total = l.added + l.norm + l.split + l.model + l.post;
+ println!(
+ " {name:<12} {regime:<4} ({:>5} chunks, {:>4} kB): added {:.3} | norm {:.3} | split {:.3} | model {:.3} | post {:.3} ns/B e2e {:.0} MB/s norm = {:.1}% of encode",
+ chunks.len(),
+ bytes / 1024,
+ l.added,
+ l.norm,
+ l.split,
+ l.model,
+ l.post,
+ l.total_mbs,
+ 100.0 * l.norm / total.max(1e-9),
+ );
+ if let Some(prepend) = replica_prepend {
+ let (count, alloc, rewrite) = decompose(&chunks, bytes, prepend);
+ println!(
+ " {:12} {regime:<4} norm decomposed: count {count:.3} + alloc {alloc:.3} + write {:.3} = replica {rewrite:.3} (ladder said {:.3})",
+ "",
+ (rewrite - count - alloc).max(0.0),
+ l.norm,
+ );
+ }
+ }
+ }
+
+ // Zero-copy A/B: the same binary and corpus, the two paths interleaved rep by rep so
+ // frequency drift hits both equally. Ids are compared over the whole corpus first.
+ if replica_prepend.is_some() {
+ let zero_copy = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(
+ zero_copy.has_zero_copy(),
+ "{file}: the zero-copy path should fire"
+ );
+ let mut written = PipelineTokenizer::try_from(&tok).unwrap();
+ written.disable_zero_copy();
+ for (name, text) in &fixtures {
+ for (regime, chunk_bytes) in [("10kB", CHUNK_BYTES), ("line", 0)] {
+ let chunks = chunks_of(text, chunk_bytes);
+ let bytes: usize = chunks.iter().map(String::len).sum();
+ for chunk in &chunks {
+ let ids = |p: &PipelineTokenizer| -> Vec {
+ p.encode(chunk, true)
+ .unwrap()
+ .iter()
+ .map(|t| t.id)
+ .collect()
+ };
+ assert_eq!(ids(&zero_copy), ids(&written), "{file} {name}: ids diverge");
+ }
+ let pass = |p: &PipelineTokenizer| {
+ let mut n = 0usize;
+ for chunk in &chunks {
+ n += p.encode(chunk, true).unwrap().len();
+ }
+ black_box(n);
+ };
+ pass(&zero_copy); // warm-up
+ pass(&written);
+ let mut zc = Vec::with_capacity(REPS);
+ let mut wr = Vec::with_capacity(REPS);
+ for _ in 0..REPS {
+ let t = Instant::now();
+ pass(&zero_copy);
+ zc.push(t.elapsed().as_secs_f64());
+ let t = Instant::now();
+ pass(&written);
+ wr.push(t.elapsed().as_secs_f64());
+ }
+ let (zc, wr) = (median(zc), median(wr));
+ println!(
+ " A/B {name:<12} {regime:<4}: zero-copy {:.0} MB/s vs written {:.0} MB/s ({:+.1}%)",
+ bytes as f64 / zc / 1e6,
+ bytes as f64 / wr / 1e6,
+ 100.0 * (wr / zc - 1.0),
+ );
+ }
+ }
+ }
+
+ // t5 and albert declare more normalizers than the space rewrite (t5 a Precompiled
+ // charsmap, albert five steps and then one). The design only removes the rewrite,
+ // so its share is measured by stripping the declared normalizer: what is left in
+ // the norm marginal comes from the Metaspace pre-tokenizer's rewriting half.
+ if replica_prepend.is_none() {
+ let mut stripped = Tokenizer::from_file(&path).unwrap();
+ let _ = stripped.with_normalizer(None::);
+ if let Ok(metaspace_only) = PipelineTokenizer::try_from(&stripped) {
+ let (name, text) = &fixtures[0];
+ let chunks = chunks_of(text, CHUNK_BYTES);
+ let bytes: usize = chunks.iter().map(String::len).sum();
+ let l = ladder(&metaspace_only, &chunks, bytes);
+ println!(
+ " declared normalizer stripped, {name} 10kB: norm {:.3} ns/B is the metaspace share",
+ l.norm,
+ );
+ }
+ }
+
+ // The second scan runs over every normalized chunk, but on an empty normalized
+ // vocabulary `Buckets::match_bytes` returns before touching the text. Injecting
+ // normalized added tokens (as `fixture_bench` does for every model) makes it a
+ // real pass; the frame marginal shows what the design's build-time gate saves.
+ let injected: Vec = NORMALIZED_WORDS
+ .iter()
+ .map(|w| AddedToken::from(*w, false).normalized(true))
+ .collect();
+ let _ = tok.add_tokens(injected);
+ if let Ok(with_normalized) = PipelineTokenizer::try_from(&tok) {
+ let (name, text) = &fixtures[0];
+ let chunks = chunks_of(text, CHUNK_BYTES);
+ let bytes: usize = chunks.iter().map(String::len).sum();
+ let before = ladder(&pipeline, &chunks, bytes);
+ let after = ladder(&with_normalized, &chunks, bytes);
+ println!(
+ " 2nd-scan probe on {name} 10kB: added {:.3} -> {:.3} ns/B with {} normalized added tokens",
+ before.added,
+ after.added,
+ NORMALIZED_WORDS.len(),
+ );
+ }
+ println!();
+ }
+}
diff --git a/tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs b/tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs
new file mode 100644
index 000000000..a19a13712
--- /dev/null
+++ b/tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs
@@ -0,0 +1,327 @@
+//! Which characters a byte-level vocabulary can emit as one token instead of as their bytes.
+//!
+//! A byte-level model's atoms are the 256 bytes, so a multi-byte character like え reaches the
+//! merge loop as three symbols that the merges then reassemble. When that reassembly is
+//! predetermined, seeding the merge loop with the character's own token skips the work. Two
+//! conditions make it predetermined:
+//!
+//! 1. The character's bytes must collapse into exactly one symbol when the merges are replayed
+//! the way BPE picks them: lowest rank first.
+//! 2. No step of that replay may be taken over from outside the character. The merge loop does
+//! not know where the character ends: if a neighbouring symbol can merge with the character's
+//! first or last symbol at a lower rank, that merge fires first and the assembly never
+//! happens.
+//!
+//! A character failing either test gets no table entry; its bytes go through the merge loop as
+//! usual, which is always exact. The fold is a shortcut, never a requirement.
+
+use std::cmp;
+
+use crate::models::bpe::MergeMap;
+use crate::utils::byte_level::CHAR_BYTES_LOOKUP;
+
+/// A character outside the table, or outside the Basic Multilingual Plane. The seeding loop
+/// falls back to per-byte symbols when it reads this.
+pub(super) const NO_FOLD: u32 = u32::MAX;
+
+/// `table[codepoint]` is the id of the token this character folds to, or [`NO_FOLD`].
+///
+/// Indexed by Basic Multilingual Plane codepoints only: a `char` above `u16::MAX` never folds.
+/// Single-byte characters are also left out, their byte symbol already is the seed.
+pub(super) type FoldTable = Box<[u32; FOLD_TABLE_LEN]>;
+pub(super) const FOLD_TABLE_LEN: usize = 1 << 16;
+
+/// Builds the fold table for a byte-level vocabulary.
+///
+/// `vocab` still spells its tokens in byte-level characters ("é" for é), which is what
+/// [`ByteLevelFold::fold`] expects; call this before the store is rebuilt on raw bytes.
+pub(super) fn build_fold_table(vocab: &[(String, u32)], merges: &MergeMap) -> FoldTable {
+ let fold = ByteLevelFold::new(vocab, merges);
+ let mut table = vec![NO_FOLD; FOLD_TABLE_LEN];
+ for (token, id) in vocab {
+ if let Fold::Folds(ch, id) = fold.fold(token, *id)
+ && ch.len_utf8() > 1
+ && (ch as usize) < FOLD_TABLE_LEN
+ {
+ table[ch as usize] = id;
+ }
+ }
+ table
+ .into_boxed_slice()
+ .try_into()
+ .expect("length is FOLD_TABLE_LEN")
+}
+
+/// What one vocab token is worth to the fold table.
+pub(super) enum Fold {
+ /// A single character whose bytes assemble to exactly this token, un-stealably.
+ Folds(char, u32),
+ /// Formable, but some step could be taken over by a neighbour.
+ Unsafe,
+ /// Not a single character, or its bytes never assemble at all. Nothing to record.
+ Skip,
+}
+
+pub(super) struct ByteLevelFold<'a> {
+ /// byte -> id of that byte's own one-character token. A byte's value is not its id
+ /// (gpt2: 0x41 -> 32, 0x20 -> 220), which is why this indirection exists.
+ byte_token: [u32; 256],
+ /// `stolen_from_left[id]`: the lowest rank at which a left neighbour merges with `id`,
+ /// counting only neighbours reachable at a character boundary. Same on the other side for
+ /// `stolen_from_right`.
+ stolen_from_left: Vec,
+ stolen_from_right: Vec,
+ merges: &'a MergeMap,
+}
+
+impl<'a> ByteLevelFold<'a> {
+ pub(super) fn new(vocab: &[(String, u32)], merges: &'a MergeMap) -> Self {
+ let mut byte_token = [u32::MAX; 256];
+ for (token, id) in vocab {
+ let mut chars = token.chars();
+ if let (Some(ch), None) = (chars.next(), chars.next())
+ && let Some(&b) = CHAR_BYTES_LOOKUP.get(&ch)
+ {
+ byte_token[b as usize] = *id;
+ }
+ }
+
+ let (stolen_from_left, stolen_from_right) = boundary_merge_ranks(vocab, merges);
+
+ Self {
+ byte_token,
+ stolen_from_left,
+ stolen_from_right,
+ merges,
+ }
+ }
+
+ /// Verdict for `token`, whose id is `id`.
+ pub(super) fn fold(&self, token: &str, id: u32) -> Fold {
+ let Some(bytes) = token
+ .chars()
+ .map(|ch| CHAR_BYTES_LOOKUP.get(&ch).copied())
+ .collect::>>()
+ else {
+ // A character with no byte-level mapping: an added token, not text.
+ return Fold::Skip;
+ };
+ // The table is keyed by codepoint, so only single-character tokens can go in it. This
+ // also drops lone bytes >= 0x80, which are not characters on their own.
+ let Ok(text) = std::str::from_utf8(&bytes) else {
+ return Fold::Skip;
+ };
+ let mut it = text.chars();
+ let (Some(ch), None) = (it.next(), it.next()) else {
+ return Fold::Skip;
+ };
+
+ let mut running: Vec = bytes.iter().map(|&b| self.byte_token[b as usize]).collect();
+ if running.contains(&u32::MAX) {
+ return Fold::Skip; // a byte with no token of its own: never assemblable
+ }
+ // Replay the assembly the way the merge loop picks: the lowest rank among the pairs
+ // still standing, until one symbol is left.
+ while running.len() > 1 {
+ let mut best: Option<(usize, u32, u32)> = None;
+ for i in 0..running.len() - 1 {
+ let pair = (running[i], running[i + 1]);
+ if let Some((rank, product)) = self.merges.get(&pair)
+ && best.is_none_or(|(_, best_rank, _)| *rank < best_rank)
+ {
+ best = Some((i, *rank, *product));
+ }
+ }
+ let Some((i, rank, product)) = best else {
+ return Fold::Skip; // stuck above one symbol: reference BPE stops here too
+ };
+ if rank >= self.stolen_from_left[running[0] as usize]
+ || rank >= self.stolen_from_right[*running.last().unwrap() as usize]
+ {
+ return Fold::Unsafe;
+ }
+ running[i] = product;
+ running.remove(i + 1);
+ }
+
+ debug_assert_eq!(running[0], id);
+ Fold::Folds(ch, running[0])
+ }
+}
+
+/// A folded character's edges are character boundaries by construction, and UTF-8
+/// pins down what may be there:
+///
+/// - right neighbour = the next character's FIRST byte -> ASCII or a lead byte, never 0x80..=0xBF
+/// - left neighbour = the previous character's LAST byte -> never a lead byte, so always < 0xC0
+fn boundary_merge_ranks(vocab: &[(String, u32)], merges: &MergeMap) -> (Vec, Vec) {
+ let max_id = vocab
+ .iter()
+ .map(|(_, id)| *id)
+ .chain(merges.iter().flat_map(|((a, b), (_, id))| [*a, *b, *id]))
+ .max()
+ .map_or(0, |id| id as usize + 1);
+
+ // First and last real byte of every token. 0xFF marks a token with no byte spelling (an
+ // added token): it counts as a possible stealer on the right, and as none on the left,
+ // which errs towards refusing a fold.
+ let (mut first, mut last) = (vec![0xFFu8; max_id], vec![0xFFu8; max_id]);
+ for (token, id) in vocab {
+ let Some(bytes) = token
+ .chars()
+ .map(|c| CHAR_BYTES_LOOKUP.get(&c).copied())
+ .collect::>>()
+ else {
+ continue;
+ };
+ if let (Some(f), Some(l)) = (bytes.first(), bytes.last()) {
+ first[*id as usize] = *f;
+ last[*id as usize] = *l;
+ }
+ }
+ // 0xC0/0xC1 are overlong lead bytes and cannot occur on either side.
+ let starts_at_boundary = |id: u32| first[id as usize] < 0x80 || first[id as usize] >= 0xC2;
+ let ends_at_boundary = |id: u32| last[id as usize] < 0xC0;
+
+ let mut stolen_from_left = vec![u32::MAX; max_id];
+ let mut stolen_from_right = vec![u32::MAX; max_id];
+ for ((a, b), (rank, _)) in merges.iter() {
+ if *a as usize >= max_id || *b as usize >= max_id {
+ continue;
+ }
+ if starts_at_boundary(*b) {
+ stolen_from_right[*a as usize] = cmp::min(stolen_from_right[*a as usize], *rank);
+ }
+ if ends_at_boundary(*a) {
+ stolen_from_left[*b as usize] = cmp::min(stolen_from_left[*b as usize], *rank);
+ }
+ }
+ (stolen_from_left, stolen_from_right)
+}
+
+#[cfg(test)]
+mod test {
+ use super::{ByteLevelFold, Fold, NO_FOLD, build_fold_table};
+ use crate::models::bpe::MergeMap;
+ use crate::utils::byte_level::BYTES_CHAR_LOOKUP;
+
+ /// 'é' is U+00E9 = bytes C3 A9; both are printable latin-1, so the byte-level names are the
+ /// identity chars 'Ã' and '©' and the vocab spells the character "é".
+ fn setup(extra_merge: bool) -> (Vec<(String, u32)>, MergeMap) {
+ let vocab = Vec::from([
+ ("Ã".to_string(), 0), // byte 0xC3
+ ("©".to_string(), 1), // byte 0xA9
+ ("é".to_string(), 2), // the character é
+ ("x".to_string(), 3),
+ ("xÃ".to_string(), 4),
+ ]);
+ let mut merges = MergeMap::new();
+ merges.insert((0, 1), (1, 2)); // Ã + © -> é at rank 1
+ if extra_merge {
+ // x + Ã at rank 0: a left neighbour "x" grabs our first byte first, so the
+ // assembly of é never happens and folding it would be wrong.
+ merges.insert((3, 0), (0, 4));
+ }
+ (vocab, merges)
+ }
+
+ #[test]
+ fn folds_when_nothing_can_steal_an_edge() {
+ let (vocab, merges) = setup(false);
+ let f = ByteLevelFold::new(&vocab, &merges);
+ assert!(matches!(f.fold("é", 2), Fold::Folds('é', 2)));
+ }
+
+ #[test]
+ fn rejects_a_boundary_steal() {
+ let (vocab, merges) = setup(true);
+ let f = ByteLevelFold::new(&vocab, &merges);
+ assert!(matches!(f.fold("é", 2), Fold::Unsafe));
+ }
+
+ #[test]
+ fn skips_what_is_not_one_character() {
+ let (vocab, merges) = setup(false);
+ let f = ByteLevelFold::new(&vocab, &merges);
+ assert!(matches!(f.fold("xÃ", 4), Fold::Skip)); // two characters once decoded
+ assert!(matches!(f.fold("<|endoftext|>", 9), Fold::Skip)); // '<' is fine, '|' is not remapped
+ assert!(matches!(f.fold("Ã", 0), Fold::Skip)); // lone 0xC3 is not valid UTF-8
+ assert!(matches!(f.fold("x", 3), Fold::Folds('x', 3))); // ASCII needs no assembly
+ }
+
+ #[test]
+ fn table_keeps_the_fold_and_drops_ascii() {
+ let (vocab, merges) = setup(false);
+ let table = build_fold_table(&vocab, &merges);
+ assert_eq!(table['é' as usize], 2);
+ // 'x' folds but single-byte characters stay out: their byte symbol already is the seed.
+ assert_eq!(table['x' as usize], NO_FOLD);
+ }
+
+ // Byte-level merges in a gpt2-like encoding. Byte level rewrites the vocab so every byte is a
+ // printable char; bytes 0x80..=0xA0 become U+0122.. and 0xAE..=0xFF stay themselves.
+ // U+671D 朝 → E6 9C 9D -> 'æ','ľ','Ŀ'
+ // U+65E5 日 → E6 97 A5 -> 'æ','Ĺ','¥'
+ //
+ // 朝 assembles with ('æ','ľ') and ('æľ','Ŀ'), so it may only merge if no merge pair can take
+ // an edge symbol first. We add such a pair: ('Ŀ','æ') 9D E6, which appears in 朝朝 and 朝日
+ // at the boundary `.. 9D | E6 ..`. It has to be a LEAD byte (E6) doing the stealing: the symbol
+ // after a complete character is always the next character's first byte.
+ enum Thief {
+ None,
+ Lead,
+ Continuation,
+ }
+
+ fn cjk_vocab(thief: Thief) -> (Vec<(String, u32)>, MergeMap) {
+ assert_eq!(
+ [0xE6u8, 0x9C, 0x9D].map(|b| BYTES_CHAR_LOOKUP[b as usize]),
+ ['æ', 'ľ', 'Ŀ']
+ );
+ let mut vocab = Vec::from([
+ ("æ".to_string(), 0), // E6
+ ("ľ".to_string(), 1), // 9C
+ ("Ŀ".to_string(), 2), // 9D
+ ("æľ".to_string(), 3), // E6 9C
+ ("æľĿ".to_string(), 4), // E6 9C 9D = 朝
+ ]);
+ let mut merges = MergeMap::new();
+ // (left, right) -> (rank, product). Ranks leave room below for the thief.
+ merges.insert((0, 1), (1, 3)); // 'æ' + 'ľ' -> "æľ"
+ merges.insert((3, 2), (2, 4)); // "æľ" + 'Ŀ' -> 朝
+ match thief {
+ Thief::None => {}
+ Thief::Lead => {
+ vocab.push(("Ŀæ".to_string(), 5)); // 9D E6, straddles a character boundary
+ merges.insert((2, 0), (0, 5)); // rank 0, below every step of 朝's assembly
+ }
+ Thief::Continuation => {
+ vocab.push(("Ģ".to_string(), 5)); // 80
+ vocab.push(("ĿĢ".to_string(), 6)); // 9D 80, never at a boundary
+ merges.insert((2, 5), (0, 6));
+ }
+ }
+ (vocab, merges)
+ }
+
+ #[test]
+ fn folds_a_cjk_char_when_no_neighbour_can_steal() {
+ let (vocab, merges) = cjk_vocab(Thief::None);
+ let f = ByteLevelFold::new(&vocab, &merges);
+ assert!(matches!(f.fold("æľĿ", 4), Fold::Folds('朝', 4)));
+ }
+
+ #[test]
+ fn refuses_the_same_char_once_a_lead_byte_can_steal() {
+ let (vocab, merges) = cjk_vocab(Thief::Lead);
+ let f = ByteLevelFold::new(&vocab, &merges);
+ assert!(matches!(f.fold("æľĿ", 4), Fold::Unsafe));
+ }
+
+ #[test]
+ fn a_continuation_byte_cannot_steal_so_it_still_folds() {
+ let (vocab, merges) = cjk_vocab(Thief::Continuation);
+ let f = ByteLevelFold::new(&vocab, &merges);
+ assert!(matches!(f.fold("æľĿ", 4), Fold::Folds('朝', 4)));
+ }
+}
diff --git a/tokenizers/tk-encode/src/models/bpe/mod.rs b/tokenizers/tk-encode/src/models/bpe/mod.rs
index 6e1cb2da9..06a7ebc5a 100644
--- a/tokenizers/tk-encode/src/models/bpe/mod.rs
+++ b/tokenizers/tk-encode/src/models/bpe/mod.rs
@@ -1,6 +1,7 @@
//! [Byte Pair Encoding](https://www.aclweb.org/anthology/P16-1162/) model.
use std::{iter, mem};
+mod bytelevel_folding;
mod model;
mod serialization;
pub mod word;
diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs
index fa7bcec23..a05fbc6dd 100644
--- a/tokenizers/tk-encode/src/models/bpe/model.rs
+++ b/tokenizers/tk-encode/src/models/bpe/model.rs
@@ -1,10 +1,11 @@
-use super::{super::OrderedVocabIter, Error, Pair, Word};
+use super::{super::OrderedVocabIter, Error, Pair, Word, bytelevel_folding};
use crate::models::bpe::Merge;
use crate::pipeline::{self, ModelScratch, PipelineToken};
use crate::tokenizer::{Model, Result, Token};
use crate::utils::byte_level::{self};
use crate::utils::cache::{DEFAULT_CACHE_CAPACITY, MAX_LENGTH};
use crate::utils::iter::ResultShunt;
+use crate::utils::word_cache::{Lookup, WordCache};
use crate::vocab::bucket_vocab_store::BucketVocabStore;
use crate::vocab_store::VocabStore;
use ahash::AHashMap;
@@ -680,11 +681,16 @@ pub struct PipelineBPE {
vocab: BucketVocabStore,
merges: MergeMap,
ignore_merges: bool,
+ /// `None` when the tokenizer asked for no cache.
+ cache_capacity: Option,
}
enum Atoms {
Bytes {
byte_to_id: [u32; 256],
+ /// Multi-byte characters proven to always assemble into one token; see
+ /// [`bytelevel_folding`]. The merge loop is seeded with that token directly.
+ fold: bytelevel_folding::FoldTable,
},
Chars {
byte_fallback: Option<[u32; 256]>,
@@ -715,6 +721,7 @@ impl PipelineBPE {
} = model;
let (vocab, atoms) = if with_byte_level {
+ let fold = bytelevel_folding::build_fold_table(&vocab.content(), &merges);
let mut vocab = BucketVocabStore::build(vocab.byte_content());
vocab = byte_level::transform_vocab(vocab);
let mut byte_to_id = [0u32; 256];
@@ -723,7 +730,7 @@ impl PipelineBPE {
.get_bytes(&[b])
.ok_or(Error::ByteAtomOutOfVocabulary(b))?;
}
- (vocab, Atoms::Bytes { byte_to_id })
+ (vocab, Atoms::Bytes { byte_to_id, fold })
} else {
let vocab = BucketVocabStore::build(vocab.byte_content());
let unk_token = if let Some(unk_str) = unk_token {
@@ -760,21 +767,68 @@ impl PipelineBPE {
ignore_merges,
merges,
vocab,
+ cache_capacity: model.cache.map(|c| c.capacity).filter(|&c| c > 0),
})
}
+ /// Does this model look the whole text up in the vocabulary before merging it?
+ pub(crate) fn ignore_merges(&self) -> bool {
+ self.ignore_merges
+ }
+
+ /// Does this model seed the merge loop with characters (rather than raw bytes)? Only
+ /// character atoms can read a span through a [`SpaceSwap`], which swaps one char.
+ pub(crate) fn char_atoms(&self) -> bool {
+ matches!(self.atoms, Atoms::Chars { .. })
+ }
+
+ /// The id of the vocabulary piece spelled exactly `bytes`, if there is one.
+ pub(crate) fn id_of_bytes(&self, bytes: &[u8]) -> Option {
+ self.vocab.get_bytes(bytes)
+ }
+
+ /// Every piece of the vocabulary, as bytes. Used at build time to work out how the text may be
+ /// cut into words; not cheap, so call it once.
+ pub(crate) fn vocab_bytes(&self) -> Vec<(Vec, u32)> {
+ self.vocab.byte_content()
+ }
+
fn merge_word(
&self,
sequence: &str,
+ swap: Option,
merge_queue: &mut QuaternaryHeap,
skip: &mut Vec,
word: &mut Word,
) {
word.clear();
match &self.atoms {
- Atoms::Bytes { byte_to_id } => {
- for &b in sequence.as_bytes() {
- word.add(byte_to_id[b as usize], 1);
+ Atoms::Bytes { byte_to_id, fold } => {
+ let bytes = sequence.as_bytes();
+ let mut i = 0;
+ while i < bytes.len() {
+ let b = bytes[i];
+ if b < 0x80 {
+ word.add(byte_to_id[b as usize], 1);
+ i += 1;
+ continue;
+ }
+ // A multi-byte character seeds as its own token when the fold proved the
+ // merges would assemble it anyway, and as its bytes otherwise.
+ let ch = sequence[i..].chars().next().unwrap();
+ let len = ch.len_utf8();
+ let id = match (ch as usize) < bytelevel_folding::FOLD_TABLE_LEN {
+ true => fold[ch as usize],
+ false => bytelevel_folding::NO_FOLD,
+ };
+ if id != bytelevel_folding::NO_FOLD {
+ word.add(id, len);
+ } else {
+ for &b in &bytes[i..i + len] {
+ word.add(byte_to_id[b as usize], 1);
+ }
+ }
+ i += len;
}
}
Atoms::Chars {
@@ -786,6 +840,14 @@ impl PipelineBPE {
.char_indices()
.map(|(i, c)| &sequence[i..i + c.len_utf8()])
{
+ // A raw `from` stands for its replacement: its atom is the replacement's
+ // id, as wide as what the rewrite would have written.
+ if let Some(swap) = swap
+ && char_str.as_bytes() == [swap.from]
+ {
+ word.add(swap.id, swap.len as usize);
+ continue;
+ }
let char_len = char_str.len();
if let Some(char_id) = self.vocab.token_to_id(char_str) {
word.add(char_id, char_len);
@@ -814,37 +876,124 @@ impl PipelineBPE {
}
}
-impl pipeline::Model for PipelineBPE {
- type Scratch = BpeScratch;
+/// One raw character read as the one a pending rewrite would swap in: the raw `from` (a
+/// single byte, all its route supports), the replacement's vocabulary id, and the byte width
+/// the rewritten form would have.
+///
+/// The zero-copy metaspace path (see [`crate::pre_tokenizers::proven_cuts`]) hands this to
+/// [`PipelineBPE::tokenize_swapped`] so the merge loop seeds spans that were never rewritten.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct CharSwap {
+ pub(crate) from: u8,
+ pub(crate) id: u32,
+ pub(crate) len: u8,
+}
- fn tokenize_pipeline(
+impl PipelineBPE {
+ /// One pre-token through the cache and the merge loop. This is the whole model
+ /// step for a span and it cannot fail; it stands alone so the fused byte-level
+ /// scan can call it per span as the span is cut, and
+ /// [`Model::tokenize_pipeline`](pipeline::Model::tokenize_pipeline) is the same
+ /// body behind the trait.
+ pub(crate) fn tokenize_span(
&self,
sequence: &str,
- scratch: &mut Self::Scratch,
+ scratch: &mut BpeScratch,
output: &mut Vec,
- ) -> Result<()> {
+ ) {
if sequence.is_empty() {
- return Ok(());
+ return;
}
+ let BpeScratch {
+ merge_queue,
+ skip,
+ word,
+ word_cache,
+ } = scratch;
+
+ let mut placement = None;
+ if let Some(cache) = word_cache.as_mut() {
+ match cache.lookup(sequence.as_bytes()) {
+ Lookup::Hit(ids) => {
+ output.extend(ids.iter().map(|&id| PipelineToken { id }));
+ return;
+ }
+ Lookup::Miss(at) => placement = at,
+ }
+ }
+ let start = output.len();
if self.ignore_merges
&& let Some(id) = self.vocab.get_bytes(sequence.as_bytes())
{
output.push(PipelineToken { id });
- return Ok(());
+ } else {
+ self.merge_word(sequence, None, merge_queue, skip, word);
+ output.extend(word.get_chars_iter().map(|id| PipelineToken { id }));
}
+ // The ids come back out of `output` because that is the only place both
+ // branches above leave them: `ignore_merges` never touches `word`.
+ if let Some(cache) = word_cache.as_mut()
+ && let Some(at) = placement
+ {
+ cache.insert(at, output[start..].iter().map(|token| token.id));
+ }
+ }
- // TODO: persistent cache mapping &str -> &[u32]
-
+ /// [`Self::tokenize_span`] for a raw span read through `swap`. The cache key is the
+ /// span's raw bytes; a rewritten form never holds the swap's `from`, so raw and rewritten
+ /// keys can only collide when they spell the same word, and the cached ids agree.
+ /// `ignore_merges` plays no part here: the zero-copy path is only built for models
+ /// without it, as the proven cuts it rides on already require.
+ pub(crate) fn tokenize_swapped(
+ &self,
+ sequence: &str,
+ swap: CharSwap,
+ scratch: &mut BpeScratch,
+ output: &mut Vec,
+ ) -> Result<()> {
+ if sequence.is_empty() {
+ return Ok(());
+ }
let BpeScratch {
merge_queue,
skip,
word,
+ word_cache,
} = scratch;
- self.merge_word(sequence, merge_queue, skip, word);
+ let mut placement = None;
+ if let Some(cache) = word_cache.as_mut() {
+ match cache.lookup(sequence.as_bytes()) {
+ Lookup::Hit(ids) => {
+ output.extend(ids.iter().map(|&id| PipelineToken { id }));
+ return Ok(());
+ }
+ Lookup::Miss(at) => placement = at,
+ }
+ }
+ let start = output.len();
+ self.merge_word(sequence, Some(swap), merge_queue, skip, word);
output.extend(word.get_chars_iter().map(|id| PipelineToken { id }));
+ if let Some(cache) = word_cache.as_mut()
+ && let Some(at) = placement
+ {
+ cache.insert(at, output[start..].iter().map(|token| token.id));
+ }
+ Ok(())
+ }
+}
+impl pipeline::Model for PipelineBPE {
+ type Scratch = BpeScratch;
+
+ fn tokenize_pipeline(
+ &self,
+ sequence: &str,
+ scratch: &mut Self::Scratch,
+ output: &mut Vec,
+ ) -> Result<()> {
+ self.tokenize_span(sequence, scratch, output);
Ok(())
}
@@ -853,6 +1002,7 @@ impl pipeline::Model for PipelineBPE {
merge_queue: QuaternaryHeap::with_capacity(64),
word: Word::with_capacity(64),
skip: Vec::new(),
+ word_cache: self.cache_capacity.map(WordCache::new),
}
}
}
@@ -861,6 +1011,8 @@ pub struct BpeScratch {
pub(crate) merge_queue: QuaternaryHeap,
pub(crate) skip: Vec,
pub(crate) word: Word,
+ /// Outlives the encode call that fills it, or it would never see a word twice.
+ pub(crate) word_cache: Option,
}
impl ModelScratch for BpeScratch {}
@@ -1438,13 +1590,61 @@ mod tests {
assert!(pipeline_ids(&pipeline, "").is_empty());
}
+ // The pool hands the SAME scratch to successive encodes. State left behind by one
+ // call (an undrained merge queue, a stale word buffer) would corrupt every call
+ // after it. Drive several inputs through one scratch and check each still matches
+ // the reference model.
+ #[test]
+ fn reused_scratch_matches_fresh() {
+ let bpe = hello_builder().build().unwrap();
+ let reference = bpe.clone();
+ let model = PipelineBPE::from_bpe(bpe, false).unwrap();
+ let mut scratch = model.init_scratch();
+ for input in ["hello", "hell", "helo", "oleh", "hello", "", "hxe"] {
+ let mut out = Vec::new();
+ pipeline::Model::tokenize_pipeline(&model, input, &mut scratch, &mut out).unwrap();
+ let got: Vec = out.iter().map(|t| t.id).collect();
+ assert_eq!(got, reference_ids(&reference, input), "{input:?}");
+ }
+ }
+
+ // A cache may forget a word, but it must never change one. Run every word twice
+ // through one scratch, the second time answered from the cache, against a model
+ // built with no cache at all.
+ #[test]
+ fn cached_ids_match_uncached() {
+ let cached = PipelineBPE::from_bpe(hello_builder().build().unwrap(), false).unwrap();
+ let uncached =
+ PipelineBPE::from_bpe(hello_builder().cache_capacity(0).build().unwrap(), false)
+ .unwrap();
+
+ let mut scratch = cached.init_scratch();
+ assert!(scratch.word_cache.is_some(), "nothing is being cached");
+ for _ in 0..2 {
+ for word in [
+ "hello",
+ "hell",
+ "o",
+ "hellohello",
+ "hello-a-word-past-fifteen-bytes",
+ "hxe",
+ ] {
+ let mut out = Vec::new();
+ pipeline::Model::tokenize_pipeline(&cached, word, &mut scratch, &mut out)
+ .unwrap();
+ let got: Vec = out.iter().map(|t| t.id).collect();
+ assert_eq!(got, pipeline_ids(&uncached, word), "{word:?}");
+ }
+ }
+ }
+
#[test]
fn unknown_char_without_unk_is_dropped() {
let bpe = hello_builder().build().unwrap();
let reference = bpe.clone();
let pipeline = PipelineBPE::from_bpe(bpe, false).unwrap();
// 'x' vanishes, making 'h' and 'e' adjacent, so the (h,e) merge
- // applies — mirrors the reference model.
+ // applies, mirroring the reference model.
assert_eq!(pipeline_ids(&pipeline, "hxe"), vec![4]);
assert_eq!(
pipeline_ids(&pipeline, "hxe"),
@@ -1615,7 +1815,7 @@ mod tests {
/// A gpt2-shaped miniature: the 256 projected single-byte tokens
/// (id == byte value) plus `extra` tokens and merges, given in raw
- /// space and projected here — like a real byte-level tokenizer.json,
+ /// space and projected here, like a real byte-level tokenizer.json,
/// whose vocab is stored in the projected alphabet.
fn byte_level_bpe(
extra: &[(&str, u32)],
@@ -1672,6 +1872,175 @@ mod tests {
);
}
+ /// Blanks a model's fold table, so it seeds every character as bytes again.
+ fn without_fold(mut model: PipelineBPE) -> PipelineBPE {
+ let Atoms::Bytes { fold, .. } = &mut model.atoms else {
+ panic!("not a byte-level model");
+ };
+ *fold = vec![bytelevel_folding::NO_FOLD; bytelevel_folding::FOLD_TABLE_LEN]
+ .into_boxed_slice()
+ .try_into()
+ .unwrap();
+ model
+ }
+
+ /// 'é' (bytes C3 A9) has a token and an unstealable assembly, so the fold seeds it as
+ /// that token; a thief merge on its first byte must blank the entry, and a character
+ /// outside the Basic Multilingual Plane never enters the table. Ids stay the same
+ /// either way.
+ #[test]
+ fn folded_char_seeds_as_its_token() {
+ let n = |b: u8| BYTES_CHAR_LOOKUP[b as usize].to_string();
+ let mut vocab: Vocab = (0..=255u8).map(|b| (n(b), u32::from(b))).collect();
+ vocab.insert(n(0xC3) + &n(0xA9), 300);
+ let merges = vec![(n(0xC3), n(0xA9))];
+ let bpe = BpeBuilder::default()
+ .vocab_and_merges(vocab.clone(), merges.clone())
+ .build()
+ .unwrap();
+ let pipeline = PipelineBPE::from_bpe(bpe, true).unwrap();
+
+ let Atoms::Bytes { fold, .. } = &pipeline.atoms else {
+ panic!("not a byte-level model");
+ };
+ assert_eq!(fold['é' as usize], 300);
+ assert_eq!(pipeline_ids(&pipeline, "é"), vec![300]);
+ // 4-byte characters stay out of the table and encode as their bytes.
+ assert_eq!(pipeline_ids(&pipeline, "🎉").len(), 4);
+
+ // "x" + C3 at rank 0 fires before é's assembly, so é may not fold; the byte
+ // path still merges it wherever no thief is adjacent.
+ vocab.insert("x".to_string() + &n(0xC3), 301);
+ let mut merges = merges;
+ merges.insert(0, ("x".to_string(), n(0xC3)));
+ let bpe = BpeBuilder::default()
+ .vocab_and_merges(vocab, merges)
+ .build()
+ .unwrap();
+ let pipeline = PipelineBPE::from_bpe(bpe, true).unwrap();
+ let Atoms::Bytes { fold, .. } = &pipeline.atoms else {
+ panic!("not a byte-level model");
+ };
+ assert_eq!(fold['é' as usize], bytelevel_folding::NO_FOLD);
+ assert_eq!(pipeline_ids(&pipeline, "xé"), vec![301, 0xA9]);
+ assert_eq!(pipeline_ids(&pipeline, "é"), vec![300]);
+ }
+
+ /// The fold is a shortcut, never a different answer: over a real vocabulary, seeding
+ /// through the table and seeding plain bytes must give identical ids. Skipped when the
+ /// file is not fetched.
+ #[test]
+ fn fold_seeding_matches_byte_seeding_on_llama3() {
+ let path = "../data/llama-3-tokenizer.json";
+ if !std::path::Path::new(path).exists() {
+ eprintln!("skip llama-3: not present (fetch with `make bench-models`)");
+ return;
+ }
+ let tok = crate::Tokenizer::from_file(path).unwrap();
+ let crate::models::ModelWrapper::BPE(bpe) = tok.get_model().clone() else {
+ panic!("llama-3 is BPE");
+ };
+ let folded = PipelineBPE::from_bpe(bpe.clone(), true).unwrap();
+ let plain = without_fold(PipelineBPE::from_bpe(bpe, true).unwrap());
+
+ let Atoms::Bytes { fold, .. } = &folded.atoms else {
+ panic!("llama-3 is byte-level");
+ };
+ let entries = fold
+ .iter()
+ .filter(|&&id| id != bytelevel_folding::NO_FOLD)
+ .count();
+ assert!(entries > 0, "no character folds in llama-3's vocabulary");
+ let cjk = fold[0x3000..]
+ .iter()
+ .filter(|&&id| id != bytelevel_folding::NO_FOLD)
+ .count();
+ eprintln!("llama-3 folds {entries} characters, {cjk} at U+3000 and above");
+
+ for text in [
+ "朝日新聞デジタルの記事一覧です。",
+ "缓存的错误信息不完整，请重试。",
+ "Καλημέρα κόσμε, добрый день",
+ "mixed ascii と 日本語 and émojis 🎉🚀 nbsp\u{a0}end",
+ " leading spaces, ▁marks and érrors ",
+ ] {
+ assert_eq!(
+ pipeline_ids(&folded, text),
+ pipeline_ids(&plain, text),
+ "{text:?}"
+ );
+ }
+ }
+
+ /// Not a test: a manual probe for what the fold is worth on CJK text, model stage only,
+ /// cache off. Run with `cargo test --release -p tk-encode --lib fold_speed -- --ignored
+ /// --nocapture`.
+ #[test]
+ #[ignore]
+ fn fold_speed_probe_llama3_cjk() {
+ let path = "../data/llama-3-tokenizer.json";
+ let jpn = std::fs::read_to_string("../data/fixtures/lang/jpn_Jpan.txt");
+ let (Ok(tok), Ok(text)) = (crate::Tokenizer::from_file(path), jpn) else {
+ eprintln!("skip: llama-3 or the jpn fixture is not fetched");
+ return;
+ };
+ let crate::models::ModelWrapper::BPE(bpe) = tok.get_model().clone() else {
+ panic!("llama-3 is BPE");
+ };
+ let mut folded = PipelineBPE::from_bpe(bpe.clone(), true).unwrap();
+ folded.cache_capacity = None;
+ let mut plain = without_fold(PipelineBPE::from_bpe(bpe, true).unwrap());
+ plain.cache_capacity = None;
+
+ // 32-character chunks as a stand-in for the pre-tokenizer's CJK runs.
+ let chunks: Vec<&str> = {
+ let mut chunks = Vec::new();
+ let mut rest = text.as_str();
+ while let Some((i, _)) = rest.char_indices().nth(32) {
+ chunks.push(&rest[..i]);
+ rest = &rest[i..];
+ }
+ chunks
+ };
+ let bytes: usize = chunks.iter().map(|c| c.len()).sum();
+
+ for (name, model) in [("folded", &folded), ("plain", &plain)] {
+ let mut scratch = model.init_scratch();
+ let mut out = Vec::new();
+ for _ in 0..2 {
+ // warm-up rep first; the second rep is the one printed
+ let start = std::time::Instant::now();
+ for chunk in &chunks {
+ out.clear();
+ pipeline::Model::tokenize_pipeline(model, chunk, &mut scratch, &mut out)
+ .unwrap();
+ }
+ let elapsed = start.elapsed();
+ eprintln!(
+ "{name}: {:.2} ns/B over {} chunks ({bytes} B)",
+ elapsed.as_nanos() as f64 / bytes as f64,
+ chunks.len(),
+ );
+ }
+ }
+ }
+
+ /// A whole-word vocab hit skips the merge loop but not the vocab lookup,
+ /// and in a byte-level vocab the words that take that path are the
+ /// commonest ones in the text. Storing the id it found turns the next
+ /// occurrence into a cache hit.
+ #[test]
+ fn ignore_merges_stores_the_whole_word_id() {
+ let bpe = byte_level_bpe(&[(" hello", 300)], &[], true);
+ let model = PipelineBPE::from_bpe(bpe, true).unwrap();
+ let mut scratch = model.init_scratch();
+ let mut out = Vec::new();
+ PipelineModel::tokenize_pipeline(&model, " hello", &mut scratch, &mut out).unwrap();
+
+ let cache = scratch.word_cache.as_mut().unwrap();
+ assert_eq!(cache.lookup(b" hello").hit(), Some(&[300u32][..]));
+ }
+
#[test]
fn byte_level_requires_full_byte_coverage() {
// An ASCII-only vocab covers no control/high bytes: building the
diff --git a/tokenizers/tk-encode/src/models/bpe/word.rs b/tokenizers/tk-encode/src/models/bpe/word.rs
index ac463d6e1..006aa6a8d 100644
--- a/tokenizers/tk-encode/src/models/bpe/word.rs
+++ b/tokenizers/tk-encode/src/models/bpe/word.rs
@@ -276,7 +276,7 @@ impl Word {
self.get_chars_iter().collect()
}
- pub fn get_chars_iter(&self) -> impl Iterator
- + '_ {
+ pub fn get_chars_iter(&self) -> impl ExactSizeIterator
- + '_ {
self.symbols.iter().map(|s| s.c)
}
diff --git a/tokenizers/tk-encode/src/models/unigram/model.rs b/tokenizers/tk-encode/src/models/unigram/model.rs
index 1cc3f4b98..7cdd2d18c 100644
--- a/tokenizers/tk-encode/src/models/unigram/model.rs
+++ b/tokenizers/tk-encode/src/models/unigram/model.rs
@@ -3,6 +3,7 @@ use super::{
trie::{Trie, TrieBuilder},
};
use crate::utils::cache::{Cache, MAX_LENGTH};
+use crate::utils::word_cache::{Lookup, WordCache};
use crate::vocab_store::VocabStore;
use crate::{
pipeline::{self, PipelineToken},
@@ -239,23 +240,33 @@ impl Unigram {
if sentence.is_empty() {
return Ok(vec![]);
}
- if self.alpha.is_none() || self.alpha == Some(0.0) {
- if let Some(result) = self.cache.get(sentence) {
- Ok(result.to_vec())
- } else {
- let result = if self.is_optimized {
- self.encode_optimized(sentence)?
- } else {
- self.encode_unoptimized(sentence)?
- };
- if sentence.len() < MAX_LENGTH {
- self.cache.set(sentence.to_owned(), result.clone());
- }
- Ok(result)
- }
+ if self.samples() {
+ return self.encode_uncached(sentence);
+ }
+ if let Some(result) = self.cache.get(sentence) {
+ return Ok(result.to_vec());
+ }
+ let result = self.encode_uncached(sentence)?;
+ if sentence.len() < MAX_LENGTH {
+ self.cache.set(sentence.to_owned(), result.clone());
+ }
+ Ok(result)
+ }
+
+ /// Whether [`Unigram::alpha`] asks for a tokenization drawn at random from the
+ /// lattice instead of its best path. Such a result is one draw out of many, so no
+ /// cache may hold it.
+ fn samples(&self) -> bool {
+ matches!(self.alpha, Some(alpha) if alpha != 0.0)
+ }
+
+ /// What `sentence` encodes to, worked out rather than looked up: the lattice's
+ /// best path, or a draw from it when [`Unigram::samples`].
+ fn encode_uncached(&self, sentence: &str) -> Result> {
+ if self.is_optimized && !self.samples() {
+ self.encode_optimized(sentence)
} else {
- let result = self.encode_unoptimized(sentence)?;
- Ok(result)
+ self.encode_unoptimized(sentence)
}
}
@@ -503,7 +514,10 @@ impl Model for Unigram {
}
}
-pub struct UnigramScratch {}
+pub struct UnigramScratch {
+ /// Outlives the encode call that fills it, or it would never see a word twice.
+ pub(crate) word_cache: Option,
+}
impl pipeline::ModelScratch for UnigramScratch {}
@@ -511,18 +525,42 @@ impl pipeline::Model for Unigram {
type Scratch = UnigramScratch;
fn init_scratch(&self) -> Self::Scratch {
- Self::Scratch {}
+ Self::Scratch {
+ word_cache: match self.cache.capacity {
+ 0 => None,
+ capacity => Some(WordCache::new(capacity)),
+ },
+ }
}
+ /// The pipeline asks for ids and nothing else, so this path caches ids, where
+ /// [`Unigram::encode`] has to hand back the pieces themselves and caches those. A hit
+ /// skips the lattice, the `String` every piece is built into, and the vocabulary
+ /// lookup that turns each one back into an id.
fn tokenize_pipeline(
&self,
sequence: &str,
- _scratch: &mut Self::Scratch,
+ scratch: &mut Self::Scratch,
output: &mut Vec,
) -> Result<()> {
- let str_tokens = self.encode(sequence)?;
+ if sequence.is_empty() {
+ return Ok(());
+ }
+ let mut placement = None;
+ if !self.samples()
+ && let Some(cache) = scratch.word_cache.as_mut()
+ {
+ match cache.lookup(sequence.as_bytes()) {
+ Lookup::Hit(ids) => {
+ output.extend(ids.iter().map(|&id| PipelineToken { id }));
+ return Ok(());
+ }
+ Lookup::Miss(at) => placement = at,
+ }
+ }
- for string in str_tokens {
+ let start = output.len();
+ for string in self.encode_uncached(sequence)? {
match self.token_to_ids.token_to_id(&string) {
Some(id) => {
output.push(PipelineToken { id });
@@ -547,6 +585,14 @@ impl pipeline::Model for Unigram {
}
};
}
+
+ // Sampling skipped the lookup, so `placement` is `None` and this stores
+ // nothing, which is what `samples` says has to happen.
+ if let Some(cache) = scratch.word_cache.as_mut()
+ && let Some(at) = placement
+ {
+ cache.insert(at, output[start..].iter().map(|token| token.id));
+ }
Ok(())
}
}
@@ -710,4 +756,123 @@ mod tests {
let tokens = unigram.tokenize("?é").unwrap();
assert_eq!(tokens[0].id, 0);
}
+
+ /// Ids 0..=8 are ``, `a`, `b`, `c`, `d`, `cd`, `ab`, `abc`, `abcd`.
+ fn abcd_vocab() -> Vocab {
+ vec![
+ ("".to_string(), 0.0),
+ ("a".to_string(), 0.0),
+ ("b".to_string(), 0.0),
+ ("c".to_string(), 0.0),
+ ("d".to_string(), 0.0),
+ ("cd".to_string(), 1.0),
+ ("ab".to_string(), 2.0),
+ ("abc".to_string(), 5.0),
+ ("abcd".to_string(), 10.0),
+ ]
+ }
+
+ fn pipeline_ids(model: &Unigram, sequence: &str, scratch: &mut UnigramScratch) -> Vec {
+ let mut output = vec![];
+ pipeline::Model::tokenize_pipeline(model, sequence, scratch, &mut output).unwrap();
+ output.iter().map(|token| token.id).collect()
+ }
+
+ #[test]
+ fn pipeline_remembers_what_a_sequence_encoded_to() {
+ let model = Unigram::from(abcd_vocab(), Some(0), false).unwrap();
+ let mut scratch = pipeline::Model::init_scratch(&model);
+
+ let ids = pipeline_ids(&model, "abcd", &mut scratch);
+
+ let cache = scratch
+ .word_cache
+ .as_mut()
+ .expect("Unigram encodes with a cache");
+ assert_eq!(cache.lookup(b"abcd").hit(), Some(&ids[..]));
+ }
+
+ #[test]
+ fn cache_hits_agree_with_a_cold_run() {
+ let model = Unigram::from(abcd_vocab(), Some(0), false).unwrap();
+ let long = "abcd".repeat(400);
+ let corpus = [
+ "abcdacdxx",
+ "ab",
+ // The same sequence again, so this one is served from the cache.
+ "abcdacdxx",
+ // Out of the vocabulary, and multibyte.
+ "東京",
+ // 1600 bytes, past the longest word the cache will store.
+ long.as_str(),
+ "abcdacdxx",
+ ];
+
+ let mut warm_scratch = pipeline::Model::init_scratch(&model);
+ let warm = corpus.map(|sequence| pipeline_ids(&model, sequence, &mut warm_scratch));
+ let cold = corpus.map(|sequence| {
+ let mut scratch = pipeline::Model::init_scratch(&model);
+ pipeline_ids(&model, sequence, &mut scratch)
+ });
+
+ assert_eq!(warm, cold);
+ }
+
+ #[test]
+ fn caches_only_the_ids_this_sequence_produced() {
+ // Every sequence the pipeline hands the model appends to one output buffer,
+ // so a sequence has to remember its own ids, not everything the buffer holds.
+ let model = Unigram::from(abcd_vocab(), Some(0), false).unwrap();
+ let mut scratch = pipeline::Model::init_scratch(&model);
+ let mut output = vec![];
+ pipeline::Model::tokenize_pipeline(&model, "ab", &mut scratch, &mut output).unwrap();
+ pipeline::Model::tokenize_pipeline(&model, "cd", &mut scratch, &mut output).unwrap();
+
+ let ids: Vec = output.iter().map(|token| token.id).collect();
+ assert_eq!(ids, [6, 5]);
+ let cache = scratch.word_cache.as_mut().unwrap();
+ assert_eq!(cache.lookup(b"cd").hit(), Some(&[5u32][..]));
+ }
+
+ #[test]
+ fn byte_fallback_ids_survive_the_cache() {
+ // A piece the vocabulary has no id for becomes one id per byte. The cache
+ // stores what came out, so a hit has to replay all of them.
+ let vocab = vec![
+ ("".to_string(), 0.0),
+ ("<0xC3>".to_string(), -0.01),
+ ("<0xA9>".to_string(), -0.03),
+ ];
+ let model = Unigram::from(vocab, Some(0), true).unwrap();
+ let mut scratch = pipeline::Model::init_scratch(&model);
+
+ let ids = pipeline_ids(&model, "é", &mut scratch);
+
+ assert_eq!(ids, [1, 2]);
+ assert_eq!(pipeline_ids(&model, "é", &mut scratch), ids);
+ }
+
+ #[test]
+ fn sampling_is_never_cached() {
+ // A sampled tokenization is one draw out of many. Remembering it would turn
+ // every later call on the same text into that same draw.
+ let mut model = Unigram::from(abcd_vocab(), Some(0), false).unwrap();
+ model.alpha = Some(0.5);
+ let mut scratch = pipeline::Model::init_scratch(&model);
+
+ pipeline_ids(&model, "abcd", &mut scratch);
+
+ let cache = scratch.word_cache.as_mut().unwrap();
+ assert_eq!(cache.lookup(b"abcd").hit(), None);
+ }
+
+ #[test]
+ fn a_capacity_of_zero_turns_the_cache_off() {
+ let mut model = Unigram::from(abcd_vocab(), Some(0), false).unwrap();
+ model.resize_cache(0);
+ let mut scratch = pipeline::Model::init_scratch(&model);
+
+ assert_eq!(pipeline_ids(&model, "abcd", &mut scratch), [8]);
+ assert!(scratch.word_cache.is_none());
+ }
}
diff --git a/tokenizers/tk-encode/src/models/wordpiece/mod.rs b/tokenizers/tk-encode/src/models/wordpiece/mod.rs
index a1286fe95..4b9a02dd0 100644
--- a/tokenizers/tk-encode/src/models/wordpiece/mod.rs
+++ b/tokenizers/tk-encode/src/models/wordpiece/mod.rs
@@ -4,6 +4,8 @@
use crate::models::bpe::BPE;
use crate::pipeline::{self, PipelineToken};
use crate::tokenizer::{Model, Result, Token};
+use crate::utils::cache::DEFAULT_CACHE_CAPACITY;
+use crate::utils::word_cache::{Lookup, WordCache};
use ahash::AHashMap;
use std::collections::HashMap;
use std::convert::TryFrom;
@@ -316,6 +318,8 @@ impl Model for WordPiece {
pub struct WordPieceScratch {
candidate_str: String,
+ /// Outlives the encode call that fills it, or it would never see a word twice.
+ word_cache: WordCache,
}
impl pipeline::ModelScratch for WordPieceScratch {}
@@ -353,23 +357,18 @@ impl TryFrom for PipelineWordPiece {
}
}
-impl pipeline::Model for PipelineWordPiece {
- type Scratch = WordPieceScratch;
-
- fn init_scratch(&self) -> Self::Scratch {
- Self::Scratch {
- candidate_str: String::with_capacity(self.max_input_chars_per_word),
- }
- }
-
- fn tokenize_pipeline(
+impl PipelineWordPiece {
+ /// One word, greedily: the longest vocabulary entry it starts with, then the
+ /// longest entry the rest of it starts with once the continuing-subword prefix
+ /// is put in front, and so on. A piece with no entry at all anywhere in the
+ /// word makes the whole word one unk token.
+ fn tokenize_word(
&self,
sequence: &str,
- scratch: &mut Self::Scratch,
- output: &mut Vec,
+ candidate: &mut String,
+ output: &mut Vec,
) -> Result<()> {
let checkpoint = output.len();
- let candidate = &mut scratch.candidate_str;
let char_len = sequence.chars().count();
if char_len > self.max_input_chars_per_word {
@@ -411,6 +410,49 @@ impl pipeline::Model for PipelineWordPiece {
}
}
+impl pipeline::Model for PipelineWordPiece {
+ type Scratch = WordPieceScratch;
+
+ fn init_scratch(&self) -> Self::Scratch {
+ Self::Scratch {
+ candidate_str: String::with_capacity(self.max_input_chars_per_word),
+ word_cache: WordCache::new(DEFAULT_CACHE_CAPACITY),
+ }
+ }
+
+ /// A hit skips `tokenize_word`: one trie search per piece of the word,
+ /// each over a fresh copy of what is left to match.
+ fn tokenize_pipeline(
+ &self,
+ sequence: &str,
+ scratch: &mut Self::Scratch,
+ output: &mut Vec,
+ ) -> Result<()> {
+ if sequence.is_empty() {
+ return Ok(());
+ }
+ let WordPieceScratch {
+ candidate_str,
+ word_cache,
+ } = scratch;
+
+ let placement = match word_cache.lookup(sequence.as_bytes()) {
+ Lookup::Hit(ids) => {
+ output.extend(ids.iter().map(|&id| PipelineToken { id }));
+ return Ok(());
+ }
+ Lookup::Miss(at) => at,
+ };
+
+ let start = output.len();
+ self.tokenize_word(sequence, candidate_str, output)?;
+ if let Some(at) = placement {
+ word_cache.insert(at, output[start..].iter().map(|token| token.id));
+ }
+ Ok(())
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -419,4 +461,93 @@ mod tests {
fn test_error_display() {
assert!(format!("{}", Error::MissingUnkToken).contains("Missing [UNK] token"));
}
+
+ /// `hello` is in the vocabulary whole and as `hell` + `##o`, so the
+ /// longest-match walk has something to choose; `world` gives a second
+ /// one-token word.
+ fn pipeline_wordpiece() -> PipelineWordPiece {
+ let vocab: Vocab = [
+ ("[UNK]", 0u32),
+ ("hell", 1),
+ ("##o", 2),
+ ("hello", 3),
+ ("world", 4),
+ ]
+ .into_iter()
+ .map(|(token, id)| (token.to_string(), id))
+ .collect();
+ let model = WordPiece::builder()
+ .vocab(vocab)
+ .max_input_chars_per_word(8)
+ .build()
+ .unwrap();
+ PipelineWordPiece::try_from(model).unwrap()
+ }
+
+ fn pipeline_ids(
+ model: &PipelineWordPiece,
+ sequence: &str,
+ scratch: &mut WordPieceScratch,
+ ) -> Vec {
+ let mut output = vec![];
+ pipeline::Model::tokenize_pipeline(model, sequence, scratch, &mut output).unwrap();
+ output.iter().map(|token| token.id).collect()
+ }
+
+ #[test]
+ fn pipeline_remembers_what_a_word_encoded_to() {
+ let model = pipeline_wordpiece();
+ let mut scratch = pipeline::Model::init_scratch(&model);
+
+ let ids = pipeline_ids(&model, "hello", &mut scratch);
+
+ assert_eq!(scratch.word_cache.lookup(b"hello").hit(), Some(&ids[..]));
+ }
+
+ #[test]
+ fn cache_hits_agree_with_a_cold_run() {
+ let model = pipeline_wordpiece();
+ let long = "hello".repeat(300);
+ let corpus = [
+ "hello",
+ // No id for `##w`, so the whole word is one unk token.
+ "hellow",
+ // Both again, so these two are served from the cache.
+ "hello",
+ "hellow",
+ // Past `max_input_chars_per_word`, which is another unk token.
+ "hellohello",
+ "hellohello",
+ // Out of the vocabulary, and multibyte.
+ "東京",
+ "東京",
+ // 1500 bytes, past the longest word the cache will store.
+ long.as_str(),
+ long.as_str(),
+ ];
+
+ let mut warm_scratch = pipeline::Model::init_scratch(&model);
+ let warm = corpus.map(|sequence| pipeline_ids(&model, sequence, &mut warm_scratch));
+ let cold = corpus.map(|sequence| {
+ let mut scratch = pipeline::Model::init_scratch(&model);
+ pipeline_ids(&model, sequence, &mut scratch)
+ });
+
+ assert_eq!(warm, cold);
+ }
+
+ #[test]
+ fn caches_only_the_ids_this_word_produced() {
+ // Every word the pipeline hands the model appends to one output buffer,
+ // so a word has to remember its own ids, not everything the buffer holds.
+ let model = pipeline_wordpiece();
+ let mut scratch = pipeline::Model::init_scratch(&model);
+ let mut output = vec![];
+ pipeline::Model::tokenize_pipeline(&model, "hello", &mut scratch, &mut output).unwrap();
+ pipeline::Model::tokenize_pipeline(&model, "world", &mut scratch, &mut output).unwrap();
+
+ let ids: Vec = output.iter().map(|token| token.id).collect();
+ assert_eq!(ids, [3, 4]);
+ assert_eq!(scratch.word_cache.lookup(b"world").hit(), Some(&[4u32][..]));
+ }
}
diff --git a/tokenizers/tk-encode/src/normalizers/metaspace.rs b/tokenizers/tk-encode/src/normalizers/metaspace.rs
index b1a9486ce..b54e79a21 100644
--- a/tokenizers/tk-encode/src/normalizers/metaspace.rs
+++ b/tokenizers/tk-encode/src/normalizers/metaspace.rs
@@ -10,22 +10,46 @@
//! `to_normalizer_and_split`, in [`crate::pre_tokenizers::metaspace`], builds that pair and spells
//! out which [`Metaspace`] settings can be rebuilt this way.
//!
+//! Other configs spell the same rewrite in their `normalizer` field instead: llama-2 as
+//! `Prepend("▁")` followed by `Replace(" " -> "▁")`, gemma-4 as the `Replace` alone. Run as
+//! declared, every one of those steps writes a full copy of the text. `MetaspaceNormalizer::fuse`
+//! recognizes both shapes when the pipeline is built and stands in for them: one pass, one
+//! right-sized allocation.
+//!
//! [`Metaspace`]: crate::pre_tokenizers::metaspace::Metaspace
//! [`Split`]: crate::pre_tokenizers::split::Split
use std::borrow::Cow;
+use crate::normalizers::NormalizerWrapper;
+use crate::normalizers::replace::{Replace, ReplacePattern};
use crate::pre_tokenizers::whitespace::WhitespaceSplit;
+use crate::tokenizer::pipeline::PendingRewrite;
use crate::tokenizer::{Result, pipeline};
+/// When [`MetaspaceNormalizer`] writes a delimiter at the start of the text it is given.
+///
+/// The two prepending modes differ only on text that already starts with the delimiter (or with a
+/// space the swap turns into one): `IfMissing` skips those, which is how a [`Metaspace`]
+/// pre-tokenizer prepends; `Unconditional` marks them again, which is what a `Prepend` normalizer
+/// running before the swap does.
+///
+/// [`Metaspace`]: crate::pre_tokenizers::metaspace::Metaspace
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub(crate) enum PrependMode {
+ Never,
+ IfMissing,
+ Unconditional,
+}
+
/// Writes the delimiter where words start (after a space)
#[derive(Debug, Clone, PartialEq)]
pub struct MetaspaceNormalizer {
/// `▁` (U+2581) for every SentencePiece model we know of, but the config is free to use
/// another character.
delimiter: char,
- /// Write the delimiter at the start of every word, not only the words that followed a space.
- prepend: bool,
+ /// When to write the delimiter at the start of the text.
+ prepend: PrependMode,
/// Throw whitespace away instead of turning it into a delimiter: tabs, newlines and repeated
/// spaces leave no trace, and each word keeps only the one delimiter `prepend` writes. This is
/// the [`WhitespaceSplit`] that t5 and albert run in front of their `Metaspace`.
@@ -33,50 +57,258 @@ pub struct MetaspaceNormalizer {
}
impl MetaspaceNormalizer {
- pub(crate) fn new(delimiter: char, prepend: bool, drop_whitespace: bool) -> Self {
+ pub(crate) fn new(delimiter: char, prepend: PrependMode, drop_whitespace: bool) -> Self {
Self {
delimiter,
prepend,
drop_whitespace,
}
}
+
+ /// The swap and prepend as a [`PendingRewrite`]. The whitespace-dropping form deletes
+ /// characters, which no character-for-character rewrite can express, so it has none.
+ pub(crate) fn as_rewrite(&self) -> Option {
+ (!self.drop_whitespace).then_some(PendingRewrite {
+ from: ' ',
+ to: self.delimiter,
+ prepend: self.prepend,
+ })
+ }
+
+ /// The one-pass stand-in for the leading `steps`, when they spell out this normalizer's job:
+ /// `Prepend(c)` followed by `Replace(" " -> c)` (llama-2's normalizer), or a `Replace(" " -> c)`
+ /// on its own (gemma-4's). Returns the stand-in and how many steps it covers; `None` when the
+ /// steps start with anything else, and the caller keeps them as declared.
+ ///
+ /// [`PipelineTokenizer::try_from`] runs the config's normalizer steps through this when the
+ /// pipeline is built. The config itself is never rewritten.
+ ///
+ /// [`PipelineTokenizer::try_from`]: crate::pipeline::PipelineTokenizer
+ pub(crate) fn fuse(steps: &[NormalizerWrapper]) -> Option<(Self, usize)> {
+ if let [
+ NormalizerWrapper::Prepend(prepend),
+ NormalizerWrapper::Replace(replace),
+ ..,
+ ] = steps
+ && let Some(delimiter) = space_swap(replace)
+ && prepend.prepend == replace.content
+ {
+ return Some((Self::new(delimiter, PrependMode::Unconditional, false), 2));
+ }
+ if let [NormalizerWrapper::Replace(replace), ..] = steps
+ && let Some(delimiter) = space_swap(replace)
+ {
+ return Some((Self::new(delimiter, PrependMode::Never, false), 1));
+ }
+ None
+ }
+}
+
+/// The delimiter a [`Replace`] swaps every space for: its pattern must be the string `" "` and its
+/// content a single char. Anything else (a regex pattern, even one only matching a space; content
+/// of any other length) is not checkable as the space swap and returns `None`.
+fn space_swap(replace: &Replace) -> Option {
+ match &replace.pattern {
+ ReplacePattern::String(pattern) if pattern == " " => {}
+ _ => return None,
+ }
+ let mut chars = replace.content.chars();
+ let delimiter = chars.next()?;
+ chars.next().is_none().then_some(delimiter)
}
impl pipeline::Normalizer for MetaspaceNormalizer {
+ fn pending_rewrite(&self) -> Option {
+ self.as_rewrite()
+ }
+
fn normalize<'a>(&self, input: &'a str) -> Result> {
+ // Only spaces become delimiters; tabs and newlines are left alone. The swap and the
+ // prepend are exactly a [`PendingRewrite`], written here.
+ if let Some(rewrite) = self.as_rewrite() {
+ return rewrite.write(input);
+ }
// Return empty input as is
if input.is_empty() {
return Ok(Cow::Borrowed(input));
}
- // The delimiter is 3 bytes where a space is 1, so the rewrite grows by 2 bytes per space, hence we allocate a bit more space
- let mut rewritten = String::with_capacity(input.len() + input.len() / 2);
- if self.drop_whitespace {
- // Whitespace is thrown away, so cut the text where `WhitespaceSplit` would and write the
- // words back one after the other, each with its own delimiter.
- let mut words = Vec::new();
- pipeline::PreTokenizer::pre_tokenize(&WhitespaceSplit, input, &mut words)?;
- for span in &words {
- let word = &input[span.range()];
- // The text may already hold delimiters of its own — never write a second one.
- if self.prepend && !word.starts_with(self.delimiter) {
- rewritten.push(self.delimiter);
- }
- rewritten.push_str(word);
- }
- } else {
- // Prepend the delimiter if self.prepend is true
- if self.prepend && !input.starts_with(' ') && !input.starts_with(self.delimiter) {
- rewritten.push(self.delimiter);
- }
- // Only spaces become delimiters; tabs and newlines are left alone
- let mut rest = input;
- while let Some(space) = memchr::memchr(b' ', rest.as_bytes()) {
- rewritten.push_str(&rest[..space]);
+ // Whitespace is thrown away, so cut the text where `WhitespaceSplit` would and write the
+ // words back one after the other, each with its own delimiter.
+ let mut words = Vec::new();
+ pipeline::PreTokenizer::pre_tokenize(&WhitespaceSplit, input, &mut words)?;
+ // Exact when every word takes a delimiter; a word that already starts with one
+ // (`IfMissing`) leaves a few bytes spare.
+ let words_len: usize = words.iter().map(|span| span.range().len()).sum();
+ let mut rewritten =
+ String::with_capacity(words_len + words.len() * self.delimiter.len_utf8());
+ for span in &words {
+ let word = &input[span.range()];
+ let prepend = match self.prepend {
+ PrependMode::Never => false,
+ // The text may already hold delimiters of its own: never write a second one.
+ PrependMode::IfMissing => !word.starts_with(self.delimiter),
+ PrependMode::Unconditional => true,
+ };
+ if prepend {
rewritten.push(self.delimiter);
- rest = &rest[space + 1..];
}
- rewritten.push_str(rest);
+ rewritten.push_str(word);
}
Ok(Cow::Owned(rewritten))
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::tokenizer::pipeline::normalize_all;
+
+ fn step(json: &str) -> NormalizerWrapper {
+ serde_json::from_str(json).unwrap()
+ }
+
+ /// llama-2's two normalizer steps.
+ const PREPEND: &str = r#"{"type":"Prepend","prepend":"▁"}"#;
+ /// gemma-4's whole normalizer, and llama-2's second step.
+ const SWAP: &str = r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#;
+
+ /// Every chunk shape the rewrite can meet: empty, already marked, leading, trailing and
+ /// repeated spaces, no spaces at all, other whitespace.
+ const TEXTS: &[&str] = &[
+ "",
+ "hello world",
+ "hello world",
+ " leading",
+ "trailing ",
+ " both ",
+ "▁already marked",
+ "▁",
+ "no_spaces",
+ "one\ttab\nand a newline",
+ " ",
+ "a▁b c",
+ ];
+
+ /// The stand-in must rewrite every text exactly as the steps it covers would.
+ fn assert_fused_matches_steps(jsons: &[&str]) {
+ let steps: Vec = jsons.iter().map(|json| step(json)).collect();
+ let (fused, covered) = MetaspaceNormalizer::fuse(&steps).expect("this shape fuses");
+ assert_eq!(covered, steps.len());
+ for text in TEXTS {
+ let expected = normalize_all(&steps, text).unwrap();
+ let got = pipeline::Normalizer::normalize(&fused, text).unwrap();
+ assert_eq!(got, expected, "{text:?}");
+ }
+ }
+
+ #[test]
+ fn fused_prepend_and_swap_match_the_two_steps() {
+ assert_fused_matches_steps(&[PREPEND, SWAP]);
+ }
+
+ #[test]
+ fn fused_swap_matches_the_lone_step() {
+ assert_fused_matches_steps(&[SWAP]);
+ }
+
+ /// Nothing ties the delimiter to `▁`; any single char must fuse the same way.
+ #[test]
+ fn fused_ascii_delimiter_matches_its_steps() {
+ assert_fused_matches_steps(&[
+ r#"{"type":"Prepend","prepend":"_"}"#,
+ r#"{"type":"Replace","pattern":{"String":" "},"content":"_"}"#,
+ ]);
+ }
+
+ /// The pair stands in for a `Prepend`, which marks a chunk even when it already starts with
+ /// the delimiter; the lone swap never prepends. `covered` tells the caller how many steps to
+ /// skip.
+ #[test]
+ fn the_pair_prepends_unconditionally_and_the_lone_swap_never_does() {
+ let steps = [step(PREPEND), step(SWAP)];
+ let (fused, covered) = MetaspaceNormalizer::fuse(&steps).unwrap();
+ assert_eq!(covered, 2);
+ assert_eq!(
+ fused,
+ MetaspaceNormalizer::new('▁', PrependMode::Unconditional, false)
+ );
+ let (fused, covered) = MetaspaceNormalizer::fuse(&steps[1..]).unwrap();
+ assert_eq!(covered, 1);
+ assert_eq!(
+ fused,
+ MetaspaceNormalizer::new('▁', PrependMode::Never, false)
+ );
+ }
+
+ /// A chunk with nothing to swap is handed back without a copy, as `Replace` hands it back.
+ #[test]
+ fn a_swap_with_no_spaces_borrows() {
+ let steps = [step(SWAP)];
+ let (fused, _) = MetaspaceNormalizer::fuse(&steps).unwrap();
+ assert!(matches!(
+ pipeline::Normalizer::normalize(&fused, "no_spaces").unwrap(),
+ Cow::Borrowed(_)
+ ));
+ }
+
+ #[test]
+ fn refuses_a_replace_that_is_not_the_space_swap() {
+ let refused = [
+ // Swaps something else entirely; running it as the space swap would corrupt the text.
+ (
+ "another literal",
+ r#"{"type":"Replace","pattern":{"String":"x"},"content":"y"}"#,
+ ),
+ // The delimiter is one char; two of them per space is not this normalizer's rewrite.
+ (
+ "multi-char content",
+ r#"{"type":"Replace","pattern":{"String":" "},"content":"▁▁"}"#,
+ ),
+ // Deletes spaces instead of marking them.
+ (
+ "empty content",
+ r#"{"type":"Replace","pattern":{"String":" "},"content":""}"#,
+ ),
+ // Only single spaces are swapped one-for-one.
+ (
+ "a two-space pattern",
+ r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#,
+ ),
+ ];
+ for (name, json) in refused {
+ assert!(MetaspaceNormalizer::fuse(&[step(json)]).is_none(), "{name}");
+ }
+ }
+
+ /// A regex spelling is refused even when it happens to match only a space: fusing checks the
+ /// pattern structurally and does not interpret regex syntax.
+ #[cfg(feature = "fancy-regex")]
+ #[test]
+ fn refuses_a_regex_pattern_even_one_matching_a_space() {
+ let json = r#"{"type":"Replace","pattern":{"Regex":" "},"content":"▁"}"#;
+ assert!(MetaspaceNormalizer::fuse(&[step(json)]).is_none());
+ }
+
+ /// A `Prepend` writing anything but the swap's delimiter is a different rewrite; the pair must
+ /// not fuse. The swap after it still can, on its own, once the caller reaches it.
+ #[test]
+ fn refuses_a_prepend_of_something_else() {
+ for prepend in [
+ r#"{"type":"Prepend","prepend":"x"}"#,
+ r#"{"type":"Prepend","prepend":"▁▁"}"#,
+ ] {
+ let steps = [step(prepend), step(SWAP)];
+ assert!(MetaspaceNormalizer::fuse(&steps).is_none(), "{prepend}");
+ assert!(MetaspaceNormalizer::fuse(&steps[1..]).is_some());
+ }
+ }
+
+ /// The pair only fuses when the two steps are adjacent: another step in between sees the text
+ /// mid-rewrite, and fusing across it would change what that step sees.
+ #[test]
+ fn refuses_a_separated_pair() {
+ let steps = [step(PREPEND), step(r#"{"type":"Lowercase"}"#), step(SWAP)];
+ assert!(MetaspaceNormalizer::fuse(&steps).is_none());
+ assert!(MetaspaceNormalizer::fuse(&steps[2..]).is_some());
+ }
+}
diff --git a/tokenizers/tk-encode/src/normalizers/mod.rs b/tokenizers/tk-encode/src/normalizers/mod.rs
index c16751ff4..effe8d888 100644
--- a/tokenizers/tk-encode/src/normalizers/mod.rs
+++ b/tokenizers/tk-encode/src/normalizers/mod.rs
@@ -237,6 +237,25 @@ impl pipeline::Normalizer for NormalizerWrapper {
Self::ByteLevel(bl) => pipeline::Normalizer::normalize(bl, input),
}
}
+
+ fn pending_rewrite(&self) -> Option {
+ match self {
+ Self::BertNormalizer(bn) => pipeline::Normalizer::pending_rewrite(bn),
+ Self::StripNormalizer(sn) => pipeline::Normalizer::pending_rewrite(sn),
+ Self::StripAccents(sn) => pipeline::Normalizer::pending_rewrite(sn),
+ Self::NFC(nfc) => pipeline::Normalizer::pending_rewrite(nfc),
+ Self::NFD(nfd) => pipeline::Normalizer::pending_rewrite(nfd),
+ Self::NFKC(nfkc) => pipeline::Normalizer::pending_rewrite(nfkc),
+ Self::NFKD(nfkd) => pipeline::Normalizer::pending_rewrite(nfkd),
+ Self::Sequence(sequence) => pipeline::Normalizer::pending_rewrite(sequence),
+ Self::Lowercase(lc) => pipeline::Normalizer::pending_rewrite(lc),
+ Self::Nmt(nmt) => pipeline::Normalizer::pending_rewrite(nmt),
+ Self::Precompiled(pc) => pipeline::Normalizer::pending_rewrite(pc),
+ Self::Replace(rp) => pipeline::Normalizer::pending_rewrite(rp),
+ Self::Prepend(pp) => pipeline::Normalizer::pending_rewrite(pp),
+ Self::ByteLevel(bl) => pipeline::Normalizer::pending_rewrite(bl),
+ }
+ }
}
#[cfg(test)]
diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs
index cc4545b47..872a64308 100644
--- a/tokenizers/tk-encode/src/normalizers/replace.rs
+++ b/tokenizers/tk-encode/src/normalizers/replace.rs
@@ -1,5 +1,6 @@
use std::borrow::Cow;
+use crate::normalizers::metaspace::PrependMode;
use crate::pipeline;
use crate::tokenizer::Decoder;
use crate::tokenizer::pattern::Pattern;
@@ -70,7 +71,7 @@ impl Search {
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", try_from = "ReplaceDeserializer")]
pub struct Replace {
- pattern: ReplacePattern,
+ pub(crate) pattern: ReplacePattern,
pub content: String,
#[serde(skip)]
search: Search,
@@ -103,6 +104,11 @@ impl Replace {
search,
})
}
+
+ /// What this normalizer looks for.
+ pub fn pattern(&self) -> &ReplacePattern {
+ &self.pattern
+ }
}
impl Normalizer for Replace {
@@ -141,14 +147,43 @@ fn replace_matches<'a>(
}
impl pipeline::Normalizer for Replace {
+ /// A swap of one character for another qualifies. Regex patterns, multi-character
+ /// literals and deletions change the text's structure and have to be written.
+ fn pending_rewrite(&self) -> Option {
+ let ReplacePattern::String(pattern) = &self.pattern else {
+ return None;
+ };
+ let mut pattern = pattern.chars();
+ let from = pattern.next().filter(|_| pattern.next().is_none())?;
+ let mut content = self.content.chars();
+ let to = content.next().filter(|_| content.next().is_none())?;
+ Some(pipeline::PendingRewrite {
+ from,
+ to,
+ prepend: PrependMode::Never,
+ })
+ }
+
fn normalize<'a>(&self, input: &'a str) -> Result> {
Ok(match &self.search {
Search::Literal(literal) => {
let width = literal.pattern().len();
- let matches = literal
- .matches(input.as_bytes())
- .map(|start| (start, start + width));
- replace_matches(input, &self.content, matches)
+ let count = literal.count_matches(input.as_bytes());
+ if count == 0 {
+ return Ok(Cow::Borrowed(input));
+ }
+ // Counting first buys the exact output length (every match swaps `width`
+ // bytes for the content), so the build below never reallocates.
+ let exact = input.len() - count * width + count * self.content.len();
+ let mut replaced = String::with_capacity(exact);
+ let mut last_end = 0;
+ literal.for_each_match(input.as_bytes(), |start| {
+ replaced.push_str(&input[last_end..start]);
+ replaced.push_str(&self.content);
+ last_end = start + width;
+ });
+ replaced.push_str(&input[last_end..]);
+ Cow::Owned(replaced)
}
Search::Regex(regex) => replace_matches(input, &self.content, regex.find_iter(input)),
Search::Nothing => Cow::Borrowed(input),
diff --git a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs
index e30f206e6..4e2ec8f40 100644
--- a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs
+++ b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs
@@ -1,3 +1,4 @@
+use atomsplit::literal::Literal;
use serde::{Deserialize, Serialize};
use crate::pipeline;
@@ -28,16 +29,18 @@ impl PreTokenizer for CharDelimiterSplit {
impl pipeline::PreTokenizer for CharDelimiterSplit {
fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> {
- // native atomsplit FSM (memchr-backed single-byte scan); `Removed` — drops the delimiter,
- // keeps the runs between, no empty spans. Byte-exact with the char-predicate split.
- let bytes = text.as_bytes();
- let mut spans = vec![pipeline::Span::default(); bytes.len() + 1];
- let n = atomsplit::fsm::CharDelimiterSplit(self.delimiter).pre_tokenize(
- bytes,
- &mut [],
- &mut spans,
+ // The same streaming literal cut as a `Split` on this char: `Removed` drops each
+ // delimiter and keeps the runs between them, with no empty spans.
+ let mut buf = [0u8; 4];
+ let delimiter = Literal::new(self.delimiter.encode_utf8(&mut buf).as_bytes())
+ .expect("a char is never empty");
+ pipeline::split_literal(
+ out,
+ &delimiter,
+ text,
+ SplitDelimiterBehavior::Removed,
+ false,
);
- out.extend_from_slice(&spans[..n]);
Ok(())
}
}
diff --git a/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs b/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs
index 3d940b475..9afd1d172 100644
--- a/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs
+++ b/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs
@@ -1,4 +1,4 @@
-use crate::normalizers::metaspace::MetaspaceNormalizer;
+use crate::normalizers::metaspace::{MetaspaceNormalizer, PrependMode};
use crate::pre_tokenizers::PreTokenizerWrapper;
use crate::pre_tokenizers::split::Split;
use crate::tokenizer::{Decoder, PreTokenizedString, PreTokenizer, Result, SplitDelimiterBehavior};
@@ -212,15 +212,17 @@ fn normalizer_and_split(
return None;
}
let prepend = match metaspace.prepend_scheme {
- PrependScheme::Always => true,
- PrependScheme::Never => false,
+ // `Metaspace` swaps spaces first and then marks only text not already starting with the
+ // delimiter, so its `always` scheme is the conditional prepend.
+ PrependScheme::Always => PrependMode::IfMissing,
+ PrependScheme::Never => PrependMode::Never,
// `First` writes the delimiter only on the piece at the very start of the text it came from.
// A normalizer is handed one chunk at a time, without that context.
PrependScheme::First => return None,
};
// Removes whitespaces and does not prepend words: nothing would show where words begin
- // The output is one big continuous blob of words hlued together
- if drop_whitespace && !prepend {
+ // The output is one big continuous blob of words glued together
+ if drop_whitespace && prepend == PrependMode::Never {
return None;
}
let delimiter = metaspace.replacement;
diff --git a/tokenizers/tk-encode/src/pre_tokenizers/mod.rs b/tokenizers/tk-encode/src/pre_tokenizers/mod.rs
index b895e3f29..757cf4ceb 100644
--- a/tokenizers/tk-encode/src/pre_tokenizers/mod.rs
+++ b/tokenizers/tk-encode/src/pre_tokenizers/mod.rs
@@ -4,6 +4,7 @@ pub mod delimiter;
pub mod digits;
pub mod fixed_length;
pub mod metaspace;
+pub mod proven_cuts;
pub mod punctuation;
pub mod sequence;
pub mod split;
diff --git a/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs b/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs
new file mode 100644
index 000000000..0e5dcd443
--- /dev/null
+++ b/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs
@@ -0,0 +1,756 @@
+//! Cutting a whole SentencePiece text into words, where the vocabulary proves it is harmless.
+//!
+//! SentencePiece vocabularies (gemma, llama-2, …) write a space as `▁` (U+2581), so `"tell me"`
+//! reaches the model as `"▁tell▁me"`. That character both opens a word and separates two, which is
+//! why the code calls it the delimiter.
+//!
+//! Some of these tokenizers ship a [`Metaspace`] pre-tokenizer, which writes the delimiters and cuts
+//! before every one of them. Those cuts are the tokenizer's own output, so
+//! [`super::metaspace::to_normalizer_and_split`] reproduces them exactly and there is nothing to
+//! decide.
+//!
+//! The tokenizers here are the other kind: their normalizer writes the delimiters and they ship no
+//! pre-tokenizer that cuts, so the model receives the whole text in one piece. Cutting it into words
+//! is only a speed-up — merging a long text costs more than merging its words one after another, and
+//! short words are friendlier to the cache. But a speed-up is worthless if it changes the tokens, so
+//! here every cut has to be proven harmless first. Two rules do that:
+//!
+//! 1. A group of delimiters stays with the word that follows it (`a▁▁▁b` → `a`, `▁▁▁b`).
+//! Vocabularies hold pieces made of several delimiters (`▁▁`, `▁▁▁`), and cutting inside a group
+//! would stop those from forming.
+//! 2. A cut is dropped when a vocabulary piece could merge across it. A merge only ever produces a
+//! piece that is in the vocabulary, so this can only happen if some piece holds a delimiter that
+//! is not at its start. gemma has exactly one such piece: `>▁ bool {
+ self.delimiter.pattern() == other.delimiter.pattern() && self.veto == other.veto
+ }
+}
+
+impl ProvenCuts {
+ fn new(delimiter: Literal, veto: Veto) -> Self {
+ Self { delimiter, veto }
+ }
+}
+
+impl pipeline::PreTokenizer for ProvenCuts {
+ fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> {
+ let delimiter = self.delimiter.pattern();
+ let bytes = text.as_bytes();
+ // The word being built runs from `start` to the next delimiter we cut at.
+ let mut start = 0usize;
+ for at in self.delimiter.matches(bytes) {
+ // A delimiter at `start` is the cut we just made, or the text opens with one: either way
+ // the word would be empty.
+ if at == start {
+ continue;
+ }
+ // A delimiter right in front of this one means we are inside a group of them, and a group
+ // stays whole with the word after it.
+ if bytes[..at].ends_with(delimiter) || self.veto.forbids(bytes, at) {
+ continue;
+ }
+ out.push(Span::new(start as u32, at as u32));
+ start = at;
+ }
+ if start < bytes.len() {
+ out.push(Span::new(start as u32, bytes.len() as u32));
+ }
+ Ok(())
+ }
+}
+
+/// The splitter for a tokenizer whose text reaches the model in one piece, or `None` when the text
+/// has to stay that way — because something already cuts it, because no normalizer provably writes
+/// the delimiters, or because the vocabulary cannot prove the cuts.
+pub(crate) fn for_tokenizer(
+ normalizer: Option<&NormalizerWrapper>,
+ pre_tokenizer: Option<&PreTokenizerWrapper>,
+ model: &PipelineModel,
+) -> Option {
+ if !leaves_the_text_whole(pre_tokenizer) {
+ return None;
+ }
+ let delimiter = delimiter_from_normalizer(normalizer?)?;
+ let delimiter = Literal::new(delimiter.to_string().as_bytes()).expect("a char is never empty");
+ let veto = veto_from_model(model, &delimiter)?;
+ Some(ProvenCuts::new(delimiter, veto))
+}
+
+/// Does this pre-tokenizer leave the text in one piece? Only the two shapes below do, and anything
+/// else is refused rather than guessed at: cutting text a tokenizer meant to keep whole changes the
+/// tokens it produces.
+fn leaves_the_text_whole(pre_tokenizer: Option<&PreTokenizerWrapper>) -> bool {
+ match pre_tokenizer {
+ // Nothing cuts the text at all. llama-2 ships this shape.
+ None => true,
+ // A `Split` that can never match, because the normalizer already replaced every space it
+ // looks for. With nothing to match, every behaviour it could carry leaves the text in one
+ // piece. gemma ships this shape.
+ Some(PreTokenizerWrapper::Split(split)) => {
+ !split.invert
+ && matches!(&split.pattern, SplitPattern::String(pattern)
+ if !pattern.is_empty() && pattern.chars().all(|c| c == ' '))
+ }
+ _ => false,
+ }
+}
+
+/// The cuts this model's vocabulary forbids, or `None` when it cannot be cut at all.
+fn veto_from_model(model: &PipelineModel, delimiter: &Literal) -> Option {
+ let PipelineModel::BPE(bpe) = model else {
+ return None;
+ };
+ // With `ignore_merges` the model first looks the whole text up in the vocabulary and emits one
+ // token when it finds it. Handing it words instead would skip that lookup.
+ if bpe.ignore_merges() {
+ return None;
+ }
+ Veto::build(&bpe.vocab_bytes(), delimiter)
+}
+
+/// The character every space becomes, if the normalizer provably rewrites all of them.
+///
+/// Accepts a `Replace` on its own, or a sequence whose last step is one and whose earlier steps only
+/// prepend text — a step running after the `Replace` could bring spaces back.
+fn delimiter_from_normalizer(normalizer: &NormalizerWrapper) -> Option {
+ match normalizer {
+ NormalizerWrapper::Replace(replace) => space_replacement(replace),
+ NormalizerWrapper::Sequence(sequence) => {
+ let (last, rest) = sequence.as_ref().split_last()?;
+ rest.iter()
+ .all(|step| matches!(step, NormalizerWrapper::Prepend(_)))
+ .then(|| match last {
+ NormalizerWrapper::Replace(replace) => space_replacement(replace),
+ _ => None,
+ })
+ .flatten()
+ }
+ _ => None,
+ }
+}
+
+/// The single character this normalizer turns every space into.
+fn space_replacement(replace: &Replace) -> Option {
+ if replace.pattern() != &ReplacePattern::String(" ".to_string()) {
+ return None;
+ }
+ let mut content = replace.content.chars();
+ let delimiter = content.next()?;
+ content.next().is_none().then_some(delimiter)
+}
+
+/// [`ProvenCuts`] without the rewrite it runs on: the same cuts, found on the raw text.
+///
+/// The tokenizers this serves normalize with a [`PendingRewrite`], swapping every space for
+/// the delimiter (llama-2 prepends one too), so their rewritten text differs from the raw
+/// text one character at a time. Every position [`ProvenCuts`] would cut is therefore
+/// visible in the raw text: it is a raw `from` or a raw `to`. Cutting there directly means
+/// the rewrite is never written; the model reads each raw span through a [`CharSwap`]
+/// ([`PipelineBPE::tokenize_swapped`]), and only the one word a prepend touches is rewritten
+/// for real.
+///
+/// The veto is taken at its prefilter's word: a cut whose preceding byte could open a veto
+/// piece is skipped without checking the piece itself. Skipping a cut never changes the ids
+/// (each cut is only ever a speed-up), it only hands the model a longer span. The full check
+/// reads the rewritten bytes around the cut, which is exactly what this path avoids building.
+#[derive(Debug, Clone)]
+pub(crate) struct ZeroCopyMetaspace {
+ /// The rewrite this path stands in for; also writes the one span the prepend touches.
+ rewrite: PendingRewrite,
+ /// The rewrite's `from`, as the single byte it encodes to. The cuts are only proven for
+ /// the space swap, so `from` is always one byte; comparing a byte (not a runtime-length
+ /// slice, which compiles to a `memcmp` call) is what keeps the scan tight.
+ from: u8,
+ /// UTF-8 bytes of the delimiter (the rewrite's `to`); only the first `to_len` are meaningful.
+ to: [u8; 4],
+ to_len: u8,
+ /// The delimiter's vocabulary id, seeded for every raw `from`.
+ to_id: u32,
+ /// [`Veto::bytes_before`]: last bytes of the pieces' `before` halves, as a 256-bit set.
+ veto_bytes_before: [u64; 4],
+}
+
+impl ZeroCopyMetaspace {
+ /// `None` when the raw text cannot stand in for the rewritten text: the cuts were proven
+ /// for the space swap only, and byte-atom models spell text in another alphabet.
+ pub(crate) fn build(
+ rewrite: PendingRewrite,
+ cuts: &ProvenCuts,
+ bpe: &PipelineBPE,
+ ) -> Option {
+ if rewrite.from != ' ' || !bpe.char_atoms() {
+ return None;
+ }
+ let mut to = [0u8; 4];
+ let encoded = rewrite.to.encode_utf8(&mut to);
+ if cuts.delimiter.pattern() != encoded.as_bytes() {
+ return None;
+ }
+ let to_id = bpe.id_of_bytes(encoded.as_bytes())?;
+ let to_len = encoded.len() as u8;
+ Some(Self {
+ rewrite,
+ from: rewrite.from as u8,
+ to,
+ to_len,
+ to_id,
+ veto_bytes_before: cuts.veto.bytes_before,
+ })
+ }
+
+ pub(crate) fn rewrite(&self) -> PendingRewrite {
+ self.rewrite
+ }
+
+ pub(crate) fn swap(&self) -> CharSwap {
+ CharSwap {
+ from: self.from,
+ id: self.to_id,
+ len: self.to_len,
+ }
+ }
+
+ /// Cuts `text` where [`ProvenCuts`] would cut its rewritten form, as raw byte offsets.
+ /// Returns whether the first span takes the prepended delimiter; the caller writes that
+ /// one span's rewrite for real and reads the rest through the swap.
+ pub(crate) fn cut(&self, text: &str, out: &mut Vec) -> bool {
+ let bytes = text.as_bytes();
+ let to = &self.to[..self.to_len as usize];
+ let mut start = 0usize;
+ let mut search = 0usize;
+ while let Some(found) = memchr::memchr2(self.from, to[0], &bytes[search..]) {
+ let at = search + found;
+ let width = if bytes[at] == self.from {
+ 1
+ } else if bytes[at..].starts_with(to) {
+ to.len()
+ } else {
+ // The delimiter's first byte opens other characters too; not a delimiter.
+ search = at + 1;
+ continue;
+ };
+ // The same three refusals as [`ProvenCuts::pre_tokenize`]: the word would be
+ // empty, the delimiter sits inside a group (its predecessor is a raw `from` or
+ // a raw `to`), or the veto's prefilter byte fires. The predecessor byte is
+ // never rewritten (a rewritten one is the group case), so raw is exact here.
+ if at != start
+ && !(bytes[at - 1] == self.from || bytes[..at].ends_with(to))
+ && !self.prefilter_forbids(bytes[at - 1])
+ {
+ out.push(Span::new(start as u32, at as u32));
+ start = at;
+ }
+ search = at + width;
+ }
+ if start < bytes.len() {
+ out.push(Span::new(start as u32, bytes.len() as u32));
+ }
+ self.rewrite.prepends(text)
+ }
+
+ fn prefilter_forbids(&self, previous: u8) -> bool {
+ self.veto_bytes_before[(previous >> 6) as usize] & (1 << (previous & 63)) != 0
+ }
+}
+
+/// How many veto pieces we put up with before giving up on cutting at all.
+const MAX_VETO_PIECES: usize = 32;
+
+/// Width of one padded half of a veto piece. Whatever length the halves really are, they are compared
+/// a full `u128` at a time.
+const HALF_WIDTH: usize = size_of::();
+
+/// One half of a [`VetoPiece`], padded out to [`HALF_WIDTH`] bytes.
+///
+/// The halves are a byte or two long, but their length is only known once the vocabulary is read, and
+/// comparing a run-time number of bytes means calling `memcmp`. Padding to a fixed width instead
+/// turns the compare into a couple of register operations, which is worth it in a loop that runs once
+/// per word.
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct Half {
+ bytes: u128,
+ /// `0xff` over the bytes of the half, `0` over the padding.
+ mask: u128,
+}
+
+impl Half {
+ /// The half at the end of the window, where the bytes running up to a cut land.
+ fn ending(half: &[u8]) -> Option {
+ Self::padded(half, HALF_WIDTH.checked_sub(half.len())?)
+ }
+
+ /// The half at the start of the window, where the bytes following a cut land.
+ fn starting(half: &[u8]) -> Option {
+ Self::padded(half, 0)
+ }
+
+ /// `None` when the half is wider than the window.
+ fn padded(half: &[u8], at: usize) -> Option {
+ let mut bytes = [0u8; HALF_WIDTH];
+ let mut mask = [0u8; HALF_WIDTH];
+ bytes.get_mut(at..at + half.len())?.copy_from_slice(half);
+ mask[at..at + half.len()].fill(0xff);
+ Some(Self {
+ bytes: u128::from_le_bytes(bytes),
+ mask: u128::from_le_bytes(mask),
+ })
+ }
+
+ fn matches(&self, window: u128) -> bool {
+ (window ^ self.bytes) & self.mask == 0
+ }
+}
+
+/// A vocabulary piece holding a delimiter that is not at its start, split at that delimiter: `>▁", after: " Option {
+ Some(Self {
+ before: Half::ending(before)?,
+ after: Half::starting(after)?,
+ })
+ }
+}
+
+/// The cuts a vocabulary does not allow.
+///
+/// A merge only ever produces a piece that is in the vocabulary, so a merge can only reach across a
+/// cut if some piece holds a delimiter that is not at its start. Those pieces are collected here, and
+/// a cut where one of them fits is dropped.
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct Veto {
+ /// Empty means no merge can reach across a cut, so every cut is allowed.
+ pieces: Vec,
+ /// Last bytes of the `before` halves, as a 256-bit set. A cut with any other byte in front of it
+ /// fits no piece, which is how most cuts get away without a single comparison.
+ bytes_before: [u64; 4],
+ /// Length of the delimiter the pieces were split at. Every `after` half starts at the byte right
+ /// behind the delimiter, so the window compared to it has to start there too.
+ delimiter_len: usize,
+}
+
+impl Veto {
+ /// Reads the vocabulary and collects the pieces a merge could use to reach across a cut.
+ ///
+ /// `None` turns cutting off: either this is not a SentencePiece vocabulary (no delimiter piece in
+ /// it), or its veto pieces are too many, too long or too tangled to rule out cheaply.
+ fn build(vocab: &[(Vec, u32)], delimiter: &Literal) -> Option {
+ let pattern = delimiter.pattern();
+ if !vocab.iter().any(|(piece, _)| piece.as_slice() == pattern) {
+ return None;
+ }
+ let mut pieces = Vec::new();
+ let mut bytes_before = [0u64; 4];
+ for (piece, _) in vocab {
+ for at in delimiter.matches(piece) {
+ // A piece starting with a delimiter sits right after a cut, not across it, and a
+ // delimiter following another one is inside a group, where we never cut.
+ if at == 0 || piece[..at].ends_with(pattern) {
+ continue;
+ }
+ let (before, after) = (&piece[..at], &piece[at + pattern.len()..]);
+ // A piece with a second delimiter reaches across two cuts at once, and checking one
+ // cut at a time no longer proves anything. None of the vocabularies we tested has
+ // one, so drop cutting instead.
+ if pieces.len() == MAX_VETO_PIECES
+ || delimiter.matches(before).next().is_some()
+ || delimiter.matches(after).next().is_some()
+ {
+ return None;
+ }
+ let previous = *before.last().expect("`at` is past the start of the piece");
+ bytes_before[(previous >> 6) as usize] |= 1 << (previous & 63);
+ pieces.push(VetoPiece::new(before, after)?);
+ }
+ }
+ Some(Self {
+ pieces,
+ bytes_before,
+ delimiter_len: pattern.len(),
+ })
+ }
+
+ /// Could a piece cover the delimiter at `at`, so that a merge reaches over a cut placed there? A
+ /// match only means such a merge is possible, not that the model performs it — either way we
+ /// leave the text in one piece. `at` is past the start of the text, so there is a byte in front
+ /// of it.
+ fn forbids(&self, text: &[u8], at: usize) -> bool {
+ let previous = text[at - 1];
+ if self.bytes_before[(previous >> 6) as usize] & (1 << (previous & 63)) == 0 {
+ return false;
+ }
+ // Both windows are padded with zeros, so a half can only match beyond the ends of `text` if
+ // the half itself holds a zero byte — and one match too many only leaves the text uncut.
+ let before = window_ending(&text[..at]);
+ let after = window_starting(&text[at + self.delimiter_len..]);
+ self.pieces
+ .iter()
+ .any(|piece| piece.before.matches(before) && piece.after.matches(after))
+ }
+}
+
+/// The last [`HALF_WIDTH`] bytes of `bytes`, padded on the left when there are fewer, packed the way
+/// [`Half::ending`] packs a half.
+fn window_ending(bytes: &[u8]) -> u128 {
+ match bytes.last_chunk() {
+ Some(window) => u128::from_le_bytes(*window),
+ None => {
+ let mut window = [0u8; HALF_WIDTH];
+ window[HALF_WIDTH - bytes.len()..].copy_from_slice(bytes);
+ u128::from_le_bytes(window)
+ }
+ }
+}
+
+/// The first [`HALF_WIDTH`] bytes of `bytes`, padded on the right when there are fewer, packed the way
+/// [`Half::starting`] packs a half.
+fn window_starting(bytes: &[u8]) -> u128 {
+ match bytes.first_chunk() {
+ Some(window) => u128::from_le_bytes(*window),
+ None => {
+ let mut window = [0u8; HALF_WIDTH];
+ window[..bytes.len()].copy_from_slice(bytes);
+ u128::from_le_bytes(window)
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::models::bpe::{BPE, BpeBuilder, Merges, PipelineBPE, Vocab};
+ use crate::tokenizer::Model;
+
+ const DELIMITER: char = '▁';
+
+ fn delimiter() -> Literal {
+ Literal::new(DELIMITER.to_string().as_bytes()).unwrap()
+ }
+
+ fn split_on(veto: Veto, text: &str) -> Vec {
+ let split = ProvenCuts::new(delimiter(), veto);
+ let mut spans = Vec::new();
+ pipeline::PreTokenizer::pre_tokenize(&split, text, &mut spans).unwrap();
+ spans.iter().map(|s| text[s.range()].to_string()).collect()
+ }
+
+ fn veto_of(vocab: &[&str]) -> Option {
+ let pieces: Vec<(Vec, u32)> = vocab
+ .iter()
+ .enumerate()
+ .map(|(id, piece)| (piece.as_bytes().to_vec(), id as u32))
+ .collect();
+ Veto::build(&pieces, &delimiter())
+ }
+
+ /// A vocabulary shaped like gemma's: `>▁ BPE {
+ let vocab: Vocab = [
+ ("▁", 0u32),
+ ("<", 1),
+ (">", 2),
+ ("/", 3),
+ ("s", 4),
+ ("p", 5),
+ ("a", 6),
+ ("b", 7),
+ ("▁", 9),
+ (">▁", "▁"),
+ (">▁", " Vec {
+ let mut spans = Vec::new();
+ pipeline::PreTokenizer::pre_tokenize(split, text, &mut spans).unwrap();
+ let mut out = Vec::new();
+ let mut scratch = pipeline::Model::init_scratch(model);
+ for span in &spans {
+ pipeline::Model::tokenize_pipeline(model, &text[span.range()], &mut scratch, &mut out)
+ .unwrap();
+ }
+ out.iter().map(|token| token.id).collect()
+ }
+
+ /// The ids the model produces for the whole text at once — what the cuts must reproduce.
+ fn ids_whole_text(model: &BPE, text: &str) -> Vec {
+ model
+ .tokenize(text)
+ .unwrap()
+ .iter()
+ .map(|token| token.id)
+ .collect()
+ }
+
+ #[test]
+ fn proven_cuts_keep_the_whole_text_ids() {
+ let reference = metaspace_bpe();
+ let model = PipelineBPE::from_bpe(reference.clone(), false).unwrap();
+ let veto = Veto::build(&model.vocab_bytes(), &delimiter()).unwrap();
+ let split = ProvenCuts::new(delimiter(), veto);
+ for text in [
+ "▁a▁b",
+ "a▁b",
+ "▁▁▁a",
+ "a▁▁▁b",
+ "▁sp▁a",
+ "sp",
+ "a▁",
+ "▁",
+ "▁a▁sp▁b▁a",
+ // The veto pieces, alone and surrounded by other cuts.
+ "▁",
+ "▁a▁▁b",
+ "p▁a",
+ "▁ap▁a▁b",
+ // Starts like a veto piece but ends differently, so the cut stands.
+ "▁a",
+ "q▁a",
+ ] {
+ assert_eq!(
+ ids_word_by_word(&model, &split, text),
+ ids_whole_text(&reference, text),
+ "{text:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn proven_cuts_keep_delimiter_groups_whole() {
+ let veto = || veto_of(&["▁"]).unwrap();
+ assert_eq!(split_on(veto(), "a▁▁▁b"), ["a", "▁▁▁b"]);
+ assert_eq!(split_on(veto(), "▁▁▁a"), ["▁▁▁a"]);
+ assert_eq!(split_on(veto(), "hello"), ["hello"]);
+ assert_eq!(split_on(veto(), ""), Vec::::new());
+ }
+
+ #[test]
+ fn proven_cuts_skip_what_the_veto_forbids() {
+ let veto = || veto_of(&["▁", ">▁▁▁"), ["▁"]);
+ // ">" in front of the `▁` again, but "▁a"), ["", "▁a"]);
+ // Only the cut the piece covers is dropped, not the other ones.
+ assert_eq!(split_on(veto(), "x▁▁▁y"), ["x", "▁▁", "▁y"]);
+ }
+
+ #[test]
+ fn veto_needs_the_delimiter_in_the_vocabulary() {
+ assert!(veto_of(&["a", "b"]).is_none());
+ assert!(veto_of(&["a", "▁"]).is_some());
+ }
+
+ #[test]
+ fn veto_collects_only_the_pieces_that_reach_over_a_cut() {
+ assert!(veto_of(&["▁", "▁hello", "▁▁"]).unwrap().pieces.is_empty());
+ assert_eq!(
+ veto_of(&["▁", "p▁"]).unwrap().pieces,
+ [VetoPiece::new(b"p", b"").unwrap()]
+ );
+ }
+
+ #[test]
+ fn veto_gives_up_when_it_cannot_prove_anything() {
+ // Two delimiters in one piece would reach over two cuts at once.
+ assert!(veto_of(&["▁", "a▁b▁c"]).is_none());
+ // More pieces than the loop is willing to walk per cut.
+ let mut vocab = vec!["▁".to_string()];
+ vocab.extend((0..=MAX_VETO_PIECES).map(|i| format!("{i}▁x")));
+ let vocab: Vec<&str> = vocab.iter().map(String::as_str).collect();
+ assert!(veto_of(&vocab).is_none());
+ // A half that does not fit the fixed-width compare.
+ let wide = "a".repeat(HALF_WIDTH + 1);
+ assert!(veto_of(&["▁", &format!("{wide}▁x")]).is_none());
+ assert!(veto_of(&["▁", &format!("a▁{wide}")]).is_none());
+ }
+
+ #[test]
+ fn veto_matches_a_half_that_fills_the_window() {
+ let before = "a".repeat(HALF_WIDTH);
+ let veto = || veto_of(&["▁", &format!("{before}▁x"), "▁x"]).unwrap();
+ let text = format!("{before}▁x");
+ assert_eq!(split_on(veto(), &text), [text.as_str()]);
+ // One byte short of the half: the window pads with a zero the half does not hold, so the
+ // piece cannot form and the cut stands.
+ let short = &before[1..];
+ assert_eq!(split_on(veto(), &format!("{short}▁x")), [short, "▁x"]);
+ }
+
+ mod for_tokenizer {
+ use super::*;
+
+ /// The shapes a `▁`-spelling tokenizer ships, copied from the files in `data/`.
+ const GEMMA_NORMALIZER: &str =
+ r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#;
+ const GEMMA_PRE_TOKENIZER: &str = r#"{"type":"Split","pattern":{"String":" "},"behavior":"MergedWithPrevious","invert":false}"#;
+ const LLAMA_NORMALIZER: &str = r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#;
+
+ fn normalizer(json: &str) -> NormalizerWrapper {
+ serde_json::from_str(json).unwrap()
+ }
+
+ fn pre_tokenizer(json: &str) -> PreTokenizerWrapper {
+ serde_json::from_str(json).unwrap()
+ }
+
+ fn bpe_model(bpe: BPE) -> PipelineModel {
+ PipelineModel::BPE(PipelineBPE::from_bpe(bpe, false).unwrap())
+ }
+
+ #[test]
+ fn the_shapes_that_leave_the_text_whole_are_cut() {
+ let model = bpe_model(metaspace_bpe());
+ for (name, normalizer_json, pre_tokenizer_json) in [
+ ("gemma", GEMMA_NORMALIZER, Some(GEMMA_PRE_TOKENIZER)),
+ ("llama-2", LLAMA_NORMALIZER, None),
+ ] {
+ let cuts = for_tokenizer(
+ Some(&normalizer(normalizer_json)),
+ pre_tokenizer_json.map(pre_tokenizer).as_ref(),
+ &model,
+ )
+ .unwrap_or_else(|| panic!("{name} should be cut into words"));
+ assert_eq!(cuts.delimiter.pattern(), "▁".as_bytes(), "{name}");
+ }
+ }
+
+ #[test]
+ fn refuses_what_it_cannot_prove() {
+ let model = bpe_model(metaspace_bpe());
+ let refused: [(&str, Option<&NormalizerWrapper>, Option<&str>); 5] = [
+ ("no normalizer and no pre-tokenizer", None, None),
+ // Nothing here turns spaces into the delimiter.
+ (
+ "a normalizer that keeps spaces",
+ Some(&normalizer(r#"{"type":"NFC"}"#)),
+ None,
+ ),
+ // A step running after the replace could bring spaces back.
+ (
+ "a replace that is not the last step",
+ Some(&normalizer(
+ r#"{"type":"Sequence","normalizers":[{"type":"Replace","pattern":{"String":" "},"content":"▁"},{"type":"NFC"}]}"#,
+ )),
+ None,
+ ),
+ // This one cuts the text itself, so the model never sees it whole.
+ (
+ "a split that does match",
+ Some(&normalizer(GEMMA_NORMALIZER)),
+ Some(
+ r#"{"type":"Split","pattern":{"String":"▁"},"behavior":"MergedWithNext","invert":false}"#,
+ ),
+ ),
+ (
+ "a byte-level pre-tokenizer",
+ Some(&normalizer(GEMMA_NORMALIZER)),
+ Some(
+ r#"{"type":"ByteLevel","add_prefix_space":false,"trim_offsets":true,"use_regex":true}"#,
+ ),
+ ),
+ ];
+ for (name, normalizer, json) in refused {
+ let declared = json.map(pre_tokenizer);
+ assert!(
+ for_tokenizer(normalizer, declared.as_ref(), &model).is_none(),
+ "{name}"
+ );
+ }
+ // A vocabulary without the delimiter proves nothing about cutting on it.
+ let plain: Vocab = [("a", 0u32), ("b", 1)]
+ .iter()
+ .map(|(piece, id)| ((*piece).into(), *id))
+ .collect();
+ let plain = BpeBuilder::default()
+ .vocab_and_merges(plain, Vec::new())
+ .build()
+ .unwrap();
+ assert!(
+ for_tokenizer(Some(&normalizer(GEMMA_NORMALIZER)), None, &bpe_model(plain))
+ .is_none(),
+ "vocabulary without the delimiter"
+ );
+ }
+
+ /// The real files, so a change to either config shape shows up here. The veto counts are the
+ /// premise the whole proof rests on — gemma-4 has one piece that can merge across a word
+ /// boundary (`>▁ Option {
+ if self.is_deepseek() {
+ return Some(pipeline::FusedScan::DeepSeek);
+ }
+ let mut work = self
+ .pre_tokenizers
+ .iter()
+ .filter(|c| !matches!(c, PipelinePreTokenizer::None));
+ match (work.next(), work.next()) {
+ (Some(PipelinePreTokenizer::Split(split)), None) => {
+ split.native_fsm().map(pipeline::FusedScan::Gpt)
+ }
+ _ => None,
+ }
+ }
+
fn is_deepseek(&self) -> bool {
use crate::pre_tokenizers::split::SplitPattern;
use crate::tokenizer::SplitDelimiterBehavior::Isolated;
diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs
index 779c167f1..a38a7b621 100644
--- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs
+++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs
@@ -126,6 +126,14 @@ impl Split {
})
}
+ /// The native atomsplit FSM this split routes to, if its shape allows it: the
+ /// same recognition and `Isolated`, not-inverted filter `pre_tokenize` applies,
+ /// so a caller planning around the FSM cannot disagree with the split itself.
+ pub(crate) fn native_fsm(&self) -> Option {
+ self.fsm
+ .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated)
+ }
+
/// Pipeline canonicalization. A recognized whole-covering GPT regex shipped
/// as `(invert=true, behavior=Removed)` — the tiktoken-conversion convention
/// used by cl100k/o200k — is byte-exactly equivalent to `(invert=false,
@@ -201,13 +209,17 @@ impl pipeline::PreTokenizer for Split {
);
return Ok(());
}
- // Not a natively-routed GPT regex: fall-back to Literal or Regex search
+ // A plain-string pattern streams: the batch scan hands each delimiter offset straight
+ // to the span fold, and nothing is built in between.
+ if let Search::Literal(literal) = &self.search {
+ pipeline::split_literal(out, literal, text, self.behavior, self.invert);
+ return Ok(());
+ }
+ // Not a natively-routed GPT regex either: fall back to the Regex search
let matches = match (&self.search, self.invert) {
- (Search::Literal(literal), false) => literal.find_matches(text)?,
- (Search::Literal(literal), true) => Invert(literal).find_matches(text)?,
(Search::Regex(regex), false) => regex.find_matches(text)?,
(Search::Regex(regex), true) => Invert(regex).find_matches(text)?,
- (Search::Unavailable, _) => {
+ _ => {
return Err(
"this `Split` pattern needs a system-regex backend; enable the `fancy-regex` feature"
.into(),
@@ -415,6 +427,84 @@ mod tests {
);
}
+ /// The literal path across every behavior, on a text with a delimiter at both ends and a
+ /// consecutive pair — the shapes the span fold can get wrong.
+ #[test]
+ fn pipeline_literal_covers_every_behavior() {
+ let text = "-a--b-";
+ #[allow(clippy::type_complexity)]
+ let cases: Vec<(SplitDelimiterBehavior, Vec<(&str, (u32, u32))>)> = vec![
+ (Removed, vec![("a", (1, 2)), ("b", (4, 5))]),
+ (
+ Isolated,
+ vec![
+ ("-", (0, 1)),
+ ("a", (1, 2)),
+ ("-", (2, 3)),
+ ("-", (3, 4)),
+ ("b", (4, 5)),
+ ("-", (5, 6)),
+ ],
+ ),
+ (
+ MergedWithPrevious,
+ vec![("-", (0, 1)), ("a-", (1, 3)), ("-", (3, 4)), ("b-", (4, 6))],
+ ),
+ (
+ MergedWithNext,
+ vec![("-a", (0, 2)), ("-", (2, 3)), ("-b", (3, 5)), ("-", (5, 6))],
+ ),
+ (
+ Contiguous,
+ vec![
+ ("-", (0, 1)),
+ ("a", (1, 2)),
+ ("--", (2, 4)),
+ ("b", (4, 5)),
+ ("-", (5, 6)),
+ ],
+ ),
+ ];
+ for (behavior, expected) in cases {
+ assert_eq!(
+ pipeline_split("-".into(), behavior, false, text),
+ expected,
+ "behavior: {behavior:?}",
+ );
+ }
+ }
+
+ /// Inverted literal search: the gaps between delimiters become the matches.
+ #[test]
+ fn pipeline_literal_inverts() {
+ // Removed drops the (inverted) matches, keeping each delimiter.
+ assert_eq!(
+ pipeline_split("-".into(), Removed, true, "-a--b-"),
+ vec![("-", (0, 1)), ("-", (2, 3)), ("-", (3, 4)), ("-", (5, 6))],
+ );
+ // Isolated keeps everything either way.
+ assert_eq!(
+ pipeline_split("-".into(), Isolated, true, "a-b"),
+ vec![("a", (0, 1)), ("-", (1, 2)), ("b", (2, 3))],
+ );
+ }
+
+ #[test]
+ fn pipeline_literal_edges() {
+ assert_eq!(
+ pipeline_split("-".into(), Removed, false, ""),
+ Vec::<(&str, (u32, u32))>::new(),
+ );
+ assert_eq!(
+ pipeline_split("-".into(), Removed, false, "---"),
+ Vec::<(&str, (u32, u32))>::new(),
+ );
+ assert_eq!(
+ pipeline_split("-".into(), Removed, false, "abc"),
+ vec![("abc", (0, 3))],
+ );
+ }
+
#[test]
fn pipeline_gpt2_uses_fsm_and_matches_legacy() {
// The gpt2 pattern is recognized -> the pipeline path routes to the native
diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs
index 5a147f5f4..a411ef265 100644
--- a/tokenizers/tk-encode/src/tokenizer/pattern.rs
+++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs
@@ -43,22 +43,27 @@ impl Pattern for &Regex {
}
/// Searching for a plain string, does not need a regex engine: [`Literal`] scans the bytes.
+/// Two batch-scan passes: [`Literal::count_matches`] sizes `splits` exactly (each match adds
+/// itself and at most one gap before it, plus the final gap), then
+/// [`Literal::for_each_match`] streams the offsets in.
impl Pattern for &Literal {
fn find_matches(&self, inside: &str) -> Result> {
if inside.is_empty() {
return Ok(vec![((0, 0), false)]);
}
+ let width = self.pattern().len();
+ let count = self.count_matches(inside.as_bytes());
+ let mut splits = Vec::with_capacity(2 * count + 1);
let mut prev = 0;
- let mut splits = Vec::with_capacity(inside.len());
- for start in self.matches(inside.as_bytes()) {
- let end = start + self.pattern().len();
+ self.for_each_match(inside.as_bytes(), |start| {
+ let end = start + width;
if prev != start {
splits.push(((prev, start), false));
}
splits.push(((start, end), true));
prev = end;
- }
+ });
if prev != inside.len() {
splits.push(((prev, inside.len()), false))
}
diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs
index 16b7430b7..3a9a0e96d 100644
--- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs
+++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs
@@ -1,5 +1,7 @@
use std::cell::RefCell;
use std::convert::TryInto;
+use std::mem;
+use std::sync::{Mutex, PoisonError};
use std::{borrow::Cow, convert::TryFrom};
use atomsplit::classify::classify;
@@ -10,19 +12,24 @@ use crate::models::wordlevel::WordLevel;
use crate::models::wordpiece::{PipelineWordPiece, WordPieceScratch};
use crate::processors::bert::BertProcessing;
use crate::processors::roberta::RobertaProcessing;
+use crate::utils::GptFsm;
use crate::utils::byte_level::GPT2_REGEX_STR;
use crate::vocab::bucket_added_vocabulary::{
AddedToken as BucketAddedToken, AddedVocabulary as BucketAddedVocabulary,
};
use crate::{
ModelWrapper, PostProcessorWrapper, PreTokenizerWrapper, Token, Tokenizer,
- normalizers::{NormalizerWrapper, metaspace::MetaspaceNormalizer},
+ normalizers::{
+ NormalizerWrapper,
+ metaspace::{MetaspaceNormalizer, PrependMode},
+ },
pre_tokenizers::{
bert::BertPreTokenizer,
delimiter::CharDelimiterSplit,
digits::Digits,
fixed_length::FixedLength,
metaspace,
+ proven_cuts::{self, ProvenCuts, ZeroCopyMetaspace},
punctuation::Punctuation,
sequence::PipelineSequence,
split::{Split as SplitPretok, SplitPattern},
@@ -56,21 +63,166 @@ pub(crate) fn classify_into_spans(
}
classify(bytes, &mut tags[..n]);
let k = fsm(bytes, &tags[..n], &mut spans[..n + 1]);
- out.extend_from_slice(&spans[..k]); // same type now — plain memcpy, no per-token conversion
+ out.extend_from_slice(&spans[..k]); // same type now: plain memcpy, no per-token conversion
});
}
pub trait Normalizer {
fn normalize<'a>(&self, input: &'a str) -> Result>;
+
+ /// The rewrite this normalizer could leave pending instead of writing, or `None` when its
+ /// rewrite restructures the text and has to be written. Only rewrites that change the
+ /// text one character at a time qualify — see [`PendingRewrite`]. Today those are
+ /// [`MetaspaceNormalizer`] (the SentencePiece space swap) and a single-character literal
+ /// [`Replace`](crate::normalizers::replace::Replace).
+ fn pending_rewrite(&self) -> Option {
+ None
+ }
}
-/// Runs `normalizers` in order, each one seeing what the one before it produced.
+/// A rewrite that changes the text one character at a time: every `from` reads as `to`, and
+/// `prepend` may put one `to` in front of the chunk.
///
-/// A normalizer returns a [`Cow`] (copy-on-write): a borrow when it had nothing to change, an owned
-/// `String` when it rewrote the text.
+/// Because positions map one to one, such a rewrite does not have to be written into a copy:
+/// a consumer can read the rewritten form off the raw text. [`NormalizedText`] carries one
+/// left pending by the normalizer chain, and the zero-copy encode path consumes it in place;
+/// [`PendingRewrite::write`] produces the rewritten text for every other consumer.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct PendingRewrite {
+ pub(crate) from: char,
+ pub(crate) to: char,
+ pub(crate) prepend: PrependMode,
+}
+
+impl PendingRewrite {
+ /// Whether the rewrite puts a `to` in front of `text`. Empty text is handed back as is,
+ /// so it never takes one.
+ pub(crate) fn prepends(&self, text: &str) -> bool {
+ !text.is_empty()
+ && match self.prepend {
+ PrependMode::Never => false,
+ // Text already opening with the rewrite's own characters counts as marked
+ // (a leading `from` becomes a `to` under the swap).
+ PrependMode::IfMissing => {
+ !text.starts_with(self.from) && !text.starts_with(self.to)
+ }
+ PrependMode::Unconditional => true,
+ }
+ }
+
+ /// Writes the rewrite into a copy: one pass, one right-sized allocation, and a borrow
+ /// when there is nothing to swap or prepend.
+ pub(crate) fn write<'a>(&self, input: &'a str) -> Result> {
+ if input.is_empty() {
+ return Ok(Cow::Borrowed(input));
+ }
+ let prepend = self.prepends(input);
+ let mut from_buf = [0u8; 4];
+ let from = self.from.encode_utf8(&mut from_buf);
+ let mut to_buf = [0u8; 4];
+ let to = self.to.encode_utf8(&mut to_buf);
+ // Counting the swaps first sizes the copy exactly and lets them stream through the
+ // batch scan instead of restarting a search at every match.
+ let pattern =
+ atomsplit::literal::Literal::new(from.as_bytes()).expect("a char is never empty");
+ let count = pattern.count_matches(input.as_bytes());
+ if !prepend && count == 0 {
+ return Ok(Cow::Borrowed(input));
+ }
+ let mut rewritten = String::with_capacity(
+ input.len() - count * from.len()
+ + count * to.len()
+ + if prepend { to.len() } else { 0 },
+ );
+ if prepend {
+ rewritten.push_str(to);
+ }
+ let mut prev = 0;
+ pattern.for_each_match(input.as_bytes(), |start| {
+ rewritten.push_str(&input[prev..start]);
+ rewritten.push_str(to);
+ prev = start + from.len();
+ });
+ rewritten.push_str(&input[prev..]);
+ Ok(Cow::Owned(rewritten))
+ }
+}
+
+impl Normalizer for PendingRewrite {
+ fn normalize<'a>(&self, input: &'a str) -> Result> {
+ self.write(input)
+ }
+
+ fn pending_rewrite(&self) -> Option {
+ Some(*self)
+ }
+}
+
+/// What a normalizer chain produced for one chunk: the normalized text, or that text one
+/// pending rewrite short of it.
+///
+/// [`NormalizedText::from_chain`] leaves a trailing [`PendingRewrite`] unwritten; the
+/// zero-copy encode path reads the rewritten form off `text`'s raw characters, and every
+/// other consumer calls [`NormalizedText::write`] to get plain text.
+pub struct NormalizedText<'a> {
+ text: Cow<'a, str>,
+ pending: Option,
+}
+
+impl<'a> NormalizedText<'a> {
+ /// Runs `normalizers` over `input`, leaving a trailing pending rewrite unwritten. A
+ /// pending rewrite anywhere else in the chain runs as declared: the step after it needs
+ /// the written text.
+ pub(crate) fn from_chain(normalizers: &[N], input: &'a str) -> Result {
+ if let [head @ .., last] = normalizers
+ && let Some(pending) = last.pending_rewrite()
+ {
+ return Ok(Self {
+ text: normalize_all(head, input)?,
+ pending: Some(pending),
+ });
+ }
+ Ok(Self {
+ text: normalize_all(normalizers, input)?,
+ pending: None,
+ })
+ }
+
+ /// The normalized text, with any pending rewrite written into it.
+ pub fn write(self) -> Result> {
+ match self.pending {
+ None => Ok(self.text),
+ Some(rewrite) => normalize_step(&rewrite, self.text),
+ }
+ }
+}
+
+/// One normalizer step over the running text: the copy-on-write bookkeeping one link of
+/// [`normalize_all`]'s chain needs.
///
-/// Chaining them needs care, because an owned `String` produced halfway through the chain is local to this function.
-/// The next normalizer may hand back a borrow of that locally owned [`String`], and that borrow cannot outlive the `String`.
+/// A normalizer returns a [`Cow`] (copy-on-write): a borrow when it had nothing to change, an
+/// owned `String` when it rewrote the text. Chaining needs care, because an owned `String`
+/// produced halfway through the chain is local: the next normalizer may hand back a borrow of
+/// that `String`, and that borrow cannot outlive it.
+fn normalize_step<'a, N: Normalizer>(normalizer: &N, cow: Cow<'a, str>) -> Result> {
+ Ok(match cow {
+ // Still the caller's input, which outlives us: pass it straight on.
+ Cow::Borrowed(s) => normalizer.normalize(s)?,
+ Cow::Owned(s) => {
+ let out = match normalizer.normalize(&s)? {
+ // Rewritten again: keep the new `String`, drop ours.
+ Cow::Owned(o) => Some(o),
+ // Handed `s` back untouched: keep the `String` we already own.
+ Cow::Borrowed(b) if b.as_ptr() == s.as_ptr() && b.len() == s.len() => None,
+ // A borrow into `s` (for example, a substring of `s`): copy it out before `s` is dropped.
+ Cow::Borrowed(b) => Some(b.to_owned()),
+ };
+ Cow::Owned(out.unwrap_or(s))
+ }
+ })
+}
+
+/// Runs `normalizers` in order, each one seeing what the one before it produced.
///
/// Text no normalizer touches is never copied: it stays a borrow of `input` throughout.
pub(crate) fn normalize_all<'a, N: Normalizer>(
@@ -79,21 +231,7 @@ pub(crate) fn normalize_all<'a, N: Normalizer>(
) -> Result> {
let mut cow: Cow<'a, str> = Cow::Borrowed(input);
for normalizer in normalizers {
- cow = match cow {
- // Still `input` itself, which outlives us: pass it straight on.
- Cow::Borrowed(s) => normalizer.normalize(s)?,
- Cow::Owned(s) => {
- let out = match normalizer.normalize(&s)? {
- // Rewritten again: keep the new `String`, drop ours.
- Cow::Owned(o) => Some(o),
- // Handed `s` back untouched: keep the `String` we already own.
- Cow::Borrowed(b) if b.as_ptr() == s.as_ptr() && b.len() == s.len() => None,
- // A borrow into `s` (for example, a substring of `s`): copy it out before `s` is dropped.
- Cow::Borrowed(b) => Some(b.to_owned()),
- };
- Cow::Owned(out.unwrap_or(s))
- }
- };
+ cow = normalize_step(normalizer, cow)?;
}
Ok(cow)
}
@@ -105,9 +243,11 @@ pub(crate) fn normalize_all<'a, N: Normalizer>(
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
enum PipelineNormalizer {
- /// The `normalizer` field of the config, as-is.
+ /// One step of the config's `normalizer` field, as-is (a `Sequence`'s members each get their
+ /// own entry).
Declared(NormalizerWrapper),
- /// The text-rewriting half of a `Metaspace` pre-tokenizer.
+ /// The one-pass stand-in for declared steps that spell out the SentencePiece space rewrite
+ /// ([`MetaspaceNormalizer::fuse`]), or the text-rewriting half of a `Metaspace` pre-tokenizer.
Metaspace(MetaspaceNormalizer),
}
@@ -118,6 +258,13 @@ impl Normalizer for PipelineNormalizer {
Self::Metaspace(normalizer) => normalizer.normalize(input),
}
}
+
+ fn pending_rewrite(&self) -> Option {
+ match self {
+ Self::Declared(normalizer) => normalizer.pending_rewrite(),
+ Self::Metaspace(normalizer) => normalizer.pending_rewrite(),
+ }
+ }
}
/// Range-based pre-tokenization: yields spans into the input rather than owned
@@ -135,6 +282,7 @@ pub enum PipelinePreTokenizer {
Delimiter(CharDelimiterSplit),
Digits(Digits),
FixedLength(FixedLength),
+ ProvenCuts(ProvenCuts),
Punctuation(Punctuation),
Sequence(PipelineSequence),
Split(SplitPretok),
@@ -144,6 +292,54 @@ pub enum PipelinePreTokenizer {
None,
}
+/// A pre-tokenizer whose whole split runs as one native FSM pass, named ahead of time
+/// so [`PipelineTokenizer::encode_fused`] can drive the scan and the model together.
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub(crate) enum FusedScan {
+ Gpt(GptFsm),
+ DeepSeek,
+}
+
+/// The route a chunk takes through [`PipelineTokenizer::encode_generic`]. See
+/// [`PipelineTokenizer::encode_route`].
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum EncodeRoute {
+ /// The pending rewrite is read in place and proven cuts split the raw text: no
+ /// rewritten copy is built and no added-token scan runs over one.
+ ZeroCopy,
+ /// One native FSM cuts the chunk and hands each span straight to the BPE model, so
+ /// the split and the model are a single pass.
+ Fused,
+ /// Normalize, write, scan for added tokens, pre-tokenize into a span buffer, then run
+ /// the model once per span.
+ Staged,
+}
+
+impl EncodeRoute {
+ /// The route's name, for a benchmark's report.
+ pub fn as_str(self) -> &'static str {
+ match self {
+ Self::ZeroCopy => "zero-copy",
+ Self::Fused => "fused",
+ Self::Staged => "staged",
+ }
+ }
+}
+
+impl PipelinePreTokenizer {
+ /// The single native FSM this whole pre-tokenizer reduces to, if any: a recognized
+ /// lone `Split`, or a `Sequence` that collapses to one. Each arm asks the child for
+ /// the routing its own `pre_tokenize` applies, so a fused caller cannot disagree
+ /// with the split it bypasses.
+ pub(crate) fn native_fsm(&self) -> Option {
+ match self {
+ Self::Split(split) => split.native_fsm().map(FusedScan::Gpt),
+ Self::Sequence(sequence) => sequence.native_fsm(),
+ _ => None,
+ }
+ }
+}
+
impl PreTokenizer for PipelinePreTokenizer {
fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> {
match self {
@@ -158,6 +354,7 @@ impl PreTokenizer for PipelinePreTokenizer {
Self::Delimiter(pretok) => pretok.pre_tokenize(text, out),
Self::Digits(pretok) => pretok.pre_tokenize(text, out),
Self::FixedLength(pretok) => pretok.pre_tokenize(text, out),
+ Self::ProvenCuts(pretok) => pretok.pre_tokenize(text, out),
Self::Punctuation(pretok) => pretok.pre_tokenize(text, out),
Self::Sequence(pretok) => pretok.pre_tokenize(text, out),
Self::Split(pretok) => pretok.pre_tokenize(text, out),
@@ -323,7 +520,7 @@ impl TryFrom<&PostProcessorWrapper> for PipelinePostProcessor {
}
}
-/// An output token. Carries only the vocabulary `id` — offsets and the token
+/// An output token. Carries only the vocabulary `id`, since offsets and the token
/// string are dropped, which is all an encode-only caller needs.
#[derive(Debug, Clone, Copy)]
pub struct PipelineToken {
@@ -448,6 +645,82 @@ pub struct PipelineTokenizer {
pre_tokenizer: PipelinePreTokenizer,
model: PipelineModel,
post_processor: PipelinePostProcessor,
+ scratch_pool: ScratchPool,
+ /// When set, a full encode runs each chunk through [`ZeroCopyMetaspace`] instead of the
+ /// normalize/pre-tokenize stages: the rewrite those stages would write is read off the
+ /// raw text. `None` for every other pipeline shape; see [`PipelineTokenizer::try_from`].
+ zero_copy: Option,
+}
+
+/// A pool of [`PipelineModelScratch`].
+///
+/// When calling [`PipelineTokenizer::encode`], an instance of [`PipelineModelScratch`] is taken out of this pool
+/// and given to the tokenizer. When the encoding is done, the scratch buffer is returned to the pool and can be
+/// reused by later calls.
+///
+/// The reusability matters because the scratch buffer may hold cache structures which are more useful when reused,
+/// and less importantly it saves an extra allocation for an fresh buffer every time.
+struct ScratchPool(Mutex>);
+
+impl ScratchPool {
+ fn new() -> Self {
+ Self(Mutex::new(Vec::new()))
+ }
+
+ /// Get a scratch buffer from the pool, wrapped in a [`ScratchGuard`].
+ /// When the [`ScratchGuard`] gets dropped, the scratch buffer is pushed back to the pool.
+ fn get<'a>(&'a self, model: &PipelineModel) -> ScratchGuard<'a> {
+ // The Mutex lock is held just long enough to pop the scratch out of the pool
+ let taken = self.0.lock().unwrap_or_else(PoisonError::into_inner).pop();
+ ScratchGuard {
+ // If there was no scratch buffer available in the pool, we build.a fresh one
+ scratch: taken.unwrap_or_else(|| model.init_scratch()),
+ pool: self,
+ }
+ }
+
+ #[cfg(test)]
+ fn len(&self) -> usize {
+ self.0.lock().unwrap_or_else(PoisonError::into_inner).len()
+ }
+}
+
+/// A wrapper around [`PipelineModelScratch`].
+/// Implements [`Deref`] and [`DerefMut`], so it behaves as [`PipelineModelScratch`].
+///
+/// When it gets dropped, it pushes [`Self::scratch`] back into the shared [`Self::pool`] so it can
+/// get reused by a later call to [`PipelineTokenizer::encode`].
+///
+/// TODO @McPatate : The Mutex can create contention, to be replaced by a better access pattern
+struct ScratchGuard<'a> {
+ scratch: PipelineModelScratch,
+ pool: &'a ScratchPool,
+}
+
+impl Drop for ScratchGuard<'_> {
+ fn drop(&mut self) {
+ // Steals the scratch buffer from self, replaces it with PipelineModelScratch::default()
+ let scratch = mem::take(&mut self.scratch);
+ // Push the scratch back in the pool
+ self.pool
+ .0
+ .lock()
+ .unwrap_or_else(PoisonError::into_inner)
+ .push(scratch);
+ }
+}
+
+impl std::ops::Deref for ScratchGuard<'_> {
+ type Target = PipelineModelScratch;
+ fn deref(&self) -> &PipelineModelScratch {
+ &self.scratch
+ }
+}
+
+impl std::ops::DerefMut for ScratchGuard<'_> {
+ fn deref_mut(&mut self) -> &mut PipelineModelScratch {
+ &mut self.scratch
+ }
}
impl TryFrom<&Tokenizer> for PipelineTokenizer {
@@ -462,7 +735,29 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer {
fn try_from(tok: &Tokenizer) -> Result {
let mut normalizers = Vec::new();
if let Some(declared) = tok.get_normalizer() {
- normalizers.push(PipelineNormalizer::Declared(declared.clone()));
+ // A `Sequence` runs its members in order, exactly what this list does, so its members
+ // join the list one by one. That puts them in [`MetaspaceNormalizer::fuse`]'s reach:
+ // steps spelling out the SentencePiece space rewrite (llama-2 as two whole-chunk
+ // copies, `Prepend("▁")` then `Replace(" " -> "▁")`; gemma-4 as the `Replace` alone)
+ // are swapped for the normalizer doing that job in one pass. Only this in-memory
+ // pipeline is rewritten; the config still serializes as declared.
+ let steps: &[NormalizerWrapper] = match declared {
+ NormalizerWrapper::Sequence(sequence) => sequence.as_ref(),
+ single => std::slice::from_ref(single),
+ };
+ let mut index = 0;
+ while index < steps.len() {
+ match MetaspaceNormalizer::fuse(&steps[index..]) {
+ Some((fused, covered)) => {
+ normalizers.push(PipelineNormalizer::Metaspace(fused));
+ index += covered;
+ }
+ None => {
+ normalizers.push(PipelineNormalizer::Declared(steps[index].clone()));
+ index += 1;
+ }
+ }
+ }
}
// A `Metaspace` pre-tokenizer does two jobs at once: it writes `▁` delimiters into the text,
@@ -559,16 +854,42 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer {
ModelWrapper::WordPiece(model) => PipelineModel::WordPiece(model.try_into()?),
};
+ // Some tokenizers spell every space as `▁` and then cut nowhere, so the model receives the
+ // whole text in one piece. Cutting it into words is faster wherever the vocabulary proves
+ // that leaves the ids alone — and that proof reads the vocabulary, hence after the model.
+ let pre_tokenizer =
+ match proven_cuts::for_tokenizer(tok.get_normalizer(), tok.get_pre_tokenizer(), &model)
+ {
+ Some(cuts) => PipelinePreTokenizer::ProvenCuts(cuts),
+ None => pre_tokenizer,
+ };
+
+ // A chain ending in a pending rewrite, followed by proven cuts: the rewrite can be
+ // read off the raw text instead of written (see [`ZeroCopyMetaspace`]). Added tokens
+ // flagged `normalized` are matched against rewritten text, which this path never
+ // builds, so one such token keeps the written stages.
+ let zero_copy = match (normalizers.last(), &pre_tokenizer, &model) {
+ (Some(last), PipelinePreTokenizer::ProvenCuts(cuts), PipelineModel::BPE(bpe))
+ if !added_vocabulary.has_normalized_tokens() =>
+ {
+ last.pending_rewrite()
+ .and_then(|rewrite| ZeroCopyMetaspace::build(rewrite, cuts, bpe))
+ }
+ _ => None,
+ };
+
Ok(Self {
added_vocabulary,
normalizers,
pre_tokenizer,
model,
+ zero_copy,
post_processor: tok
.get_post_processor()
.map(PipelinePostProcessor::try_from)
.transpose()?
.unwrap_or_default(),
+ scratch_pool: ScratchPool::new(),
})
}
}
@@ -588,6 +909,37 @@ impl PipelineTokenizer {
&self.model
}
+ pub fn get_pre_tokenizer(&self) -> &PipelinePreTokenizer {
+ &self.pre_tokenizer
+ }
+
+ /// The single native FSM the whole split reduces to, when the model behind it is one
+ /// the fused scan can feed. [`encode_fused`](Self::encode_fused) and
+ /// [`encode_route`](Self::encode_route) both ask this, so the route a benchmark
+ /// reports cannot disagree with the route an encode takes.
+ fn fused_scan(&self) -> Option {
+ let scan = self.pre_tokenizer.native_fsm()?;
+ matches!(self.model, PipelineModel::BPE(_)).then_some(scan)
+ }
+
+ /// Which route this tokenizer's chunks take through [`encode_generic`].
+ ///
+ /// The three routes decompose into stages differently. The fused scan does the split
+ /// and the model in one pass, and the zero-copy route replaces the rewrite write and
+ /// the added-token scan over the rewritten text as well as the split. A stage
+ /// breakdown is only meaningful against the route it measured, so a benchmark reports
+ /// this next to the numbers.
+ pub fn encode_route(&self) -> EncodeRoute {
+ // Same order `encode_generic` tests them in.
+ if self.zero_copy.is_some() {
+ EncodeRoute::ZeroCopy
+ } else if self.fused_scan().is_some() {
+ EncodeRoute::Fused
+ } else {
+ EncodeRoute::Staged
+ }
+ }
+
/// Encode `input` into token ids.
///
/// Special tokens are matched in two passes:
@@ -599,7 +951,7 @@ impl PipelineTokenizer {
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();
+ let mut scratch = self.scratch_pool.get(&self.model);
self.encode_generic::<{ Self::STAGE_POSTPROCESS }>(
input,
@@ -613,7 +965,7 @@ impl PipelineTokenizer {
/// Decode token ids back to a `String`.
///
- /// Not implemented yet — the pipeline decode path is being built. It fails
+ /// Not implemented yet: the pipeline decode path is being built. It fails
/// loud (rather than returning a plausible-but-wrong string) so the oracle
/// test and the comparative benchmark report decode as *pending* instead of
/// silently validating garbage. Implementing this flips the ignored
@@ -623,7 +975,7 @@ impl PipelineTokenizer {
}
/// Decode several id sequences at once, one `String` per input. Mirrors the
- /// released `decode_batch`; sequential (KISS) — behavior-identical to a
+ /// released `decode_batch`; sequential (KISS), behavior-identical to a
/// parallel map, since each [`decode`](Self::decode) is independent.
pub fn decode_batch(
&self,
@@ -638,7 +990,7 @@ impl PipelineTokenizer {
/// Incremental decode: feed ids one at a time via [`PipelineDecodeStream::step`].
/// Same prefix-tracking scheme as the released `DecodeStream`, built on
- /// [`decode`](Self::decode) — so it is correct exactly where `decode` is.
+ /// [`decode`](Self::decode), so it is correct exactly where `decode` is.
pub fn decode_stream(&self, skip_special_tokens: bool) -> PipelineDecodeStream<'_> {
PipelineDecodeStream {
tokenizer: self,
@@ -651,7 +1003,7 @@ impl PipelineTokenizer {
/// 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
+ /// the disabled stages are compiled out, so the full specialization
/// ([`STAGE_POSTPROCESS`], 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.
@@ -660,7 +1012,7 @@ impl PipelineTokenizer {
/// [`STAGE_POSTPROCESS`]: Self::STAGE_POSTPROCESS
///
/// `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
+ /// 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(
@@ -684,11 +1036,42 @@ impl PipelineTokenizer {
output.push(PipelineToken { id: token });
}
Segment::Text(chunk) => {
- let normalized: Cow = if STAGE >= Self::STAGE_NORMALIZE {
- normalize_all(&self.normalizers, chunk)?
+ let normalized = if STAGE >= Self::STAGE_NORMALIZE {
+ NormalizedText::from_chain(&self.normalizers, chunk)?
} else {
- Cow::Borrowed(chunk)
+ NormalizedText {
+ text: Cow::Borrowed(chunk),
+ pending: None,
+ }
};
+ // A pipeline with the zero-copy route reads the pending rewrite in
+ // place at *every* stage: it never writes the rewritten text and never
+ // scans it for added tokens, so no rung may charge this route for work
+ // it does not do. The cut is its split stage and the model its own.
+ //
+ // Whether the route applies is a property of the pipeline, not of
+ // the chunk. That matters below `STAGE_NORMALIZE`: there the ladder has
+ // not computed the rewrite yet, so `pending` is `None` for a reason
+ // that says nothing about the route. Falling through on a missing
+ // `pending` above that keeps ids right if a chunk ever lacks one.
+ if let Some(zero_copy) = &self.zero_copy {
+ if STAGE < Self::STAGE_NORMALIZE {
+ continue;
+ }
+ if normalized.pending.is_some() {
+ if STAGE >= Self::STAGE_SPLIT {
+ self.encode_chunk_zero_copy::(
+ zero_copy,
+ &normalized.text,
+ pre_tokens,
+ scratch,
+ output,
+ )?;
+ }
+ continue;
+ }
+ }
+ let normalized: Cow = normalized.write()?;
// Extract special tokens from the normalized input
for segment in
@@ -700,18 +1083,27 @@ impl PipelineTokenizer {
}
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,
- )?;
+ if self.encode_fused::(
+ normalized_chunk,
+ pre_tokens,
+ scratch,
+ output,
+ ) {
+ // Split and model ran as one pass; nothing left to do.
+ } else {
+ // 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,
+ )?;
+ }
}
}
}
@@ -727,6 +1119,148 @@ impl PipelineTokenizer {
}
Ok(())
}
+
+ /// The fused fast path: cut `chunk` with its native FSM and hand each span
+ /// straight to the BPE model the moment it is emitted. The span buffer, the
+ /// per-span model dispatch and the per-span trait call of the staged path all
+ /// disappear, and the model reads each word while its bytes are still hot from
+ /// the scan. Returns `false` when this tokenizer is not that shape, and the
+ /// staged path must run instead.
+ ///
+ /// Below [`STAGE_MODEL`](Self::STAGE_MODEL) the scan still runs, because it *is* the
+ /// split stage for a fused pipeline, and each span lands in `pre_tokens` instead of
+ /// going to the model. So the ladder's split rung times the scanner a full encode
+ /// really uses, and the model rung is a difference between two runs of the same scan.
+ /// Timing the staged array FSM here instead would report a splitter that no longer
+ /// runs, and make the model rung a subtraction across two different pipelines.
+ ///
+ /// [`STAGE_MODEL`]: Self::STAGE_MODEL
+ fn encode_fused(
+ &self,
+ chunk: &str,
+ pre_tokens: &mut Vec,
+ scratch: &mut PipelineModelScratch,
+ output: &mut Vec,
+ ) -> bool {
+ let Some(scan) = self.fused_scan() else {
+ return false;
+ };
+ let PipelineModel::BPE(model) = &self.model else {
+ return false;
+ };
+ let PipelineModelScratch::BPE(scratch) = scratch else {
+ return false;
+ };
+ thread_local! {
+ static TAGS: RefCell> = const { RefCell::new(Vec::new()) };
+ }
+ if STAGE < Self::STAGE_MODEL {
+ pre_tokens.clear();
+ }
+ let bytes = chunk.as_bytes();
+ TAGS.with(|cell| {
+ let tags = &mut *cell.borrow_mut();
+ if tags.len() < bytes.len() {
+ tags.resize(bytes.len(), 0); // grow-only, as in `classify_into_spans`
+ }
+ classify(bytes, &mut tags[..bytes.len()]);
+ let tags = &tags[..bytes.len()];
+ use atomsplit::fsm::{
+ scan_byte_level_masked, scan_cl100k_cap_masked, scan_deepseek_masked,
+ scan_o200k_masked, scan_tekken_masked,
+ };
+ /// One span, either through the model or into the span buffer. `STAGE` is a
+ /// const generic, so a full encode compiles to the model call alone.
+ #[inline(always)]
+ fn feed(
+ model: &PipelineBPE,
+ chunk: &str,
+ span: Span,
+ pre_tokens: &mut Vec,
+ scratch: &mut BpeScratch,
+ output: &mut Vec,
+ ) {
+ if STAGE < PipelineTokenizer::STAGE_MODEL {
+ pre_tokens.push(span);
+ return;
+ }
+ model.tokenize_span(&chunk[span.range()], scratch, output);
+ }
+ // One emit closure literal per arm: each scan gets its own instance by
+ // value, which is what lets the emit inline into the scan loop.
+ match scan {
+ FusedScan::Gpt(GptFsm::Gpt2) => scan_byte_level_masked(bytes, tags, |span| {
+ feed::(model, chunk, span, pre_tokens, scratch, output);
+ }),
+ FusedScan::Gpt(GptFsm::Cl100k { digit_cap }) => {
+ scan_cl100k_cap_masked(bytes, tags, digit_cap, |span| {
+ feed::(model, chunk, span, pre_tokens, scratch, output);
+ })
+ }
+ FusedScan::Gpt(GptFsm::O200k) => scan_o200k_masked(bytes, tags, |span| {
+ feed::(model, chunk, span, pre_tokens, scratch, output);
+ }),
+ FusedScan::Gpt(GptFsm::Tekken) => scan_tekken_masked(bytes, tags, |span| {
+ feed::(model, chunk, span, pre_tokens, scratch, output);
+ }),
+ FusedScan::DeepSeek => scan_deepseek_masked(bytes, tags, |span| {
+ feed::(model, chunk, span, pre_tokens, scratch, output);
+ }),
+ }
+ });
+ true
+ }
+
+ /// One chunk through [`ZeroCopyMetaspace`]: cut the raw text, then read each span through
+ /// the swap. The rewritten chunk is never built, and the special-token pass over it is
+ /// skipped with it — the path is only built when no added token matches normalized text.
+ ///
+ /// The cut is this route's split stage, so below [`STAGE_MODEL`](Self::STAGE_MODEL) it
+ /// returns with `pre_tokens` filled and the model unrun.
+ ///
+ /// [`STAGE_MODEL`]: Self::STAGE_MODEL
+ fn encode_chunk_zero_copy(
+ &self,
+ zero_copy: &ZeroCopyMetaspace,
+ chunk: &str,
+ pre_tokens: &mut Vec,
+ scratch: &mut PipelineModelScratch,
+ output: &mut Vec,
+ ) -> Result<()> {
+ let (PipelineModel::BPE(bpe), PipelineModelScratch::BPE(scratch)) = (&self.model, scratch)
+ else {
+ unreachable!("`try_from` only builds `zero_copy` for a BPE model");
+ };
+ pre_tokens.clear();
+ let prepend = zero_copy.cut(chunk, pre_tokens);
+ if STAGE < Self::STAGE_MODEL {
+ return Ok(());
+ }
+ let mut spans = pre_tokens.iter();
+ if prepend && let Some(first) = spans.next() {
+ // The prepend rewrites this one word for real, so its cache key is the rewritten
+ // form; every other span's key is its raw bytes.
+ let rewritten = zero_copy.rewrite().write(&chunk[first.range()])?;
+ bpe.tokenize_pipeline(&rewritten, scratch, output)?;
+ }
+ let swap = zero_copy.swap();
+ for span in spans {
+ bpe.tokenize_swapped(&chunk[span.range()], swap, scratch, output)?;
+ }
+ Ok(())
+ }
+
+ /// Turns the zero-copy metaspace path off, so a full encode runs the written rewrite.
+ #[doc(hidden)] // public only so `examples/normalize_claims.rs` can A/B the two paths
+ pub fn disable_zero_copy(&mut self) {
+ self.zero_copy = None;
+ }
+
+ /// Whether full encodes take the zero-copy metaspace path.
+ #[doc(hidden)] // public only so `examples/normalize_claims.rs` can verify the path fires
+ pub fn has_zero_copy(&self) -> bool {
+ self.zero_copy.is_some()
+ }
}
/// Streaming decoder over a [`PipelineTokenizer`]; see [`PipelineTokenizer::decode_stream`].
@@ -785,7 +1319,7 @@ pub enum SplitPolicy {
Isolate,
}
-/// Splits `text` into same-class groups, emitting each as a [`Split`]
+/// Splits `text` into same-class groups, emitting each as a [`Span`]
/// according to its [`SplitPolicy`].
///
/// `classify` maps each char to a small `Copy + Eq` class, the current
@@ -907,75 +1441,131 @@ pub fn split_matches(
matches: Vec<((usize, usize), bool)>,
behavior: SplitDelimiterBehavior,
) {
- use SplitDelimiterBehavior::*;
+ let mut fold = SplitFold::new(out, behavior);
+ for (offsets, is_match) in matches {
+ fold.segment(offsets, is_match);
+ }
+ fold.finish();
+}
- // (offsets, should_remove) — mirrors `NormalizedString::split`.
- let splits: Vec<((usize, usize), bool)> = match behavior {
- Isolated => matches.into_iter().map(|(o, _)| (o, false)).collect(),
- Removed => matches, // should_remove == is_match
- Contiguous => {
- let mut previous_match = false;
- matches
- .into_iter()
- .fold(vec![], |mut acc, (offsets, is_match)| {
- if is_match == previous_match {
- if let Some(((_, end), _)) = acc.last_mut() {
- *end = offsets.1;
- } else {
- acc.push((offsets, false));
- }
- } else {
- acc.push((offsets, false));
- }
- previous_match = is_match;
- acc
- })
+/// [`split_matches`] for a plain-byte pattern, shared by every pre-tokenizer that cuts on a
+/// literal (`Split` with a string pattern, `CharDelimiterSplit`). The count pass sizes `out`
+/// for the worst case, one span per match plus the pieces around them; the batch scan then
+/// hands each match offset straight to the fold, so no segment vector is built. `invert`
+/// flips which side counts as the delimiter, as in `Split`.
+pub(crate) fn split_literal(
+ out: &mut Vec,
+ literal: &atomsplit::literal::Literal,
+ text: &str,
+ behavior: SplitDelimiterBehavior,
+ invert: bool,
+) {
+ let width = literal.pattern().len();
+ let count = literal.count_matches(text.as_bytes());
+ out.reserve(2 * count + 1);
+ let mut fold = SplitFold::new(out, behavior);
+ let mut prev = 0;
+ literal.for_each_match(text.as_bytes(), |start| {
+ if prev != start {
+ fold.segment((prev, start), invert);
}
- MergedWithPrevious => {
- let mut previous_match = false;
- matches
- .into_iter()
- .fold(vec![], |mut acc, (offsets, is_match)| {
- if is_match && !previous_match {
- if let Some(((_, end), _)) = acc.last_mut() {
- *end = offsets.1;
- } else {
- acc.push((offsets, false));
- }
- } else {
- acc.push((offsets, false));
- }
- previous_match = is_match;
- acc
- })
+ fold.segment((start, start + width), !invert);
+ prev = start + width;
+ });
+ if prev != text.len() {
+ fold.segment((prev, text.len()), invert);
+ }
+ fold.finish();
+}
+
+/// The streaming form of [`split_matches`]: feed it the covering `(offsets, is_match)`
+/// segments left to right with [`SplitFold::segment`], then call [`SplitFold::finish`].
+/// Callers that produce segments one at a time ([`split_literal`]) drive it
+/// directly and never build the segment vector.
+///
+/// One span under construction and the previous segment's flag are the only state. A span
+/// stays pending as long as the next segment might still extend it (a delimiter merging with
+/// its neighbour, a contiguous run); a segment that cannot extend it flushes it to `out` and
+/// takes its place.
+pub(crate) struct SplitFold<'a> {
+ out: &'a mut Vec,
+ behavior: SplitDelimiterBehavior,
+ pending: Option<(usize, usize)>,
+ previous_match: bool,
+}
+
+impl<'a> SplitFold<'a> {
+ pub(crate) fn new(out: &'a mut Vec, behavior: SplitDelimiterBehavior) -> Self {
+ Self {
+ out,
+ behavior,
+ pending: None,
+ previous_match: false,
}
- MergedWithNext => {
- let mut previous_match = false;
- let mut splits =
- matches
- .into_iter()
- .rev()
- .fold(vec![], |mut acc, (offsets, is_match)| {
- if is_match && !previous_match {
- if let Some(((start, _), _)) = acc.last_mut() {
- *start = offsets.0;
- } else {
- acc.push((offsets, false));
- }
- } else {
- acc.push((offsets, false));
- }
- previous_match = is_match;
- acc
- });
- splits.reverse();
- splits
+ }
+
+ pub(crate) fn segment(&mut self, offsets: (usize, usize), is_match: bool) {
+ use SplitDelimiterBehavior::*;
+ match self.behavior {
+ Isolated => self.emit(Some(offsets)),
+ Removed => {
+ if !is_match {
+ self.emit(Some(offsets));
+ }
+ }
+ // A run of equal flags becomes one span.
+ Contiguous => {
+ if is_match == self.previous_match {
+ self.extend_pending_to(offsets);
+ } else {
+ let done = self.pending.replace(offsets);
+ self.emit(done);
+ }
+ }
+ // A delimiter glues onto the piece before it; in a run of delimiters only the
+ // first one glues, the rest stand alone.
+ MergedWithPrevious => {
+ if is_match && !self.previous_match {
+ self.extend_pending_to(offsets);
+ } else {
+ let done = self.pending.replace(offsets);
+ self.emit(done);
+ }
+ }
+ // A delimiter glues onto the piece after it, so it stays pending until that
+ // piece arrives; a delimiter followed by another delimiter stands alone.
+ MergedWithNext => {
+ if !is_match && self.previous_match {
+ self.extend_pending_to(offsets);
+ } else {
+ let done = self.pending.replace(offsets);
+ self.emit(done);
+ }
+ }
}
- };
+ self.previous_match = is_match;
+ }
+
+ /// Flushes the last pending span.
+ pub(crate) fn finish(mut self) {
+ let last = self.pending.take();
+ self.emit(last);
+ }
- for ((start, end), should_remove) in splits {
- if !should_remove && start != end {
- out.push(Span {
+ /// Stretches the pending span to cover `offsets` too, or starts one from it.
+ fn extend_pending_to(&mut self, offsets: (usize, usize)) {
+ match &mut self.pending {
+ Some((_, end)) => *end = offsets.1,
+ None => self.pending = Some(offsets),
+ }
+ }
+
+ /// Empty pieces are dropped, matching the fold in `NormalizedString::split`.
+ fn emit(&mut self, span: Option<(usize, usize)>) {
+ if let Some((start, end)) = span
+ && start != end
+ {
+ self.out.push(Span {
start: start as u32,
end: end as u32,
});
@@ -1045,11 +1635,19 @@ impl Model for PipelineModel {
}
}
+/// A set of buffers and other state the model needs to encode efficiently,
+/// reused among calls to [`PipelineTokenizer::encode`].
+///
+/// Each model gets its own variant.
+#[derive(Default)]
pub enum PipelineModelScratch {
BPE(BpeScratch),
WordLevel(()),
WordPiece(WordPieceScratch),
Unigram(UnigramScratch),
+ /// We need a default value to be able to use [`mem::take`] in [`ScratchGuard::drop`]
+ #[default]
+ None,
}
impl ModelScratch for PipelineModelScratch {}
@@ -1100,11 +1698,504 @@ mod tests {
.iter()
.map(|t| t.id)
.collect();
- // Not the unk id: both the `Replace` and the `Split` really ran on the literal path.
+ // Not the unk id: the string-pattern config ran end-to-end with no regex backend (the
+ // `Replace` fuses into the Metaspace normalizer; the `Split` runs on the literal path).
assert_eq!(ids, [1, 2]);
assert_pipeline_matches_reference(&tok, "hello world");
}
+ /// llama-2's declared normalizer is `Prepend("▁")` then `Replace(" " -> "▁")`, two steps that
+ /// each copy the chunk. Building the pipeline must swap them for the one-pass
+ /// [`MetaspaceNormalizer`] (see [`MetaspaceNormalizer::fuse`]), and must not touch the source
+ /// config while doing so.
+ #[test]
+ fn try_from_fuses_the_declared_space_rewrite() {
+ let normalizer: NormalizerWrapper = serde_json::from_str(
+ r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#,
+ )
+ .unwrap();
+ let mut tok = wordlevel_tokenizer(vec![("", 0)], None);
+ tok.with_normalizer(Some(normalizer)).unwrap();
+
+ let declared = serde_json::to_string(tok.get_normalizer().unwrap()).unwrap();
+ let pipeline = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(
+ matches!(
+ pipeline.normalizers.as_slice(),
+ [PipelineNormalizer::Metaspace(_)]
+ ),
+ "{:?}",
+ pipeline.normalizers
+ );
+ assert_eq!(
+ serde_json::to_string(tok.get_normalizer().unwrap()).unwrap(),
+ declared,
+ );
+ }
+
+ /// The real llama-2 and gemma-4 files qualify for two independent rewrites: their declared
+ /// space rewrite fuses into the one-pass [`MetaspaceNormalizer`], and the whole-chunk text
+ /// they hand the model is cut into words by [`ProvenCuts`]. Both read the declared config,
+ /// not each other's output, so building the pipeline must apply both; this pins that neither
+ /// detection hides the other. Skipped when the files are not fetched.
+ #[test]
+ fn real_sentencepiece_configs_get_the_fused_normalizer_and_proven_cuts() {
+ for file in ["llama-2.json", "gemma-4.json"] {
+ let path = format!("../data/{file}");
+ if !std::path::Path::new(&path).exists() {
+ eprintln!("skip {file}: not present (fetch with `make bench-models`)");
+ continue;
+ }
+ let tok = crate::Tokenizer::from_file(&path).unwrap();
+ let pipeline = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(
+ matches!(
+ pipeline.normalizers.as_slice(),
+ [PipelineNormalizer::Metaspace(_)]
+ ),
+ "{file}: the declared space rewrite should fuse, got {:?}",
+ pipeline.normalizers,
+ );
+ assert!(
+ matches!(pipeline.pre_tokenizer, PipelinePreTokenizer::ProvenCuts(_)),
+ "{file}: the whole-chunk text should get proven cuts",
+ );
+ }
+ }
+
+ /// Normalization runs on each chunk between special tokens, so the fused prepend must too:
+ /// the text after `
` starts with a space, and `Prepend` + swap turn that into two
+ /// delimiters. Only `"▁▁world"` is in the vocabulary; a prepend skipping the marked chunk
+ /// would produce the unk id instead.
+ #[test]
+ fn fused_metaspace_normalizer_prepends_per_chunk() {
+ let normalizer: NormalizerWrapper = serde_json::from_str(
+ r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#,
+ )
+ .unwrap();
+ let mut tok = wordlevel_tokenizer(vec![("", 0), ("▁hello▁", 1), ("▁▁world", 2)], None);
+ tok.with_normalizer(Some(normalizer)).unwrap();
+ tok.with_pre_tokenizer(None::);
+ let _ = tok.add_special_tokens([crate::AddedToken::from("", true)]);
+
+ let ids: Vec = PipelineTokenizer::try_from(&tok)
+ .unwrap()
+ .encode("hello world", false)
+ .unwrap()
+ .iter()
+ .map(|t| t.id)
+ .collect();
+ assert_eq!(ids, [1, 3, 2]);
+ assert_pipeline_matches_reference(&tok, "hello world");
+ }
+
+ /// A `Replace` swapping anything but a single space stays a declared step, and the ids stay
+ /// those of the declared rewrite.
+ #[test]
+ fn a_replace_that_is_not_the_space_swap_is_left_declared() {
+ let normalizer: NormalizerWrapper =
+ serde_json::from_str(r#"{"type":"Replace","pattern":{"String":"x"},"content":"y"}"#)
+ .unwrap();
+ let mut tok = wordlevel_tokenizer(vec![("", 0), ("y a", 1)], None);
+ tok.with_normalizer(Some(normalizer)).unwrap();
+ tok.with_pre_tokenizer(None::);
+
+ let pipeline = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(
+ matches!(
+ pipeline.normalizers.as_slice(),
+ [PipelineNormalizer::Declared(_)]
+ ),
+ "{:?}",
+ pipeline.normalizers
+ );
+ let ids: Vec = pipeline
+ .encode("x a", false)
+ .unwrap()
+ .iter()
+ .map(|t| t.id)
+ .collect();
+ assert_eq!(ids, [1]);
+ assert_pipeline_matches_reference(&tok, "x a");
+ }
+
+ /// The metaspace-shaped BPE from the proven-cuts tests, plus an unk token and a `▁▁`
+ /// piece: the swap, the delimiter-group rule and the veto prefilter all have something
+ /// to act on.
+ fn metaspace_bpe_tokenizer(
+ normalizer_json: &str,
+ pre_tokenizer_json: Option<&str>,
+ ) -> Tokenizer {
+ use crate::models::bpe::{BpeBuilder, Merges, Vocab};
+
+ let vocab: Vocab = [
+ ("", 0u32),
+ ("▁", 1),
+ ("<", 2),
+ (">", 3),
+ ("/", 4),
+ ("s", 5),
+ ("p", 6),
+ ("a", 7),
+ ("b", 8),
+ ("▁", 10),
+ (">▁", "▁"),
+ (">▁", "".to_string())
+ .build()
+ .unwrap();
+ let mut tok = Tokenizer::new(bpe);
+ tok.with_normalizer(Some(
+ serde_json::from_str::(normalizer_json).unwrap(),
+ ))
+ .unwrap();
+ if let Some(json) = pre_tokenizer_json {
+ tok.with_pre_tokenizer(Some(
+ serde_json::from_str::(json).unwrap(),
+ ));
+ }
+ tok
+ }
+
+ /// The two spellings of the space rewrite that reach the proven cuts, copied from the
+ /// files in `data/`.
+ const LLAMA_STYLE_NORMALIZER: &str = r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#;
+ const GEMMA_STYLE_NORMALIZER: &str =
+ r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#;
+ const GEMMA_STYLE_PRE_TOKENIZER: &str = r#"{"type":"Split","pattern":{"String":" "},"behavior":"MergedWithPrevious","invert":false}"#;
+
+ /// Chunk shapes the raw-coordinate cuts must reproduce: leading, trailing and grouped
+ /// spaces, raw delimiters, the veto text (cut by the written path, kept whole by the
+ /// prefilter), characters outside the vocabulary, and no spaces at all.
+ const ZERO_COPY_TEXTS: &[&str] = &[
+ "sp a",
+ " sp",
+ "a b ",
+ " b a ",
+ "",
+ " ",
+ " ",
+ "sp",
+ "▁",
+ "▁a b",
+ "a▁b sp",
+ " ",
+ "a b",
+ "p a",
+ "a\tb\nsp",
+ "x y",
+ ];
+
+ /// Both space-rewrite spellings reach the fused-normalizer + proven-cuts shape, so both
+ /// must take the zero-copy path, and its ids must match the written rewrite's on every
+ /// chunk shape — including the veto text, where the two paths cut differently and only
+ /// the ids agree.
+ #[test]
+ fn zero_copy_ids_match_the_written_rewrite() {
+ for (name, normalizer, pre_tokenizer) in [
+ ("llama-2 shape", LLAMA_STYLE_NORMALIZER, None),
+ (
+ "gemma-4 shape",
+ GEMMA_STYLE_NORMALIZER,
+ Some(GEMMA_STYLE_PRE_TOKENIZER),
+ ),
+ ] {
+ let tok = metaspace_bpe_tokenizer(normalizer, pre_tokenizer);
+ let zero_copy = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(
+ zero_copy.zero_copy.is_some(),
+ "{name} should take the zero-copy path"
+ );
+ let mut written = PipelineTokenizer::try_from(&tok).unwrap();
+ written.disable_zero_copy();
+ for text in ZERO_COPY_TEXTS {
+ let got: Vec = zero_copy
+ .encode(text, false)
+ .unwrap()
+ .iter()
+ .map(|t| t.id)
+ .collect();
+ let want: Vec = written
+ .encode(text, false)
+ .unwrap()
+ .iter()
+ .map(|t| t.id)
+ .collect();
+ assert_eq!(got, want, "{name}: {text:?}");
+ assert_pipeline_matches_reference(&tok, text);
+ }
+ }
+ }
+
+ /// An added token flagged `normalized` is matched against rewritten text, which the
+ /// zero-copy path never builds; one such token must keep the written stages.
+ #[test]
+ fn a_normalized_added_token_keeps_the_written_rewrite() {
+ let mut tok = metaspace_bpe_tokenizer(LLAMA_STYLE_NORMALIZER, None);
+ let _ = tok.add_tokens([crate::AddedToken::from("spa", false).normalized(true)]);
+ let pipeline = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(pipeline.zero_copy.is_none());
+ assert_pipeline_matches_reference(&tok, "b spa sp a");
+ }
+
+ /// The chain hands a trailing pending rewrite back unwritten; `write` must produce
+ /// exactly what running the chain as declared produces.
+ #[test]
+ fn a_chain_ending_in_a_pending_rewrite_leaves_it_unwritten() {
+ let chain = [
+ PipelineNormalizer::Declared(serde_json::from_str(r#"{"type":"Lowercase"}"#).unwrap()),
+ PipelineNormalizer::Metaspace(MetaspaceNormalizer::new(
+ '▁',
+ PrependMode::Unconditional,
+ false,
+ )),
+ ];
+ for text in ["Hello World", "", " A ", "no_spaces", "▁Marked"] {
+ let normalized = NormalizedText::from_chain(&chain, text).unwrap();
+ assert!(normalized.pending.is_some(), "{text:?}");
+ assert_eq!(
+ NormalizedText::from_chain(&chain, text)
+ .unwrap()
+ .write()
+ .unwrap(),
+ normalize_all(&chain, text).unwrap(),
+ "{text:?}"
+ );
+ }
+ }
+
+ /// A pending rewrite anywhere but last runs as declared: the step after it needs the
+ /// written text.
+ #[test]
+ fn a_mid_chain_rewrite_is_written() {
+ let chain = [
+ PipelineNormalizer::Metaspace(MetaspaceNormalizer::new(
+ '▁',
+ PrependMode::Unconditional,
+ false,
+ )),
+ PipelineNormalizer::Declared(serde_json::from_str(r#"{"type":"Lowercase"}"#).unwrap()),
+ ];
+ let normalized = NormalizedText::from_chain(&chain, "Hello World").unwrap();
+ assert!(normalized.pending.is_none());
+ assert_eq!(
+ NormalizedText::from_chain(&chain, "Hello World")
+ .unwrap()
+ .write()
+ .unwrap(),
+ normalize_all(&chain, "Hello World").unwrap()
+ );
+ }
+
+ /// A `Replace` swapping one character for another is a pending rewrite too; anything
+ /// changing the text's structure is not.
+ #[test]
+ fn a_single_char_replace_is_a_pending_rewrite() {
+ let single: NormalizerWrapper =
+ serde_json::from_str(r#"{"type":"Replace","pattern":{"String":"x"},"content":"y"}"#)
+ .unwrap();
+ assert_eq!(
+ Normalizer::pending_rewrite(&single),
+ Some(PendingRewrite {
+ from: 'x',
+ to: 'y',
+ prepend: PrependMode::Never,
+ })
+ );
+ let refused = [
+ (
+ "a multi-char pattern",
+ r#"{"type":"Replace","pattern":{"String":"xy"},"content":"y"}"#,
+ ),
+ (
+ "multi-char content",
+ r#"{"type":"Replace","pattern":{"String":"x"},"content":"yy"}"#,
+ ),
+ (
+ "a deletion",
+ r#"{"type":"Replace","pattern":{"String":"x"},"content":""}"#,
+ ),
+ ];
+ for (name, json) in refused {
+ let replace: NormalizerWrapper = serde_json::from_str(json).unwrap();
+ assert_eq!(Normalizer::pending_rewrite(&replace), None, "{name}");
+ }
+ }
+
+ /// Skipped when the fixtures have not been fetched.
+ #[test]
+ fn zero_copy_is_built_for_the_real_sentencepiece_configs() {
+ for file in ["llama-2.json", "gemma-4.json"] {
+ let path = format!("../data/{file}");
+ if !std::path::Path::new(&path).exists() {
+ eprintln!("skip {file}: not present (fetch with `make bench-models`)");
+ continue;
+ }
+ let tok = crate::Tokenizer::from_file(&path).unwrap();
+ let pipeline = PipelineTokenizer::try_from(&tok).unwrap();
+ assert!(pipeline.zero_copy.is_some(), "{file}");
+ }
+ }
+
+ /// The benchmark's stage ladder subtracts one rung from the next, which is only
+ /// meaningful if every rung runs the route a full encode runs. These three pin that
+ /// for each route, on the real configs.
+ ///
+ /// Without them the ladder can silently drift back to timing a splitter nothing uses:
+ /// the `pipeline_oracle` tests only ever drive `STAGE_POSTPROCESS`, so they prove ids
+ /// and say nothing about which path produced them.
+ mod stage_ladder {
+ use super::*;
+
+ fn real_pipeline(file: &str) -> Option {
+ let path = format!("../data/{file}");
+ if !std::path::Path::new(&path).exists() {
+ eprintln!("skip {file}: not present (fetch with `make bench-models`)");
+ return None;
+ }
+ let tok = crate::Tokenizer::from_file(&path).unwrap();
+ Some(PipelineTokenizer::try_from(&tok).unwrap())
+ }
+
+ const TEXT: &str =
+ "The quick brown fox jumps 123 times, don't it?\n\tif x == 1: return 'a'\n";
+
+ fn run(pipeline: &PipelineTokenizer) -> (Vec