Input: "Hello  world!"
→ UTF-8: [72, 101, 108, 108, 111, 32, 240, 159, 140, 141, 32, 119, ...]
→ BPE on bytes → merges → tokens

import regex as re
from collections import defaultdict

class ByteBPE:
    def __init__(self):
        self.merges = {}           # (byte1, byte2) → new_id
        self.vocab = {}            # id → bytes
        self.pattern = r"""'s|'t|'re|'ve|'m|'ll|'d| ?[a-zA-Z]+| ?[0-9]+| ?[^\s\w]+|\s+"""

    def get_stats(self, ids):
        pairs = defaultdict(int)
        for i in range(len(ids)-1):
            pairs[(ids[i], ids[i+1])] += 1
        return pairs

    def merge(self, ids, pair, new_id):
        new_ids = []
        i = 0
        while i < len(ids):
            if i < len(ids)-1 and (ids[i], ids[i+1]) == pair:
                new_ids.append(new_id)
                i += 2
            else:
                new_ids.append(ids[i])
                i += 1
        return new_ids

    def train(self, text, vocab_size=50257, verbose=False):
        # 1. Pre-tokenize with regex
        utf8_bytes = text.encode("utf-8")
        ids = list(utf8_bytes)  # list of integers in range(0, 256)

        # 2. Build initial vocab: 256 bytes
        self.vocab = {i: bytes([i]) for i in range(256)}
        num_merges = vocab_size - 256

        for i in range(num_merges):
            stats = self.get_stats(ids)
            if not stats: break
            pair = max(stats, key=stats.get)
            new_id = 256 + i
            ids = self.merge(ids, pair, new_id)
            self.merges[pair] = new_id
            self.vocab[new_id] = self.vocab[pair[0]] + self.vocab[pair[1]]

            if verbose and i < 5:
                print(f"Merge {i+1}: {pair} → {new_id} | '{self.vocab[new_id].decode('utf-8', errors='replace')}'")

        print(f"Trained {len(self.merges)} merges. Final vocab size: {len(self.vocab)}")

    def encode_chunk(self, text):
        """Encode a single chunk using BPE merges"""
        ids = list(text.encode("utf-8"))
        while len(ids) >= 2:
            # Get all possible pairs and their ranks
            pair_to_rank = {pair: rank for rank, pair in enumerate(self.merges.keys()) if pair in self.merges}
            pairs = [(ids[i], ids[i+1]) for i in range(len(ids)-1)]
            if not any(p in pair_to_rank for p in pairs):
                break
            # Get lowest rank (earliest merge)
            min_rank = min(pair_to_rank.get(p, float('inf')) for p in pairs)
            pair = next(p for p in pairs if pair_to_rank.get(p) == min_rank)
            idx = pairs.index(pair)
            ids = ids[:idx] + [self.merges[pair]] + ids[idx+2:]
        return ids

    def encode(self, text):
        """Full encode with regex pre-tokenization"""
        if not text: return []
        chunks = re.findall(self.pattern, text)
        ids = []
        for chunk in chunks:
            ids.extend(self.encode_chunk(chunk))
        return ids

    def decode(self, ids):
        """Decode list of ids → text"""
        if not ids: return ""
        bytes_list = [self.vocab[id] for id in ids if id in self.vocab]
        return b''.join(bytes_list).decode("utf-8", errors="replace")

# Download TinyShakespeare
import urllib.request
url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
text = urllib.request.urlopen(url).read().decode("utf-8")

# Train
tokenizer = ByteBPE()
tokenizer.train(text[:100_000], vocab_size=1024, verbose=True)

# Test
text = "Hello  world! This is a test."
ids = tokenizer.encode(text)
decoded = tokenizer.decode(ids)

print(f"Text:  {text}")
print(f"IDs:   {ids}")
print(f"Decoded: {decoded}")
print(f"Match: {text == decoded}")

# Use GPT-2's official merges and vocab
import json, requests

merges_url = "https://huggingface.co/gpt2/resolve/main/merges.txt"
vocab_url = "https://huggingface.co/gpt2/resolve/main/vocab.json"

merges = requests.get(merges_url).text.strip().split('\n')[1:]  # skip header
vocab = json.loads(requests.get(vocab_url).text)

# Build merges dict
gpt2_merges = {}
for i, merge in enumerate(merges):
    p1, p2 = merge.split()
    gpt2_merges[(p1, p2)] = 50257 + i  # GPT-2 vocab starts at 50257?

print("Loaded GPT-2 merges and vocab")

class TrieNode:
    def __init__(self):
        self.children = {}
        self.token_id = None

class FastByteBPE(ByteBPE):
    def __init__(self):
        super().__init__()
        self.trie = TrieNode()

    def build_trie(self):
        root = self.trie
        for token_id, token_bytes in self.vocab.items():
            node = root
            for byte in token_bytes:
                if byte not in node.children:
                    node.children[byte] = TrieNode()
                node = node.children[byte]
            node.token_id = token_id
        return root

    def encode_chunk_trie(self, text):
        bytes_in = list(text.encode("utf-8"))
        ids = []
        i = 0
        while i < len(bytes_in):
            node = self.trie
            j = i
            best_id = None
            best_len = 0
            while j < len(bytes_in) and bytes_in[j] in node.children:
                node = node.children[bytes_in[j]]
                if node.token_id is not None:
                    best_id = node.token_id
                    best_len = j - i + 1
                j += 1
            if best_id is not None:
                ids.append(best_id)
                i += best_len
            else:
                ids.append(bytes_in[i])
                i += 1
        return ids

import time

# Train
tokenizer = ByteBPE()
tokenizer.train(text[:50_000], vocab_size=512)

# Trie version
fast_tokenizer = FastByteBPE()
fast_tokenizer.merges = tokenizer.merges
fast_tokenizer.vocab = tokenizer.vocab
fast_tokenizer.build_trie()

test_text = "Hello world! " * 100

# List-based
start = time.time()
for _ in range(1000):
    tokenizer.encode(test_text)
list_time = time.time() - start

# Trie-based
start = time.time()
for _ in range(1000):
    fast_tokenizer.encode_chunk_trie(test_text)
trie_time = time.time() - start

print(f"List: {list_time:.3f}s | Trie: {trie_time:.3f}s | Speedup: {list_time/trie_time:.1f}x")

class GPT2ByteBPE(ByteBPE):
    def __init__(self):
        super().__init__()
        self.pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"

    def encode(self, text):
        if not text: return []
        import regex as re
        chunks = re.findall(self.pattern, text)
        ids = []
        for chunk in chunks:
            ids.extend(self.encode_chunk(chunk))
        return ids

import re
from collections import defaultdict

class ByteBPE:
    def __init__(self):
        self.merges = {}
        self.vocab = {i: bytes([i]) for i in range(256)}
        self.pattern = r"""'s|'t|'re|'ve|'m|'ll|'d| ?[a-zA-Z]+| ?[0-9]+| ?[^\s\w]+|\s+"""

    def train(self, text, vocab_size=1024):
        ids = list(text.encode("utf-8"))
        for i in range(vocab_size - 256):
            pairs = defaultdict(int)
            for j in range(len(ids)-1):
                pairs[(ids[j], ids[j+1])] += 1
            if not pairs: break
            pair = max(pairs, key=pairs.get)
            new_id = 256 + i
            ids = [new_id if (a,b) == pair else a for a, b in zip(ids, ids[1:]) if not ((a,b) == pair)] + [new_id] * ids.count(pair[0])  # simplified
            self.merges[pair] = new_id
            self.vocab[new_id] = self.vocab[pair[0]] + self.vocab[pair[1]]

    def encode(self, text):
        chunks = re.findall(self.pattern, text)
        ids = []
        for chunk in chunks:
            bytes_in = list(chunk.encode("utf-8"))
            while len(bytes_in) > 1:
                pairs = [(bytes_in[i], bytes_in[i+1]) for i in range(len(bytes_in)-1)]
                merge_pair = min(pairs, key=lambda p: self.merges.get(p, float('inf')))
                if merge_pair not in self.merges: break
                idx = pairs.index(merge_pair)
                bytes_in = bytes_in[:idx] + [self.merges[merge_pair]] + bytes_in[idx+2:]
            ids.extend(bytes_in)
        return ids

    def decode(self, ids):
        return ''.join(self.vocab[i].decode('utf-8', errors='replace') for i in ids)