Byte Pair Encoding from Scratch

from collections import Counter, deque
from functools import lru_cache
import json
 
 
class BPETokenizerSimple:
    def __init__(self):
        # Maps token_id to token_str (e.g., {11246: "some"})
        self.vocab = {}
        # Maps token_str to token_id (e.g., {"some": 11246})
        self.inverse_vocab = {}
        # Dictionary of BPE merges: {(token_id1, token_id2): merged_token_id}
        self.bpe_merges = {}
 
    def train(self, text, vocab_size, allowed_special={"<|endoftext|>"}):
        """Train the BPE tokenizer from scratch."""
        # Pre-tokenize: split into words, marking leading spaces with 'Ġ'
        tokens = self.pretokenize_text(text)
 
        # Initialize vocab with the 256 byte values + extra chars + 'Ġ'
        unique_chars = [chr(i) for i in range(256)]
        unique_chars.extend(
            char for char in sorted({char for token in tokens for char in token})
            if char not in unique_chars
        )
        if "Ġ" not in unique_chars:
            unique_chars.append("Ġ")
 
        self.vocab = {i: char for i, char in enumerate(unique_chars)}
        self.inverse_vocab = {char: i for i, char in self.vocab.items()}
 
        # Add allowed special tokens (e.g. <|endoftext|>)
        if allowed_special:
            for token in allowed_special:
                if token not in self.inverse_vocab:
                    new_id = len(self.vocab)
                    self.vocab[new_id] = token
                    self.inverse_vocab[token] = new_id
 
        # Each pre-token becomes a list of character IDs
        token_id_sequences = [
            [self.inverse_vocab[char] for char in token]
            for token in tokens
        ]
 
        # BPE steps 1-3: repeatedly find and replace frequent pairs
        for new_id in range(len(self.vocab), vocab_size):
            pair_id = self.find_freq_pair(token_id_sequences, mode="most")
            if pair_id is None:
                break
            token_id_sequences = self.replace_pair(
                token_id_sequences, pair_id, new_id)
            self.bpe_merges[pair_id] = new_id
 
        # Build the vocabulary entries for the merged tokens
        for (p0, p1), new_id in self.bpe_merges.items():
            merged_token = self.vocab[p0] + self.vocab[p1]
            self.vocab[new_id] = merged_token
            self.inverse_vocab[merged_token] = new_id
 
    def encode(self, text, allowed_special=None):
        """Encode text into a list of token IDs."""
        token_ids = []
        # (special-token passthrough handling omitted here for brevity —
        #  see the original notebook for the tiktoken-style guard rails)
        tokens = self.pretokenize_text(text)
 
        for tok in tokens:
            if tok in self.inverse_vocab:
                token_ids.append(self.inverse_vocab[tok])
            else:
                token_ids.extend(self.tokenize_with_bpe(tok))
        return token_ids
 
    def tokenize_with_bpe(self, token):
        """Tokenize a single word using the learned BPE merges."""
        # Start from individual characters
        token_ids = [self.inverse_vocab.get(char, None) for char in token]
        if None in token_ids:
            missing_chars = [char for char, tid in zip(token, token_ids)
                             if tid is None]
            raise ValueError(f"Characters not found in vocab: {missing_chars}")
 
        # Repeatedly merge adjacent pairs that exist in bpe_merges
        can_merge = True
        while can_merge and len(token_ids) > 1:
            can_merge = False
            new_tokens = []
            i = 0
            while i < len(token_ids) - 1:
                pair = (token_ids[i], token_ids[i + 1])
                if pair in self.bpe_merges:
                    merged_token_id = self.bpe_merges[pair]
                    new_tokens.append(merged_token_id)
                    i += 2  # skip the next token; it was merged
                    can_merge = True
                else:
                    new_tokens.append(token_ids[i])
                    i += 1
            if i < len(token_ids):
                new_tokens.append(token_ids[i])
            token_ids = new_tokens
        return token_ids
 
    def decode(self, token_ids):
        """Decode a list of token IDs back into a string."""
        out = []
        for tid in token_ids:
            if tid not in self.vocab:
                raise ValueError(f"Token ID {tid} not found in vocab.")
            tok = self.vocab[tid]
            if tid == 198 or tok == "\n":
                out.append("\n")
            elif tok.startswith("Ġ"):
                out.append(" " + tok[1:])
            else:
                out.append(tok)
        return "".join(out)
 
    @staticmethod
    def find_freq_pair(token_id_sequences, mode="most"):
        pairs = Counter(
            pair
            for token_ids in token_id_sequences
            for pair in zip(token_ids, token_ids[1:])
        )
        if not pairs:
            return None
        if mode == "most":
            return max(pairs.items(), key=lambda x: x[1])[0]
        elif mode == "least":
            return min(pairs.items(), key=lambda x: x[1])[0]
        else:
            raise ValueError("Invalid mode. Choose 'most' or 'least'.")
 
    @staticmethod
    def replace_pair(token_id_sequences, pair_id, new_id):
        replaced_sequences = []
        for token_ids in token_id_sequences:
            dq = deque(token_ids)
            replaced = []
            while dq:
                current = dq.popleft()
                if dq and (current, dq[0]) == pair_id:
                    replaced.append(new_id)
                    dq.popleft()  # remove the 2nd token of the pair
                else:
                    replaced.append(current)
            replaced_sequences.append(replaced)
        return replaced_sequences