gigatoken 0.2.2 → 0.4.0

data/benches/pretokenize.rs

@@ -0,0 +1,151 @@
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+#[cfg(all(
+    target_arch = "x86_64",
+    target_feature = "avx512bw",
+    target_feature = "avx512vl"
+))]
+use gigatoken_rs::pretokenize::reference::avx512::Avx512PretokenizerIter;
+use gigatoken_rs::pretokenize::{
+    reference::combinator::pretokens_iterator, FastCl100kPretokenizer, FastQwen2Pretokenizer,
+    FastQwen35Pretokenizer, FastR50kPretokenizer, PretokenizerIter,
+};
+use gigatoken_rs::pretokenize::reference::simd::SimdPretokIter;
+use std::hint::black_box;
+
+const TARGET_BENCH_SIZE: usize = 100_000_000; // ~100 MB
+
+/// Load OWT data, truncated to a UTF-8-safe boundary near `max_bytes`.
+fn load_owt(max_bytes: usize) -> Vec<u8> {
+    let data_dir = std::env::home_dir().unwrap().join("data");
+    let all_bytes =
+        std::fs::read(data_dir.join("owt_train.txt")).expect("Could not read ~/data/owt_train.txt");
+    let mut end = max_bytes.min(all_bytes.len());
+    // Back up to a UTF-8 character boundary
+    while end > 0 && !std::str::from_utf8(&all_bytes[..end]).is_ok() {
+        end -= 1;
+    }
+    all_bytes[..end].to_vec()
+}
+
+fn pretokenize_benches(c: &mut Criterion) {
+    let input = load_owt(TARGET_BENCH_SIZE);
+    let input_len = input.len() as u64;
+    eprintln!("Benchmark input size: {:.1} MB", input_len as f64 / 1e6);
+
+    let mut group = c.benchmark_group("pretokenize");
+    group.throughput(Throughput::Bytes(input_len));
+    group.sample_size(10);
+
+    group.bench_function("state_machine", |b| {
+        b.iter(|| {
+            let count = PretokenizerIter::new(&input).count();
+            black_box(count);
+        });
+    });
+
+    group.bench_function("winnow", |b| {
+        b.iter(|| {
+            let mut input_str = unsafe { std::str::from_utf8_unchecked(&input) };
+            let count = pretokens_iterator(&mut input_str).count();
+            black_box(count);
+        });
+    });
+
+    #[cfg(all(
+        target_arch = "x86_64",
+        target_feature = "avx512bw",
+        target_feature = "avx512vl"
+    ))]
+    group.bench_function("avx512", |b| {
+        b.iter(|| {
+            let mut iter = Avx512PretokenizerIter::new(&input);
+            let mut count = 0;
+            while iter.next().is_some() {
+                count += 1;
+            }
+            black_box(count);
+        });
+    });
+    group.bench_function("simd", |b| {
+        b.iter(|| {
+            let count = SimdPretokIter::new(&input).count();
+            black_box(count);
+        });
+    });
+
+    group.bench_function("fast_scalar", |b| {
+        b.iter(|| {
+            let mut iter = FastR50kPretokenizer::new(&input);
+            let mut count = 0;
+            while iter.next().is_some() {
+                count += 1;
+            }
+            black_box(count);
+        });
+    });
+
+    group.bench_function("cl100k_fast_scalar", |b| {
+        b.iter(|| {
+            let mut iter = FastCl100kPretokenizer::new(&input);
+            let mut count = 0;
+            while iter.next().is_some() {
+                count += 1;
+            }
+            black_box(count);
+        });
+    });
+
+    group.bench_function("qwen2_fast_scalar", |b| {
+        b.iter(|| {
+            let mut iter = FastQwen2Pretokenizer::new(&input);
+            let mut count = 0;
+            while iter.next().is_some() {
+                count += 1;
+            }
+            black_box(count);
+        });
+    });
+
+    group.bench_function("qwen3_5_fast_scalar", |b| {
+        b.iter(|| {
+            let mut iter = FastQwen35Pretokenizer::new(&input);
+            let mut count = 0;
+            while iter.next().is_some() {
+                count += 1;
+            }
+            black_box(count);
+        });
+    });
+
+    let re = fancy_regex::Regex::new(
+        r"'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+",
+    )
+    .unwrap();
+
+    group.bench_function("regex", |b| {
+        b.iter(|| {
+            let text = unsafe { std::str::from_utf8_unchecked(&input) };
+            let count = re.find_iter(text).count();
+            black_box(count);
+        });
+    });
+
+    // Backtracking-compatible equivalent of the possessive cl100k pattern
+    let re_cl100k = fancy_regex::Regex::new(
+        r"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s+$|\s*[\r\n]|\s+(?!\S)|\s+",
+    )
+    .unwrap();
+
+    group.bench_function("cl100k_regex", |b| {
+        b.iter(|| {
+            let text = unsafe { std::str::from_utf8_unchecked(&input) };
+            let count = re_cl100k.find_iter(text).count();
+            black_box(count);
+        });
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, pretokenize_benches);
+criterion_main!(benches);

data/benches/pretokenize_profile.rs

@@ -0,0 +1,36 @@
+//! Profiling target for the `FastR50kPretokenizer` hot loop in isolation: a plain
+//! single-pass `main` (no criterion, no BPE encode) that `black_box`es every
+//! yielded pretoken slice, so the slice production can't be optimized away.
+
+use gigatoken_rs::pretokenize::FastR50kPretokenizer;
+use std::hint::black_box;
+use std::time::Instant;
+
+mod common;
+fn main() {
+    let input = common::load_owt_input(None);
+    let size_gb = input.len() as f64 / 1e9;
+
+    // Feed the entire buffer to one pretokenizer in a single pass — this matches
+    // the real encode path (`pretokenize_as_iter(text.as_bytes())`), which does
+    // not pre-split on newlines.
+    let buf: &[u8] = &input;
+
+    eprintln!("Pretokenizing (fast_scalar, single-threaded, whole buffer)...");
+    let start = Instant::now();
+    let mut total_tokens: usize = 0;
+    // Hand each real pretoken slice to black_box so the bounds computation can't
+    // be optimized down to a counter.
+    let mut iter = FastR50kPretokenizer::new(buf);
+    for pretoken in iter {
+        black_box(pretoken);
+        total_tokens += 1;
+    }
+    let elapsed = start.elapsed().as_secs_f64();
+    let throughput_gb = size_gb / elapsed;
+
+    eprintln!(
+        "{total_tokens} tokens in {elapsed:.2}s — {throughput_gb:.2} GB/s ({:.0} MB/s)",
+        throughput_gb * 1000.0
+    );
+}

data/benches/pretokenize_profile_all.rs

@@ -0,0 +1,52 @@
+//! Per-scheme variant of `pretokenize_profile`: same single-pass loop over
+//! OWT with every yielded pretoken black_boxed, scheme selected via the
+//! SCHEME env var (r50k | cl100k | olmo3 | qwen2 | qwen3_5 | deepseek_v3).
+//! r50k is also what ByteLevel tokenizers like ModernBERT resolve to. Used
+//! for interleaved A/B runs of the mask-scanner schemes.
+
+use gigatoken_rs::pretokenize::{
+    FastCl100kPretokenizer, FastDeepSeekV3Pretokenizer, FastOlmo3Pretokenizer,
+    FastQwen2Pretokenizer, FastQwen35Pretokenizer, FastR50kPretokenizer,
+};
+use std::hint::black_box;
+use std::time::Instant;
+
+mod common;
+
+macro_rules! drive {
+    ($ty:ty, $buf:expr) => {{
+        let mut total_tokens: usize = 0;
+        let mut iter = <$ty>::new($buf);
+        while let Some(pretoken) = iter.next() {
+            black_box(pretoken);
+            total_tokens += 1;
+        }
+        total_tokens
+    }};
+}
+
+fn main() {
+    let input = common::load_owt_input(None);
+    let size_gb = input.len() as f64 / 1e9;
+    let buf: &[u8] = &input;
+    let scheme = std::env::var("SCHEME").unwrap_or_else(|_| "r50k".to_string());
+
+    eprintln!("Pretokenizing ({scheme}, single-threaded, whole buffer)...");
+    let start = Instant::now();
+    let total_tokens = match scheme.as_str() {
+        "r50k" => drive!(FastR50kPretokenizer, buf),
+        "cl100k" => drive!(FastCl100kPretokenizer, buf),
+        "olmo3" => drive!(FastOlmo3Pretokenizer, buf),
+        "qwen2" => drive!(FastQwen2Pretokenizer, buf),
+        "qwen3_5" => drive!(FastQwen35Pretokenizer, buf),
+        "deepseek_v3" => drive!(FastDeepSeekV3Pretokenizer, buf),
+        other => panic!("unknown SCHEME {other:?}"),
+    };
+    let elapsed = start.elapsed().as_secs_f64();
+    let throughput_gb = size_gb / elapsed;
+
+    eprintln!(
+        "{total_tokens} tokens in {elapsed:.2}s — {throughput_gb:.2} GB/s ({:.0} MB/s)",
+        throughput_gb * 1000.0
+    );
+}

data/benches/simdutf_transcode.rs

@@ -0,0 +1,82 @@
+use criterion::{Criterion, Throughput, criterion_group, criterion_main};
+use rayon::prelude::*;
+use std::hint::black_box;
+
+const TARGET_BENCH_SIZE: usize = 100_000_000; // ~100 MB
+
+/// Load OWT data, truncated to a UTF-8-safe boundary near `max_bytes`.
+fn load_owt(max_bytes: usize) -> Vec<u8> {
+    let data_dir = std::env::home_dir().unwrap().join("data");
+    let all_bytes =
+        std::fs::read(data_dir.join("owt_train.txt")).expect("Could not read ~/data/owt_train.txt");
+    let mut end = max_bytes.min(all_bytes.len());
+    while end > 0 && std::str::from_utf8(&all_bytes[..end]).is_err() {
+        end -= 1;
+    }
+    all_bytes[..end].to_vec()
+}
+
+fn simdutf_transcode_benches(c: &mut Criterion) {
+    let input = load_owt(TARGET_BENCH_SIZE);
+    let input_len = input.len() as u64;
+    eprintln!("Benchmark input size: {:.1} MB", input_len as f64 / 1e6);
+
+    let mut group = c.benchmark_group("simdutf_transcode");
+    group.throughput(Throughput::Bytes(input_len));
+    group.sample_size(10);
+
+    for num_threads in [1, 2, 4, 8, 12, 16] {
+        let pool = rayon::ThreadPoolBuilder::new()
+            .num_threads(num_threads)
+            .build()
+            .unwrap();
+
+        let chunk_size = input.len() / num_threads;
+
+        // Pre-compute chunk boundaries aligned to UTF-8 char boundaries.
+        let mut boundaries = Vec::with_capacity(num_threads + 1);
+        boundaries.push(0);
+        for i in 1..num_threads {
+            let mut pos = chunk_size * i;
+            while pos < input.len() && input[pos] & 0b1100_0000 == 0b1000_0000 {
+                pos += 1;
+            }
+            boundaries.push(pos);
+        }
+        boundaries.push(input.len());
+
+        // Pre-allocate one destination buffer per thread.
+        let mut dst_bufs: Vec<Vec<u32>> = boundaries
+            .windows(2)
+            .map(|w| vec![0u32; simdutf::utf32_length_from_utf8(&input[w[0]..w[1]])])
+            .collect();
+
+        group.bench_function(format!("utf8_to_utf32_{num_threads}t"), |b| {
+            b.iter(|| {
+                pool.install(|| {
+                    let total: usize = boundaries
+                        .par_windows(2)
+                        .zip(&mut dst_bufs)
+                        .map(|(w, dst)| {
+                            let chunk = &input[w[0]..w[1]];
+                            let written = unsafe {
+                                simdutf::convert_valid_utf8_to_utf32(
+                                    chunk.as_ptr(),
+                                    chunk.len(),
+                                    dst.as_mut_ptr(),
+                                )
+                            };
+                            black_box(written)
+                        })
+                        .sum();
+                    black_box(total);
+                })
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, simdutf_transcode_benches);
+criterion_main!(benches);

data/benches/unicode.rs

@@ -0,0 +1,88 @@
+use criterion::{Criterion, criterion_group, criterion_main};
+use icu::properties::{CodePointMapDataBorrowed, props::EnumeratedProperty};
+use std::hint::black_box;
+
+use rand::{self, Rng, RngExt};
+
+// pub fn fibonacci(n: u64) -> u64 {
+//     let mut a = 0;
+//     let mut b = 1;
+//     for _ in 0..n {
+//         let c = a + b;
+//         a = b;
+//         b = c;
+//     }
+//     a
+// }
+
+// Removed dependency since icu is ~95% faster
+// use unicode_properties::{GeneralCategoryGroup, UnicodeGeneralCategory};
+// pub fn unicode_properties_classify(c: char) -> bool {
+//     c.general_category_group() == GeneralCategoryGroup::Letter
+// }
+
+pub fn icu4x_classify(c: char) -> bool {
+    icu::properties::props::GeneralCategoryGroup::Letter
+        .contains(icu::properties::props::GeneralCategory::for_char(c))
+}
+
+pub fn icu4x_classify_table(c: char) -> bool {
+    let gc: CodePointMapDataBorrowed<icu::properties::props::GeneralCategory> =
+        icu::properties::CodePointMapData::new();
+    icu::properties::props::GeneralCategoryGroup::Letter.contains(gc.get(c))
+}
+
+pub fn criterion_benchmark(c: &mut Criterion) {
+    // c.bench_function("fib 20", |b| b.iter(|| fibonacci(black_box(20))));
+    let mut group = c.benchmark_group("unicode_classify");
+
+    let chars_input: Vec<char> = rand::rng()
+        .sample_iter::<char, _>(rand::distr::StandardUniform)
+        .take(4096)
+        .collect();
+    group.bench_with_input("icu4x", chars_input.as_slice(), |b, chars: &[char]| {
+        b.iter(|| {
+            for c in chars {
+                icu4x_classify(*c);
+            }
+        });
+    });
+    // group.bench_with_input(
+    //     "unicode_properties",
+    //     chars_input.as_slice(),
+    //     |b, chars: &[char]| {
+    //         b.iter(|| {
+    //             for c in chars {
+    //                 unicode_properties_classify(*c);
+    //             }
+    //         });
+    //     },
+    // );
+    group.bench_with_input(
+        "icu4x table",
+        chars_input.as_slice(),
+        |b, chars: &[char]| {
+            b.iter(|| {
+                for c in chars {
+                    icu4x_classify_table(*c);
+                }
+            });
+        },
+    );
+
+    // c.bench_function("unicode classify letter", |b| {
+    //     b.iter_batched(
+    //         || {},
+    //         |chars: Vec<char>| {
+    //             for c in chars {
+    //                 unicode_classify(c);
+    //             }
+    //         },
+    //         criterion::BatchSize::SmallInput,
+    //     )
+    // });
+    // c.bench_function("unicode pretokenize")
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

data/ext/gigatoken/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gigatoken-rb"
-version = "0.2.2"
+version = "0.4.0"
 edition = "2021"
 
 [lib]

data/ext/gigatoken/src/error.rs

@@ -1,20 +1,40 @@
-//! `Gigatoken::Error` is defined in Ruby (`lib/gigatoken.rb`), loaded before
-//! this extension is required — looked up fresh at each raise site rather
-//! than cached, so no magnus `Value` needs a GC-registered static home.
+//! `Gigatoken::Error` and its subclasses are defined in Ruby
+//! (`lib/gigatoken.rb`), loaded before this extension is required — looked up
+//! fresh at each raise site rather than cached, so no magnus `Value` needs a
+//! GC-registered static home.
 
 use magnus::{exception::ExceptionClass, prelude::*, Error, RModule, Ruby};
 
-fn error_class(ruby: &Ruby) -> Result<ExceptionClass, Error> {
+fn error_class(ruby: &Ruby, name: &str) -> Result<ExceptionClass, Error> {
     ruby.class_object()
         .const_get::<_, RModule>("Gigatoken")?
-        .const_get("Error")
+        .const_get(name)
 }
 
-/// Raise `Gigatoken::Error` with `message` (core load/encode failures surface
-/// through this — never a Rust panic across the Ruby boundary).
-pub fn raise(ruby: &Ruby, message: impl Into<String>) -> Error {
-    match error_class(ruby) {
+fn raise_as(ruby: &Ruby, name: &str, message: impl Into<String>) -> Error {
+    match error_class(ruby, name) {
         Ok(class) => Error::new(class, message.into()),
         Err(e) => e,
     }
 }
+
+/// Raise `Gigatoken::Error` with `message` — what a failure that is neither
+/// the caller's document nor the model itself surfaces through (never a Rust
+/// panic across the Ruby boundary). Prefer [`input_error`] or [`model_error`]
+/// where one of them fits.
+pub fn raise(ruby: &Ruby, message: impl Into<String>) -> Error {
+    raise_as(ruby, "Error", message)
+}
+
+/// Raise `Gigatoken::InputError`: a document the tokenizer cannot take —
+/// invalid UTF-8 on the SentencePiece path, an id outside the vocabulary in
+/// `decode`.
+pub fn input_error(ruby: &Ruby, message: impl Into<String>) -> Error {
+    raise_as(ruby, "InputError", message)
+}
+
+/// Raise `Gigatoken::ModelError`: a tokenizer that cannot be loaded — bad or
+/// hostile JSON, an unknown pretokenizer scheme, a malformed `.tiktoken`.
+pub fn model_error(ruby: &Ruby, message: impl Into<String>) -> Error {
+    raise_as(ruby, "ModelError", message)
+}