bpe.py

"""Byte pair encoding, trained from scratch — the reference implementation.

This is the teaching version of the algorithm that produced the vocabularies on
/tokens/. It mirrors trainer.js exactly, symbol for symbol, and checkbpe.py
fails the build if the two ever diverge.

Two honest simplifications versus a production tokenizer:

  * It merges characters, not UTF-8 bytes. Real tokenizers start from the 256
    possible bytes so that any input at all is representable. Starting from
    characters is easier to watch and makes no difference to the mechanism.
  * Pre-tokenization is one simple rule (optional whitespace, then a run of
    non-whitespace) rather than a long regex. That rule is why leading spaces
    end up welded to words — which is the point worth seeing.

Tie-breaking: highest pair count wins; ties go to whichever pair was seen first
in a left-to-right scan. First-seen order avoids comparing strings, which
Python and JavaScript order differently for astral characters.
"""

import re

PRETOKEN = re.compile(r"\s*\S+|\s+")


def pretokenize(text):
    """Split into chunks of optional leading whitespace plus a run of non-space."""
    return PRETOKEN.findall(text)


def word_counts(text):
    """Unique chunks and how often each occurs, in first-seen order."""
    counts = {}
    for chunk in pretokenize(text):
        counts[chunk] = counts.get(chunk, 0) + 1
    return counts


def _pair_stats(words):
    """pair -> total count, in first-seen order (dicts preserve insertion order,
    and so do JavaScript Maps — that is what makes tie-breaking identical)."""
    counts = {}
    for symbols, freq in words:
        for i in range(len(symbols) - 1):
            pair = (symbols[i], symbols[i + 1])
            counts[pair] = counts.get(pair, 0) + freq
    return counts


def _merge_in(symbols, a, b):
    """Replace every adjacent a,b with the joined symbol."""
    out = []
    i = 0
    n = len(symbols)
    while i < n:
        if i < n - 1 and symbols[i] == a and symbols[i + 1] == b:
            out.append(a + b)
            i += 2
        else:
            out.append(symbols[i])
            i += 1
    return out


def train(text, num_merges):
    """Run BPE. Returns the merge list and a step-by-step record."""
    words = [(list(chunk), freq) for chunk, freq in word_counts(text).items()]

    alphabet = {}
    for chunk in word_counts(text):
        for ch in chunk:
            alphabet[ch] = True

    merges, steps = [], []
    for _ in range(num_merges):
        counts = _pair_stats(words)
        if not counts:
            break
        # Strict >, scanning in first-seen order: the earliest of any tied pairs
        # wins. No string comparison, so Python and JavaScript cannot disagree.
        best = None
        for pair, count in counts.items():
            if best is None or count > counts[best]:
                best = pair
        if counts[best] < 2:
            break                      # nothing repeats; merging is pointless

        a, b = best
        words = [(_merge_in(sym, a, b), freq) for sym, freq in words]
        merges.append((a, b))
        steps.append({
            "pair": [a, b],
            "token": a + b,
            "count": counts[best],
            "vocab": len(alphabet) + len(merges),
        })

    return {
        "merges": merges,
        "steps": steps,
        "alphabet": sorted(alphabet),
        "vocab_size": len(alphabet) + len(merges),
    }


def encode(word, merges):
    """Apply a trained merge list to one chunk, in the order learned."""
    symbols = list(word)
    for a, b in merges:
        symbols = _merge_in(symbols, a, b)
    return symbols


def encode_text(text, merges):
    out = []
    for chunk in pretokenize(text):
        out.extend(encode(chunk, merges))
    return out