wordchipper - HPC Rust BPE Tokenizer

Overview

wordchipper is a high-performance Rust byte-pair encoder tokenizer for the OpenAI GPT-2 tokenizer family. It achieves throughput speedups relative to tiktoken-rs in rust on a 64 core machine of ~4.3-5.7x (4 to 64 cores) for general regex BPE vocabularies, and ~6.9x-9.2x when using custom DFA lexers for specific OpenAI vocabularies. Under python wrappers, we see a range of ~2x-4x (4 to 64 cores) speedups over tiktoken.

Suite Crates

This is the main crate for the wordchipper project.

The core additional user-facing crates are:

• wordchipper-cli - a multi-tool tokenizer binary; notably:
  • wordchipper-cli cat - in-line encoder/decoder tool.
  • wordchipper-cli train - tokenizer training tool.
• wordchipper-training - an extension crate for training tokenizers.

Encode/Decode Side-by-Side Benchmarks

Read the full performance paper:

• wordchipper: Fast BPE Tokenization with Substitutable Internals

Client Usage

Pretrained Vocabularies

• OpenAI OATokenizer

Encoders and Decoders

• Token Encoders
• Token Decoders

Loading Pretrained Models

Loading a pre-trained model requires reading the vocabulary, as well as configuring the spanning (regex and special words) configuration.

For a number of pretrained models, simplified constructors are available to download, cache, and load the vocabulary.

See: wordchipper::get_model

use std::sync::Arc;

use wordchipper::{
    get_model,
    TokenDecoder,
    TokenEncoder,
    UnifiedTokenVocab,
    disk_cache::WordchipperDiskCache,
};

fn example() -> wordchipper::errors::Result<(Arc<dyn TokenEncoder<u32>>, Arc<dyn TokenDecoder<u32>>)> {
    let mut disk_cache = WordchipperDiskCache::default();
    let vocab: UnifiedTokenVocab<u32> = get_model("openai/o200k_harmony", &mut disk_cache)?;

    let encoder = vocab.to_default_encoder();
    let decoder = vocab.to_default_decoder();

    Ok((encoder, decoder))
}

Acknowledgements

• Thank you to @karpathy and nanochat for the work on rustbpe.
• Thank you to tiktoken for their initial work in the rust tokenizer space.