Byte Pair Encoding › BPE Implementation from Scratch

class BPETokenizer:
    """
    A complete BPE tokenizer trained from raw text.
    Supports encode (text -> IDs) and decode (IDs -> text).
    """

    END_OF_WORD = "</w>"

    def __init__(self, vocab_size=500):
        self.vocab_size = vocab_size
        self.merge_rules = []  # ordered list of (a, b) pairs
        self.token_to_id = {}  # str -> int
        self.id_to_token = {}  # int -> str

    # ------------------------------------------------------------------ #
    #  Internal helpers                                                    #
    # ------------------------------------------------------------------ #

    def _build_word_freqs(self, texts):
        wf = {}
        for text in texts:
            for word in text.lower().split():
                key = " ".join(list(word)) + f" {self.END_OF_WORD}"
                wf[key] = wf.get(key, 0) + 1
        return wf

    @staticmethod
    def _count_pairs(wf):
        pairs = {}
        for word, freq in wf.items():
            tokens = word.split()
            for i in range(len(tokens) - 1):
                p = (tokens[i], tokens[i + 1])
                pairs[p] = pairs.get(p, 0) + freq
        return pairs

    @staticmethod
    def _merge_pair(wf, pair):
        merged = pair[0] + pair[1]
        return {
            word.replace(f"{pair[0]} {pair[1]}", merged): freq
            for word, freq in wf.items()
        }

    # ------------------------------------------------------------------ #
    #  Training                                                            #
    # ------------------------------------------------------------------ #

    def train(self, texts):
        wf = self._build_word_freqs(texts)

        # Collect base characters
        base_tokens = set()
        for word in wf:
            base_tokens.update(word.split())

        num_merges = self.vocab_size - len(base_tokens)

        for _ in range(max(0, num_merges)):
            pairs = self._count_pairs(wf)
            if not pairs:
                break
            best = max(pairs.items(), key=lambda x: x[1])[0]
            self.merge_rules.append(best)
            wf = self._merge_pair(wf, best)

        # Build vocabulary from final token set
        all_tokens = set(base_tokens)
        for a, b in self.merge_rules:
            all_tokens.add(a + b)

        for idx, tok in enumerate(sorted(all_tokens)):
            self.token_to_id[tok] = idx
            self.id_to_token[idx] = tok

        return self

    # ------------------------------------------------------------------ #
    #  Encoding                                                            #
    # ------------------------------------------------------------------ #

    def _encode_word(self, word):
        tokens = list(word)
        for pair in self.merge_rules:
            i = 0
            while i < len(tokens) - 1:
                if tokens[i] == pair[0] and tokens[i + 1] == pair[1]:
                    tokens = tokens[:i] + [pair[0] + pair[1]] + tokens[i + 2 :]
                else:
                    i += 1
        return tokens

    def tokenize(self, text):
        """Return list of subword token strings."""
        result = []
        for word in text.lower().split():
            result.extend(self._encode_word(word))
        return result

    def encode(self, text):
        """Return list of integer token IDs."""
        unk_id = self.token_to_id.get("<unk>", 0)
        return [self.token_to_id.get(t, unk_id) for t in self.tokenize(text)]

    # ------------------------------------------------------------------ #
    #  Decoding                                                            #
    # ------------------------------------------------------------------ #

    def decode(self, ids):
        """Reconstruct text from integer IDs."""
        tokens = [self.id_to_token.get(i, "<unk>") for i in ids]
        return "".join(tokens).replace(self.END_OF_WORD, " ").strip()

training_corpus = [
    "the cat sat on the mat",
    "the dog ran in the park",
    "cats and dogs are pets",
    "the cat chased the dog",
    "dogs like to run and play",
    "cats prefer to sleep all day",
    "the park has many dogs and cats",
    "running dogs and sleeping cats",
    "the fastest dog won the race",
    "slower cats prefer the mat",
    "cats run and dogs play in parks",
    "the dog chased a cat running fast",
]

bpe = BPETokenizer(vocab_size=120)
bpe.train(training_corpus)

test_inputs = [
    "the cat ran",
    "dogs chase cats",
    "running and sleeping",
    "unrecognized word",  # novel word: not in training
]