File size: 6,955 Bytes
a2b8b4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """
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
# GPT-2 pre-tokenization pattern: contractions, words, numbers, punctuation, whitespace
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()
# ------------------------------------------------------------------ setup
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]
# --------------------------------------------------------------- training
@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)
# max-heap with lazy invalidation: stale entries are skipped on pop
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 # stale
new_id = 256 + len(merges)
merges.append(pair)
for idx in list(where[pair]):
w, f = words[idx], freqs[idx]
# remove the old pairs' contribution
for p in zip(w, w[1:]):
pair_counts[p] -= f
where[p].discard(idx)
# merge the pair
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
# add the new pairs' contribution
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
# ------------------------------------------------------------ encode/decode
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:
# lowest-rank pair = the merge learned earliest
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")
# ---------------------------------------------------------------- persist
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"])
|