Build a BPE Tokenizer From Scratch in Python (and See Why Some Languages Cost More)

"""A byte-level BPE tokenizer, the kind behind GPT and Claude, written from scratch for Revontuli's chat."""
import collections
import re

# Text is first cut into chunks: words with their leading space, numbers, punctuation runs and whitespace. Merges
# never cross a chunk boundary, so "dog." and "dog!" share the token for " dog".
PATTERN = re.compile(r"'(?:s|t|re|ve|m|ll|d)| ?[^\W\d_]+| ?\d{1,3}| ?[^\s\w]+|\s+(?!\S)|\s+")


def chunks(text):
    return PATTERN.findall(text)


# ---------- Step 1: text is bytes ----------
def to_bytes(text):
    """The UTF-8 bytes of text, as a list of ints 0-255."""
    # TODO (Step 1): text.encode("utf-8") gives bytes; list() of it gives the ints.
    raise NotImplementedError("Step 1: write to_bytes()")


def pair_counts(ids, counts=None, weight=1):
    """Count every adjacent pair in ids (overlapping: [1, 1, 1] has the pair (1, 1) twice), adding `weight` per
    occurrence into `counts` (a new collections.Counter when None). Returns the counter."""
    # TODO (Step 1): zip(ids, ids[1:]) walks the adjacent pairs.
    raise NotImplementedError("Step 1: write pair_counts()")


# ---------- Step 2: one merge ----------
def merge(ids, pair, new_id):
    """A new list where every occurrence of `pair`, scanning left to right without overlaps, is replaced by
    new_id: merge([7, 7, 7], (7, 7), 256) == [256, 7]."""
    raise NotImplementedError("merge() arrives in Step 2")


# ---------- Step 3: training ----------
def train(text, vocab_size):
    """Learn vocab_size - 256 merges from text. Returns {pair: new_id} in the order learned (ids 256, 257, ...).
    Work on distinct chunks with their counts: each round, count pairs over all chunks (weighted by how often the
    chunk occurs), take the most frequent pair (ties: the smallest pair, e.g. (32, 116) before (32, 119)), give it
    the next id and merge it in every chunk. Stop early if no pair is left."""
    raise NotImplementedError("train() arrives in Step 3")


# ---------- Step 4: encode and decode ----------
def encode(text, merges):
    """Token ids for text. Per chunk: start from its bytes, then repeatedly apply the learned merge with the
    lowest id among the pairs present (the order training learned them), until none applies."""
    raise NotImplementedError("encode() arrives in Step 4")


def vocab(merges):
    """{id: bytes}: 0-255 are single bytes; every merged id is the bytes of its two parts joined."""
    raise NotImplementedError("vocab() arrives in Step 4")


def decode(ids, merges):
    """The text for token ids. A slice of ids can end in the middle of a character, so undecodable bytes become
    U+FFFD (errors="replace") instead of raising."""
    raise NotImplementedError("decode() arrives in Step 4")


# ---------- Step 5: what a token costs ----------
def tokens_per_word(encode_fn, text):
    """len(encode_fn(text)) divided by the number of words (text.split())."""
    raise NotImplementedError("tokens_per_word() arrives in Step 5")


def compare(encoders, texts):
    """{encoder name: {text name: tokens_per_word}} for every encoder and text."""
    raise NotImplementedError("compare() arrives in Step 5")