BPE Trainer

mediumFREETokenization AlgorithmsLLM 1: Foundations

Implement BPE training for a given number of merges.

Starting with a character-level tokenization of the input text (each character is a token, words separated by spaces are kept separate), repeatedly find the most frequent adjacent token pair and merge them.

Input:

text: the text
num_merges: number of merges to perform

Output: Return the final vocabulary as a set of the unique tokens after all merges. The tests sort it alphabetically and print it space-separated.

Rules:

Split text into words by spaces; each word is tokenized into characters
At each step, count adjacent pairs across all words, find the most frequent
If there's a tie, pick the pair that comes first alphabetically (by concatenation)
Merge that pair in all words
Repeat for the given number of merges

Example

Input

low lower newest
2

Output

e low n r s t w

Reasoning

Inputs: text is three words, and num_merges is $2$ . Each word starts as its own character sequence: l·o·w, l·o·w·e·r, n·e·w·e·s·t. Pairs never cross a space, so the last letter of one word is never paired with the first letter of the next.
Merge 1: counting adjacent pairs over all three words gives l·o $=2$ , o·w $=2$ , w·e $=2$ (once in "lower", once in "newest"), and every other pair $=1$ . Three pairs tie at $2$ . Comparing their concatenations, "lo" < "ow" < "we", so l and o fuse into lo: lo·w, lo·w·e·r, n·e·w·e·s·t.
Merge 2: now lo·w $=2$ and w·e $=2$ tie. "low" < "we", so low is formed: low, low·e·r, n·e·w·e·s·t.
Output: the distinct tokens still present are low, e, r, n, w, s, t. Sorted alphabetically they read e low n r s t w.
Notice that lo, l and o are missing. The vocabulary holds only the tokens that are still present after the last merge, not every token that ever existed.

Constraints

1 <= number of merges <= 20
Text contains lowercase letters and spaces
Words are separated by single spaces
Output sorted unique tokens after all merges

Learn this in

Tokenization & EmbeddingsLLM Study Plan Text PreprocessingNLP Study Plan Word EmbeddingsNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Multi-Head Attentionmedium

Concept primer

Background Knowledge

The BPE (Byte Pair Encoding) algorithm is a lossless text encoding technique that operates by merging the most frequent adjacent pairs of bytes in a string. In the context of tokenization, BPE is used to create a vocabulary of subwords, which are contiguous sequences of characters within a word. This approach helps to reduce the dimensionality of the input text and improve the efficiency of language models. The BPE algorithm starts with a character-level tokenization of the input text and iteratively merges adjacent tokens based on their frequency.

The key concept in BPE is the idea of tokenization, which refers to the process of breaking down text into individual units, such as characters, words, or subwords. In this problem, we start with a character-level tokenization and gradually merge adjacent tokens to form subwords. The frequency of adjacent token pairs is used to determine which pairs to merge at each step. This process is repeated for a specified number of merges, which controls the granularity of the resulting vocabulary.

The BPE algorithm has several advantages, including its ability to handle out-of-vocabulary (OOV) words and its efficiency in representing text data. By merging frequent adjacent pairs, BPE can create a vocabulary that captures common subword patterns in the input text, which is useful for language modeling and other natural language processing tasks.

Algorithm/Approach

The general approach to solving this problem involves implementing the BPE algorithm, which consists of the following pattern:

Initialize the vocabulary with a character-level tokenization of the input text
Repeat for the specified number of merges:
- Count the frequency of adjacent token pairs across all words
- Find the most frequent pair (or the pair that comes first alphabetically in case of a tie)
- Merge the selected pair in all words
Output the final vocabulary, sorted alphabetically

Step-by-Step Strategy

To implement the solution, follow these steps:

Tokenize the input text: Split the text into words and tokenize each word into characters.
Initialize the vocabulary: Create a set of unique tokens from the character-level tokenization.
Count adjacent pairs: Iterate through each word and count the frequency of adjacent token pairs.
Find the most frequent pair: Identify the pair with the highest frequency (or the pair that comes first alphabetically in case of a tie).
Merge the pair: Replace all occurrences of the selected pair with a new token that represents the merged pair.
Repeat steps 3-5: Repeat the process for the specified number of merges.
Output the final vocabulary: Sort the resulting vocabulary alphabetically and output it as a space-separated string.

Common Pitfalls

When implementing the solution, watch out for the following:

Incorrect tokenization: Ensure that the input text is tokenized correctly into characters and words.
Inconsistent merging: Make sure to merge the selected pair consistently across all words.
Tie-breaking: Handle ties correctly by choosing the pair that comes first alphabetically.

Time & Space Complexity

The expected time complexity of the solution is $O(n \cdot m \cdot k)$ , where $n$ is the length of the input text, $m$ is the number of merges, and $k$ is the average length of a word. The space complexity is $O(n + m)$ , where $n$ is the length of the input text and $m$ is the number of unique tokens in the final vocabulary. Note that these complexities assume a naive implementation and may be optimized further with more efficient data structures and algorithms.

solution.py

Python 3.13

Test results