📘

BPE Trainer

MediumTokenization Algorithms 📚 LLM 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:

Line 1: The text
Line 2: Number of merges to perform

Output: The final vocabulary (unique tokens after all merges), sorted alphabetically, 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 er l lo n ow s t w

Reasoning:

The input text "low lower newest" is first tokenized into characters, resulting in the tokens: l-o-w, l-o-w-e-r, n-e-w-e-s-t
The most frequent adjacent token pair is found to be "o-w", which appears in both "low" and "lower", so it is merged into a single token "ow", resulting in the tokens: l-ow, l-ow-er, n-e-w-e-s-t
After the first merge, the most frequent adjacent token pair is found to be "l-ow" is not more frequent than "n-e" or "e-w" or "w-e" but "l-ow" is not the most frequent, "e-w" is, however "lo" and "ow" and "ew" have the same frequency and "lo" comes first alphabetically (by concatenation), so "lo" is not merged, "ow" is already merged and "ew" is merged into a single token "ew" is not, "lo" is, resulting in the tokens: l-o-w, l-o-w-er, n-e-w-e-s-t, then "lo" is merged, resulting in the tokens: lo-w, lo-w-er, n-e-w-e-s-t
The final vocabulary after the two merges is: e, er, l, lo, n, ow, s, t, w, which when sorted alphabetically and space-separated gives the output: e er l lo n ow s t w

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

More from LLM 1: Foundations

Transformer Block Forward PassHard Feed-Forward NetworkEasy Layer NormalizationEasy KV-Cache AttentionHard Multi-Head AttentionMedium

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

0/0

Run code to see test results.