feat: whitespace optimizer

b00f · b00f · commit 971be68b499e · 2025-08-06T17:02:24.000+08:00
diff --git a/examples/whitespace_optimization_demo.rs b/examples/whitespace_optimization_demo.rs
@@ -0,0 +1,116 @@
+use std::time::Instant;
+use tokenizers::pre_tokenizers::whitespace::{Whitespace, WhitespaceOptimized};
+use tokenizers::{OffsetReferential, OffsetType, PreTokenizer, PreTokenizedString};
+
+fn main() {
+    println!("Whitespace Pre-Tokenizer Optimization Demo");
+    println!("==========================================\n");
+
+    // Test cases with different characteristics
+    let test_cases = vec![
+        ("Simple text", "Hello world! How are you doing?"),
+        ("Mixed content", "This is a test with numbers 123 and symbols @#$% and unicode: café résumé"),
+        ("Whitespace heavy", "Multiple    spaces\tand\nnewlines\r\nhere"),
+        ("Symbol heavy", "Hello!@#$%^&*()world?><>{}[]|\\"),
+        ("Word heavy", "This is a very long sentence with many words that should be tokenized properly"),
+        ("Unicode heavy", "αβγ δέζ ηθι κλμ νξο πρσ τυφ χψω"),
+        ("Mixed unicode", "Hello 123 αβγ !@# world δέζ ηθι"),
+    ];
+
+    for (name, text) in test_cases {
+        println!("Test case: {}", name);
+        println!("Input: '{}'", text);
+
+        // Test original implementation
+        let start = Instant::now();
+        let mut original = PreTokenizedString::from(text);
+        let original_pretok = Whitespace {};
+        original_pretok.pre_tokenize(&mut original).unwrap();
+        let original_duration = start.elapsed();
+
+        let original_splits = original
+            .get_splits(OffsetReferential::Original, OffsetType::Byte)
+            .into_iter()
+            .map(|(s, o, _)| (s, o))
+            .collect::<Vec<_>>();
+
+        // Test optimized implementation
+        let start = Instant::now();
+        let mut optimized = PreTokenizedString::from(text);
+        let optimized_pretok = WhitespaceOptimized {};
+        optimized_pretok.pre_tokenize(&mut optimized).unwrap();
+        let optimized_duration = start.elapsed();
+
+        let optimized_splits = optimized
+            .get_splits(OffsetReferential::Original, OffsetType::Byte)
+            .into_iter()
+            .map(|(s, o, _)| (s, o))
+            .collect::<Vec<_>>();
+
+        // Verify compatibility
+        let compatible = original_splits == optimized_splits;
+
+        println!("  Original tokens: {:?}", original_splits);
+        println!("  Optimized tokens: {:?}", optimized_splits);
+        println!("  Compatible: {}", compatible);
+        println!("  Original time: {:?}", original_duration);
+        println!("  Optimized time: {:?}", optimized_duration);
+
+        if original_duration > optimized_duration {
+            let speedup = original_duration.as_nanos() as f64 / optimized_duration.as_nanos() as f64;
+            println!("  Speedup: {:.2}x", speedup);
+        } else {
+            let slowdown = optimized_duration.as_nanos() as f64 / original_duration.as_nanos() as f64;
+            println!("  Slowdown: {:.2}x", slowdown);
+        }
+        println!();
+    }
+
+    // Performance test with large text
+    println!("Large text performance test:");
+    let base_text = "Hello world! This is a test with numbers 123 and symbols @#$% and unicode: café résumé. ";
+    let large_text: String = base_text.repeat(1000); // ~50KB of text
+    println!("Text size: {} bytes", large_text.len());
+
+    // Warm up
+    for _ in 0..10 {
+        let mut _warmup = PreTokenizedString::from(&large_text);
+        let _pretok = Whitespace {};
+        // Don't actually call pre_tokenize to avoid affecting results
+    }
+
+    // Benchmark original
+    let iterations = 100;
+    let start = Instant::now();
+    for _ in 0..iterations {
+        let mut pretokenized = PreTokenizedString::from(&large_text);
+        let pretok = Whitespace {};
+        pretok.pre_tokenize(&mut pretokenized).unwrap();
+    }
+    let original_total = start.elapsed();
+    let original_avg = original_total / iterations;
+
+    // Benchmark optimized
+    let start = Instant::now();
+    for _ in 0..iterations {
+        let mut pretokenized = PreTokenizedString::from(&large_text);
+        let pretok = WhitespaceOptimized {};
+        pretok.pre_tokenize(&mut pretokenized).unwrap();
+    }
+    let optimized_total = start.elapsed();
+    let optimized_avg = optimized_total / iterations;
+
+    println!("  Original average: {:?}", original_avg);
+    println!("  Optimized average: {:?}", optimized_avg);
+
+    if original_avg > optimized_avg {
+        let speedup = original_avg.as_nanos() as f64 / optimized_avg.as_nanos() as f64;
+        println!("  Overall speedup: {:.2}x", speedup);
+    } else {
+        let slowdown = optimized_avg.as_nanos() as f64 / original_avg.as_nanos() as f64;
+        println!("  Overall slowdown: {:.2}x", slowdown);
+    }
+
+    println!("\nNote: Performance results may vary depending on hardware and system load.");
+    println!("For accurate benchmarks, run: cargo bench --bench whitespace_benchmark");
+}
diff --git a/tokenizers/benches/whitespace_benchmark.rs b/tokenizers/benches/whitespace_benchmark.rs
@@ -0,0 +1,103 @@
+#[macro_use]
+extern crate criterion;
+
+use criterion::{Criterion, Throughput};
+use tokenizers::pre_tokenizers::whitespace::{Whitespace, WhitespaceOptimized};
+use tokenizers::{OffsetReferential, OffsetType, PreTokenizer, PreTokenizedString};
+
+fn bench_whitespace_comparison(c: &mut Criterion) {
+    let mut group = c.benchmark_group("whitespace-pre-tokenizers");
+
+    // Test data with various characteristics
+    let test_cases = vec![
+        ("simple", "Hello world! How are you doing?"),
+        ("mixed", "This is a test with numbers 123 and symbols @#$% and unicode: café résumé"),
+        ("whitespace_heavy", "Multiple    spaces\tand\nnewlines\r\nhere"),
+        ("symbol_heavy", "Hello!@#$%^&*()world?><>{}[]|\\"),
+        ("word_heavy", "This is a very long sentence with many words that should be tokenized properly"),
+        ("unicode_heavy", "αβγ δέζ ηθι κλμ νξο πρσ τυφ χψω"),
+        ("mixed_unicode", "Hello 123 αβγ !@# world δέζ ηθι"),
+    ];
+
+    for (name, text) in test_cases {
+        let data_len = text.len() as u64;
+        group.throughput(Throughput::Bytes(data_len));
+
+        // Benchmark original regex-based implementation
+        group.bench_function(&format!("{}-original", name), |b| {
+            b.iter(|| {
+                let mut pretokenized = PreTokenizedString::from(text);
+                let pretok = Whitespace {};
+                pretok.pre_tokenize(&mut pretokenized).unwrap();
+                let _result = pretokenized
+                    .get_splits(OffsetReferential::Original, OffsetType::Byte)
+                    .into_iter()
+                    .map(|(s, o, _)| (s, o))
+                    .collect::<Vec<_>>();
+            })
+        });
+
+        // Benchmark optimized byte-level implementation
+        group.bench_function(&format!("{}-optimized", name), |b| {
+            b.iter(|| {
+                let mut pretokenized = PreTokenizedString::from(text);
+                let pretok = WhitespaceOptimized {};
+                pretok.pre_tokenize(&mut pretokenized).unwrap();
+                let _result = pretokenized
+                    .get_splits(OffsetReferential::Original, OffsetType::Byte)
+                    .into_iter()
+                    .map(|(s, o, _)| (s, o))
+                    .collect::<Vec<_>>();
+            })
+        });
+    }
+
+    group.finish();
+}
+
+fn bench_large_text(c: &mut Criterion) {
+    let mut group = c.benchmark_group("whitespace-large-text");
+
+    // Create a large text by repeating patterns
+    let base_text = "Hello world! This is a test with numbers 123 and symbols @#$% and unicode: café résumé. ";
+    let large_text: String = base_text.repeat(1000); // ~50KB of text
+    let data_len = large_text.len() as u64;
+
+    group.throughput(Throughput::Bytes(data_len));
+
+    group.bench_function("large-original", |b| {
+        b.iter(|| {
+            let mut pretokenized = PreTokenizedString::from(large_text.as_str());
+            let pretok = Whitespace {};
+            pretok.pre_tokenize(&mut pretokenized).unwrap();
+            let _result = pretokenized
+                .get_splits(OffsetReferential::Original, OffsetType::Byte)
+                .into_iter()
+                .map(|(s, o, _)| (s, o))
+                .collect::<Vec<_>>();
+        })
+    });
+
+    group.bench_function("large-optimized", |b| {
+        b.iter(|| {
+            let mut pretokenized = PreTokenizedString::from(large_text.as_str());
+            let pretok = WhitespaceOptimized {};
+            pretok.pre_tokenize(&mut pretokenized).unwrap();
+            let _result = pretokenized
+                .get_splits(OffsetReferential::Original, OffsetType::Byte)
+                .into_iter()
+                .map(|(s, o, _)| (s, o))
+                .collect::<Vec<_>>();
+        })
+    });
+
+    group.finish();
+}
+
+criterion_group! {
+    name = whitespace_benches;
+    config = Criterion::default().sample_size(20);
+    targets = bench_whitespace_comparison, bench_large_text
+}
+
+criterion_main!(whitespace_benches);
diff --git a/tokenizers/src/pre_tokenizers/whitespace.rs b/tokenizers/src/pre_tokenizers/whitespace.rs
