Homework · Runs in your browser

Build a tiny byte-pair-encoding tokenizer

Byte-pair encoding starts from single characters and repeatedly merges the most frequent adjacent pair. Implement the three steps in starter.py: count pairs, merge one pair, and train for a number of merges. Then run the tests. In the browser runner, press Run tests. In Colab or on your own machine, put starter.py and test_starter.py in one folder and run pytest. The same files work everywhere, and no libraries are needed.

Hints:

Work on a list of strings. “abab” starts as [“a”, “b”, “a”, “b”].
Count pairs of neighbours with zip(tokens, tokens[1:]).
When two pairs tie, pick the one that appears first in the text. This keeps results the same on every run.
To merge, walk the list left to right. If the current and next token match the pair, output their join and skip both.
This is a toy. Real tokenizers work on bytes and use a huge training corpus, but the loop is the same.

Running the tests needs JavaScript. You can still read the starter code below and run it on your own machine with pytest.

Your code starter.py

def most_frequent_pair(tokens):
    """Return the most frequent adjacent pair as a tuple, or None if there is no pair.

    On a tie, return the pair that appears first in the list.
    """
    # Wrong on purpose: always claims there is nothing to merge.
    return None


def merge_pair(tokens, pair):
    """Return a new list where every adjacent occurrence of pair is joined into one token.

    Scan left to right; merged tokens do not overlap.
    """
    raise NotImplementedError


def train_bpe(text, num_merges):
    """Start from single characters and apply up to num_merges merges.

    Return (tokens, merges): the final token list and the list of merged pairs
    in the order they were learned. Stop early if no pair is left.
    """
    raise NotImplementedError

Tab inserts 4 spaces. To move on with the keyboard, press Esc, then Tab. Shift+Tab always moves back.

Show a solution

Try it yourself first. This site is static, so the solution is sent to your browser with the page; it is hidden here, not secret.

def most_frequent_pair(tokens):
    """Return the most frequent adjacent pair as a tuple, or None if there is no pair."""
    counts = {}
    for pair in zip(tokens, tokens[1:]):
        counts[pair] = counts.get(pair, 0) + 1
    if not counts:
        return None
    best = None
    for pair, n in counts.items():  # dicts keep first-seen order, so ties go to the first
        if best is None or n > counts[best]:
            best = pair
    return best


def merge_pair(tokens, pair):
    """Return a new list where every adjacent occurrence of pair is joined into one token."""
    out = []
    i = 0
    while i < len(tokens):
        if i + 1 < len(tokens) and (tokens[i], tokens[i + 1]) == pair:
            out.append(tokens[i] + tokens[i + 1])
            i += 2
        else:
            out.append(tokens[i])
            i += 1
    return out


def train_bpe(text, num_merges):
    """Start from single characters and apply up to num_merges merges."""
    tokens = list(text)
    merges = []
    for _ in range(num_merges):
        pair = most_frequent_pair(tokens)
        if pair is None:
            break
        tokens = merge_pair(tokens, pair)
        merges.append(pair)
    return tokens, merges

Show the tests (test_starter.py)

Python

from starter import merge_pair, most_frequent_pair, train_bpe


def test_most_frequent_pair_counts_neighbours():
    assert most_frequent_pair(list("abab")) == ("a", "b"), (
        "('a','b') appears twice, ('b','a') once"
    )


def test_tie_goes_to_the_first_pair():
    assert most_frequent_pair(list("abcd")) == ("a", "b"), (
        "all pairs appear once; the first one in the text wins"
    )


def test_no_pair_in_short_input():
    assert most_frequent_pair(["a"]) is None, "one token has no neighbours"


def test_merge_pair_joins_without_overlap():
    assert merge_pair(list("aaab"), ("a", "a")) == ["aa", "a", "b"], (
        "scan left to right; the third 'a' is left alone"
    )


def test_train_bpe_learns_merges_in_order():
    tokens, merges = train_bpe("low lower lowest", 2)
    assert merges == [("l", "o"), ("lo", "w")], "most frequent pair first, then its result"
    assert tokens[:3] == ["low", " ", "low"], "frequent words become single tokens"


def test_train_bpe_stops_when_nothing_left_to_merge():
    tokens, merges = train_bpe("ab", 5)
    assert tokens == ["ab"] and merges == [("a", "b")], (
        "stop early instead of looping when only one token remains"
    )