Byte-pair encoding tokenizer: learn merges, encode and decode

Implement a byte-pair encoding (BPE) tokenizer from scratch: learn merge rules from a training text, then use them to turn text into token ids and back. Training starts from single characters. Each round, count the adjacent token pairs, take the most frequent one and merge it into a new token. A pair is merged only if its count is at least threshold; training stops when no pair reaches it. Ties on count go to the lexicographically smallest pair.

Interface

Unless you agree on a different interface with the interviewer, assume:

• train(text: str, threshold: int) -> None — learns the ordered list of merges.
• encode(text: str) -> list[int] — token ids for text, using the learned merges.
• decode(ids: list[int]) -> str — the inverse of encode.
• Pairs are counted at every adjacent position, overlapping ones included ("aaa" has the pair (a, a) twice). A merge replaces occurrences left to right without overlap.
• Ties compare pairs as (first token, second token).
• Ids: the distinct training characters get 0, 1, 2, ... in sorted order, and each merged token gets the next id in the order it was learned.

Part 1: train

train("abababa", 2)
    round 1: pairs (a,b) = 3, (b,a) = 3  -> tie, merge (a, b)    tokens: ab ab ab a
    round 2: (ab,ab) = 2, (ab,a) = 1     -> merge (ab, ab)        tokens: abab ab a
    round 3: (abab,ab) = 1, (ab,a) = 1   -> below threshold, stop
    vocabulary: a = 0, b = 1, ab = 2, abab = 3

Round 2 reaches the threshold only because overlapping positions count; check your counting rule against it.

Part 2: encode and decode

Apply the merges to new text in the order they were learned.

encode("ababa")     ->  [3, 0]      # abab, a
encode("bab")       ->  [1, 2]      # b, ab
encode("abababab")  ->  [3, 3]
decode([3, 0])      ->  "ababa"

decode(encode(s)) == s for any s made of characters seen in training.

Clarifying questions

• Are overlapping occurrences counted when choosing a pair? Which occurrences are merged when they overlap?
• Is the tie-break on the pair or on the concatenated string?
• How are ids assigned to base characters and merged tokens?
• What does encode do with a character not seen in training, and decode with an unknown id?

What a strong answer covers

• Separate pieces: pair counting, applying one merge to a token list, the training loop, the vocabulary.
• The rules (overlap counting, tie-break, id order) stated up front and checked round by round against "abababa".
• Encode applies merges in learned order (or by merge rank), so it segments like training did; a stated policy for unseen characters (error, or an unknown token).
• Complexity: naive training is O(rounds * n); faster training keeps pair counts and pair positions and updates only the neighbours of each merge, with a heap of counts.
• Tests from small hand-checked strings such as "abababa" and "aaa", plus round-trip tests.

Follow-ups

• Make training fast enough for tens of millions of characters.
• Use bytes as base tokens so no input is ever unknown.
• Split the text into words first so merges never cross a space.
• Stop at a target vocabulary size instead of a threshold.