Implement a Byte-Pair Encoding Tokenizer: Threshold Training, Encode and Decode

Implement a byte-pair encoding (BPE) tokenizer from scratch with three functions:

• train(text: str, threshold: int) learns merge rules from text.
• encode(text: str) -> list[int] turns text into a list of integer token ids using the learned merges. The interviewer called this list the "embedding"; it is the sequence of token ids.
• decode(ids: list[int]) -> str turns a list of token ids back into text.

Training starts from single characters. In each round, count the adjacent pairs of tokens, pick the pair that occurs most often, and merge it into one new token. A pair may be merged only if its count is at least threshold, and training stops when no pair reaches the threshold. If several pairs share the highest count, choose by lexicographical order.

The interviewer's example: for text = "abababa" and threshold = 2, the first round chooses the pair (a, b) and the second round chooses the pair (ab, ab).

Clarifying Questions

• How are overlapping occurrences counted? For example, in "aaa", does the pair (a, a) occur once or twice?
• When pairs tie on count, is the lexicographical comparison made on the pair as (first token, second token), or on the concatenated string?
• When occurrences of the chosen pair overlap, which of them are merged?
• How are ids assigned to the base characters and to the merged tokens?
• What should encode do with a character that never appeared in the training text, and what should decode do with an unknown id?

Part 1 — Train

Implement train. Check your implementation against the example above, round by round, and state its time complexity.

What This Part Should Cover

• Pair counting and a merge step that reproduce the example exactly
• The threshold and the lexicographic tie-break, applied in the right order
• The cost of straightforward training, and how it could be reduced for long texts

Part 2 — Encode and decode

Implement encode and decode on top of the learned merges. decode(encode(s)) must return s for any string made of characters seen in training.

What This Part Should Cover

• A vocabulary that maps every base character and merged token to an id
• Applying the learned merges in a consistent order
• Behaviour for unseen characters and invalid ids