🦀 Rust BPE Tokenizer (rust-bpe-tokenizer)

High-performance Byte-Pair Encoding (BPE) tokenizer implemented in Rust with Python FFI bindings (via PyO3) for Tolstoy-CLI.

🚀 Key Features

• Extreme Speed: ~7–10x faster training and ~10–20x faster inference by bypassing Python GIL and leveraging Rayon multi-threading.
• Drop-in replacement: 100% compatible Python API (train, encode, decode, save, load).
• Standalone Rust API: Pure Rust file IO (BpeCore::save, BpeCore::load, to_bytes, from_bytes) without Python runtime dependency.
• Python Protocols Support: Native support for pickle, copy.deepcopy, and multiprocessing (ProcessPoolExecutor).
• Dual Inference Modes: Parallel Rayon mode (use_parallel=True) and sequential mode (use_parallel=False) for low-overhead tokenization of short texts.
• Memory Safety: 100% safe, zero segmentation faults or memory leaks.
• Compact Persistence: Bincode binary format (.bpe) with automatic fallback migration for legacy Python .pkl files.

📦 Installation

pip install .

💻 Python Usage

import pickle, copy
from rust_bpe_tokenizer import BPETokenizer

tokenizer = BPETokenizer()

# Train on text corpus
# Note: n_workers is accepted for API compatibility with the Python
# implementation and ignored — the Rayon thread pool is managed automatically.
tokenizer.train("Hello world! High performance BPE tokenizer.", vocab_size=300)

# Vocabulary inspection
print(f"Vocab size: {len(tokenizer)}, Merges: {tokenizer.merges_len}")

# Encode text into token IDs
tokens = tokenizer.encode("Hello world!", use_parallel=True)
print("Token IDs:", tokens)

# Decode back to text or raw bytes
decoded_text = tokenizer.decode(tokens)
raw_bytes = tokenizer.decode_bytes(tokens)
print("Decoded text:", decoded_text)

# Save & Load
tokenizer.save("model.bpe")
loaded_tokenizer = BPETokenizer.load("model.bpe")

# Multiprocessing & Pickle support
pickled = pickle.dumps(tokenizer)
restored = pickle.loads(pickled)

🦀 Pure Rust Usage

use rust_bpe_tokenizer::BpeCore;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let text = "High-performance Rust BPE tokenizer.";
    let mut tok = BpeCore::new();

    // Training
    tok.train_core::<fn(f64, &str)>(text, 300, None)?;

    // Encoding & Decoding
    let tokens = tok.encode(text)?;
    let decoded = tok.decode(&tokens);
    assert_eq!(decoded, text);

    // Save and Load binary model
    tok.save("model.bpe")?;
    let loaded = BpeCore::load("model.bpe")?;
    
    Ok(())
}

📊 Benchmarks

Run the benchmark suite:

python benchmarks/bench_bpe.py

• Training: ~7.1x speedup (up to 5.6 MB/s)
• Encoding (Rayon Parallel): ~10.4x speedup (up to 9.4 MB/s)
• Decoding: ~3.2x speedup
• Bincode Model Loading: ~2.3x speedup compared to Python pickle