| """ |
| Byte-level BPE from scratch (lecture: Tokenization). |
| |
| Core idea: |
| 1. Start with 256 tokens = every byte. There is no "unknown token"; anything can be encoded. |
| 2. Pre-tokenize text with a regex (GPT-2 style) so merges never cross spaces/punctuation. |
| 3. Repeat: find the most frequent adjacent token pair -> merge it into a new token. |
| 4. Stop when the vocab reaches its target size. |
| |
| Training is incremental (heap + pair->word index) so thousands of merges finish in seconds |
| instead of recounting every pair on every iteration. |
| """ |
| from __future__ import annotations |
|
|
| import heapq |
| import json |
| from collections import Counter, defaultdict |
|
|
| import regex as re |
|
|
| |
| GPT2_SPLIT = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""") |
|
|
| EOS = "<|eos|>" |
| USER = "<|user|>" |
| ASSISTANT = "<|assistant|>" |
| SPECIAL_TOKENS = [EOS, USER, ASSISTANT] |
|
|
|
|
| class BPETokenizer: |
| def __init__(self, merges: list[tuple[int, int]] | None = None, special_tokens: list[str] | None = None): |
| self.merges: list[tuple[int, int]] = merges or [] |
| self.special_tokens = special_tokens or list(SPECIAL_TOKENS) |
| self._build() |
|
|
| |
| def _build(self) -> None: |
| self.ranks = {pair: i for i, pair in enumerate(self.merges)} |
| self.vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)} |
| for i, (a, b) in enumerate(self.merges): |
| self.vocab[256 + i] = self.vocab[a] + self.vocab[b] |
| base = 256 + len(self.merges) |
| self.special_to_id = {tok: base + i for i, tok in enumerate(self.special_tokens)} |
| self.id_to_special = {v: k for k, v in self.special_to_id.items()} |
| self._special_re = re.compile("(" + "|".join(re.escape(t) for t in self.special_tokens) + ")") |
| self._cache: dict[str, list[int]] = {} |
|
|
| @property |
| def vocab_size(self) -> int: |
| return 256 + len(self.merges) + len(self.special_tokens) |
|
|
| @property |
| def eos_id(self) -> int: |
| return self.special_to_id[EOS] |
|
|
| |
| @classmethod |
| def train(cls, text: str, vocab_size: int, verbose: bool = True) -> "BPETokenizer": |
| n_merges = vocab_size - 256 - len(SPECIAL_TOKENS) |
| assert n_merges > 0, "vocab_size is too small" |
|
|
| word_freq = Counter(GPT2_SPLIT.findall(text)) |
| words = [list(w.encode("utf-8")) for w in word_freq] |
| freqs = list(word_freq.values()) |
|
|
| pair_counts: dict[tuple[int, int], int] = defaultdict(int) |
| where: dict[tuple[int, int], set[int]] = defaultdict(set) |
| for idx, w in enumerate(words): |
| for p in zip(w, w[1:]): |
| pair_counts[p] += freqs[idx] |
| where[p].add(idx) |
|
|
| |
| heap = [(-c, p) for p, c in pair_counts.items()] |
| heapq.heapify(heap) |
|
|
| merges: list[tuple[int, int]] = [] |
| while len(merges) < n_merges and heap: |
| neg_c, pair = heapq.heappop(heap) |
| if pair_counts.get(pair, 0) != -neg_c or -neg_c <= 0: |
| continue |
| new_id = 256 + len(merges) |
| merges.append(pair) |
|
|
| for idx in list(where[pair]): |
| w, f = words[idx], freqs[idx] |
| |
| for p in zip(w, w[1:]): |
| pair_counts[p] -= f |
| where[p].discard(idx) |
| |
| out, i = [], 0 |
| while i < len(w): |
| if i < len(w) - 1 and w[i] == pair[0] and w[i + 1] == pair[1]: |
| out.append(new_id) |
| i += 2 |
| else: |
| out.append(w[i]) |
| i += 1 |
| words[idx] = out |
| |
| touched = set() |
| for p in zip(out, out[1:]): |
| pair_counts[p] += f |
| where[p].add(idx) |
| touched.add(p) |
| for p in touched: |
| heapq.heappush(heap, (-pair_counts[p], p)) |
| pair_counts.pop(pair, None) |
| where.pop(pair, None) |
|
|
| if verbose and len(merges) % 500 == 0: |
| print(f" merge {len(merges)}/{n_merges}") |
|
|
| tok = cls(merges) |
| if verbose: |
| print(f"Tokenizer done: vocab {tok.vocab_size} ({len(merges)} merges)") |
| return tok |
|
|
| |
| def _encode_chunk(self, chunk: str) -> list[int]: |
| cached = self._cache.get(chunk) |
| if cached is not None: |
| return cached |
| ids = list(chunk.encode("utf-8")) |
| while len(ids) >= 2: |
| |
| best = min(zip(ids, ids[1:]), key=lambda p: self.ranks.get(p, 1 << 30)) |
| rank = self.ranks.get(best) |
| if rank is None: |
| break |
| new_id = 256 + rank |
| out, i = [], 0 |
| while i < len(ids): |
| if i < len(ids) - 1 and ids[i] == best[0] and ids[i + 1] == best[1]: |
| out.append(new_id) |
| i += 2 |
| else: |
| out.append(ids[i]) |
| i += 1 |
| ids = out |
| if len(self._cache) < 500_000: |
| self._cache[chunk] = ids |
| return ids |
|
|
| def encode(self, text: str, allow_special: bool = True) -> list[int]: |
| ids: list[int] = [] |
| parts = self._special_re.split(text) if allow_special else [text] |
| for part in parts: |
| if not part: |
| continue |
| if allow_special and part in self.special_to_id: |
| ids.append(self.special_to_id[part]) |
| continue |
| for chunk in GPT2_SPLIT.findall(part): |
| ids.extend(self._encode_chunk(chunk)) |
| return ids |
|
|
| def decode(self, ids: list[int]) -> str: |
| out = bytearray() |
| for i in ids: |
| if i in self.id_to_special: |
| out.extend(self.id_to_special[i].encode("utf-8")) |
| else: |
| out.extend(self.vocab[i]) |
| return out.decode("utf-8", errors="replace") |
|
|
| |
| def save(self, path: str) -> None: |
| with open(path, "w") as f: |
| json.dump({"merges": self.merges, "special_tokens": self.special_tokens}, f) |
|
|
| @classmethod |
| def load(cls, path: str) -> "BPETokenizer": |
| with open(path) as f: |
| d = json.load(f) |
| return cls([tuple(m) for m in d["merges"]], d["special_tokens"]) |
|
|
