Performance ideas

I made a toy GPT2 tokenizer as a python rust extension. It seems to be slightly faster than tiktoken in my tests. It looks like https://github.com/openai/tiktoken/pull/31 may get most or all the way there, but I thought I'd post the results from this script:

import os
import time
from typing import Any, cast

import numpy as np
import tiktoken

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

def benchmark_batch(documents: list[bytes]) -> None:
    num_threads = 1
    num_bytes = sum(map(len, documents))
    print(f"num_threads: {num_threads}, num_bytes: {num_bytes}")
    documents_decoded = [doc.decode("utf8") for doc in documents]

    enc = tiktoken.get_encoding("gpt2")
    enc.encode("warmup")

    start = time.perf_counter_ns()
    tiktoken_output = enc.encode_ordinary_batch(documents_decoded, num_threads=num_threads)
    end = time.perf_counter_ns()
    print(f"tiktoken \t{num_bytes / (end - start) * 1e9} bytes / s")

    import transformers

    hf_enc = cast(Any, transformers).GPT2TokenizerFast.from_pretrained("gpt2")
    hf_enc.model_max_length = 1e30  # silence!
    hf_enc.encode("warmup")

    start = time.perf_counter_ns()
    hf_enc_output = hf_enc(documents_decoded)
    end = time.perf_counter_ns()
    print(f"huggingface \t{num_bytes / (end - start) * 1e9} bytes / s")

    import csh_bpe.codec
    csh_bpe_enc = csh_bpe.codec.RustGPTCodec(word_encoder_kind="bigram", doc_splitter_kind="direct")
    csh_bpe_enc.encode(np.frombuffer(b"warmup", dtype=np.uint8))
    
    start = time.perf_counter_ns()
    csh_bpe_output = csh_bpe_enc.encode(np.frombuffer(documents[0], dtype=np.uint8))
    end = time.perf_counter_ns()
    print(f"csh_bpe \t{num_bytes / (end - start) * 1e9} bytes / s")

    assert hf_enc_output["input_ids"][0] == tiktoken_output[0]
    assert csh_bpe_output.tolist() == tiktoken_output[0]


def main():
    with open(os.path.join(SCRIPT_DIR, "..", "local-data", "64MB.txt"), "rb") as f:
        contents = f.read()
    benchmark_batch([contents])


if __name__ == "__main__":
    main()

The text is 64MiB of wikipedia wikitext, probably enwik8, but I just found it on my hard drive.

python -m csh_bpe.compare_tiktoken
num_threads: 1, num_bytes: 67108864
tiktoken        6004366.360373783 bytes / s
huggingface     1120214.7857500792 bytes / s
csh_bpe         17070974.6114367 bytes / s

There are no fancy optimizations here (like SIMD stuff), the library has a few things it might do differently from tiktoken:

1. The word splitting regular expression is implemented using rust code instead of a regexp library. It uses Go's unicode tables: https://github.com/golang/go/blob/19309779ac5e2f5a2fd3cbb34421dafb2855ac21/src/unicode/tables.go and this seems to produce the same output at least for this 64MB file. The splitting is done with a function that takes a u8 numpy array and start offset and returns the end offset.
2. The bigram encoder takes a u8 slice for the word, a HashMap<(i32, i32), i32> mergelist, an i32 slice mapping bytes to tokens (used to populate the initial output), and a mutable i32 slice of output tokens. It keeps a list of skip lengths for each index of the output tokens (initially all 1s), which it updates whenever it merges two tokens together, then compacts the output tokens when it is done.
3. (I think tiktoken does this) after splitting, before encoding a word, it will check the vocab hashmap to see if the word is already a single token.
4. The interface uses numpy arrays instead of bytes, and the output array is provided as one of the inputs so the caller can manage more memory allocations (not sure if this has any performance impact)

I didn't implement rust regexps so I don't know if the word splitting matters, though I could benchmark just the splitting part.

Start by reviewing the benchmark in the issue, especially tiktoken's encode_ordinary_batch comparison and the csh_bpe.compare_tiktoken script. The issue names no tiktoken file or test and provides no concrete optimization target or completion criterion, so the intended scope would need to be clarified before work begins.