$ show karpathy/rustbpe
rustbpe is a Rust library with Python bindings that implements a BPE (Byte Pair Encoding) tokenizer for training and inference. The system provides efficient training of BPE models using parallel processing (Rayon), GPT-4 style regex pre-tokenization, and incremental update algorithms to avoid recomputing pair counts from scratch. It exports trained tokenizers to TikTok format for fast inference. The architecture centers around a Tokenizer struct that maintains merge rules and patterns, with methods for training, encoding, decoding, and batch processing. The Python interface mirrors the core Rust implementation, allowing seamless integration into Python NLP workflows while leveraging Rust's performance for computationally intensive operations.
The LLM identified 4 files worth highlighting for reuse:
The train_core_incremental method and its helpers (count_pairs_parallel, Word::merge_pair) implement an optimized BPE training algorithm that avoids O(n^2) recomputation. It uses parallel processing for initial counting, a dary heap for efficient max-pair selection, and incremental delta updates to maintain pair counts after each merge. This approach significantly reduces training time compared to naive implementations while maintaining correctness.
The train_from_iterator method demonstrates a sophisticated pattern for bridging Python iterators with Rust parallel processing. It buffers data from Python under the GIL, releases the GIL for Rayon-based parallel processing of text chunks, and efficiently merges results. This pattern enables high-throughput text processing while safely managing the Python-Rust boundary, making it reusable for other NLP tasks requiring parallel execution.
The count_pairs_parallel function efficiently counts token pairs across text segments while tracking their positions for incremental updates. Using ahash for performance and parallel processing via Rayon, it returns both pair counts and the set of positions where each pair occurs. This enables O(1) updates to pair counts when merges are applied, rather than O(n) rescans, and is a reusable component for any algorithm requiring frequent pair counting with position awareness.
The Word::merge_pair method applies a merge operation (replacing a token pair with a new token) to a sequence of token IDs and returns precise delta changes to pair counts. By tracking exactly which pairs are removed and created during the merge operation, it enables incremental updates to global pair statistics without rescanning the entire sequence. This technique is valuable for any string replacement algorithm that needs to maintain frequency statistics efficiently.