Implement a Byte-Pair Encoding Tokenizer: Threshold Training, Encode and Decode
Company: Glean
Role: Software Engineer
Category: Software Engineering Fundamentals
Difficulty: medium
Interview Round: Onsite
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.
```hint Read the example closely
Write out the token sequence after the first merge and count its pairs. Check which counting rule lets the second round reach the threshold.
```
#### 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.
```hint Order matters
Encoding should segment a text the way training would have. Think about the order in which the learned merges must be applied.
```
#### 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
### What a Strong Answer Covers
- Every ambiguity in the rules (overlap counting, tie-break, id assignment) settled explicitly and checked against the example
- Clean separation of pair counting, single-merge application and the training loop
- Encode and decode that round-trip, with a stated policy for unknown input
- Correct complexity for training and encoding, plus a credible plan to make training fast
- Tests built from small hand-checked strings such as `"abababa"` and `"aaa"`
### Follow-up Questions
- How would you make training fast enough for a text of tens of millions of characters?
- How would byte-level base tokens change the handling of characters never seen in training?
- If the text were first split into words, so that merges never cross a space, how would the training loop change?
Overview: Implement a byte-pair encoding tokenizer with train, encode and decode, where training repeatedly merges the most frequent adjacent token pair whose count reaches a threshold, breaking ties lexicographically. Tests precise reading of the merge rules, round-trip correctness and the complexity of training on long text.