"""Byte-level BPE tokenizer: training, encoding and decoding. IDs 0-255 are raw bytes, then one ID per learned merge, and the last 21 IDs of the vocabulary are special tokens. """ from __future__ import annotations import heapq import json from collections import Counter from collections.abc import Iterable from pathlib import Path import regex PRETOKENIZE_PATTERN = r" ?[\p{L}\p{M}]+|\p{N}|[^\s\p{L}\p{M}\p{N}]+|\s+(?!\S)|\s+" SPECIAL_TOKENS = ["<|bos|>", "<|eos|>", "<|pad|>", "<|user|>", "<|spec|>", "<|result|>", "<|answer|>", "<|end|>"] + [ f"<|reserved_{i}|>" for i in range(13) ] _PRETOKEN_RE = regex.compile(PRETOKENIZE_PATTERN) _DIGIT_BYTES = frozenset(range(ord("0"), ord("9") + 1)) _CACHE_LIMIT = 200_000 def pretokenize(text: str) -> list[str]: return _PRETOKEN_RE.findall(text) class BPETokenizer: def __init__(self, merges: list[tuple[int, int]], vocab_size: int): n_merges_max = vocab_size - 256 - len(SPECIAL_TOKENS) if len(merges) > n_merges_max: raise ValueError(f"{len(merges)} merges do not fit in vocab_size {vocab_size}") self.merges = [tuple(m) for m in merges] self.vocab_size = vocab_size first_special = vocab_size - len(SPECIAL_TOKENS) self.special_ids = {s: first_special + i for i, s in enumerate(SPECIAL_TOKENS)} self._special_by_id = {i: s for s, i in self.special_ids.items()} self._ranks = {pair: 256 + i for i, pair in enumerate(self.merges)} self._bytes: list[bytes] = [bytes([i]) for i in range(256)] for a, b in self.merges: self._bytes.append(self._bytes[a] + self._bytes[b]) self._special_re = regex.compile("(" + "|".join(regex.escape(s) for s in SPECIAL_TOKENS) + ")") self._cache: dict[str, list[int]] = {} def token_bytes(self, id: int) -> bytes: if id in self._special_by_id: return self._special_by_id[id].encode() if 0 <= id < len(self._bytes): return self._bytes[id] raise KeyError(f"token id {id} is not assigned") def _encode_pretoken(self, piece: str) -> list[int]: cached = self._cache.get(piece) if cached is not None: return cached ids = list(piece.encode("utf-8", errors="surrogatepass")) while len(ids) > 1: best = min( ((self._ranks.get((ids[i], ids[i + 1]), None), i) for i in range(len(ids) - 1)), key=lambda r: (r[0] is None, r[0] or 0), ) rank = best[0] if rank is None: break pair = self.merges[rank - 256] merged, i = [], 0 while i < len(ids): if i < len(ids) - 1 and (ids[i], ids[i + 1]) == pair: merged.append(rank) i += 2 else: merged.append(ids[i]) i += 1 ids = merged if len(self._cache) < _CACHE_LIMIT: self._cache[piece] = ids return ids def _encode_ordinary(self, text: str) -> list[int]: out: list[int] = [] for piece in _PRETOKEN_RE.findall(text): out.extend(self._encode_pretoken(piece)) return out def encode(self, text: str, allowed_special: bool = False) -> list[int]: if not allowed_special: return self._encode_ordinary(text) out: list[int] = [] for part in self._special_re.split(text): if part in self.special_ids: out.append(self.special_ids[part]) elif part: out.extend(self._encode_ordinary(part)) return out def decode(self, ids: list[int]) -> str: # Unassigned IDs (possible from an untrained model) decode to U+FFFD. parts = [] for i in ids: try: parts.append(self.token_bytes(i)) except KeyError: parts.append("�".encode()) return b"".join(parts).decode("utf-8", errors="replace") def save(self, path) -> None: data = { "version": 1, "pattern": PRETOKENIZE_PATTERN, "vocab_size": self.vocab_size, "merges": [list(m) for m in self.merges], "special": SPECIAL_TOKENS, } Path(path).write_text(json.dumps(data)) @classmethod def load(cls, path) -> BPETokenizer: data = json.loads(Path(path).read_text()) if data["pattern"] != PRETOKENIZE_PATTERN or data["special"] != SPECIAL_TOKENS: raise ValueError(f"{path} was saved with a different pattern or special-token list") return cls([tuple(m) for m in data["merges"]], data["vocab_size"]) def train_bpe(texts: Iterable[str], vocab_size: int = 16384) -> BPETokenizer: """Learn merges by repeatedly joining the most frequent adjacent pair. Ties go to the smallest (a, b) so training is deterministic. Pairs that contain a digit byte are never merged. """ counts: Counter[str] = Counter() for text in texts: counts.update(_PRETOKEN_RE.findall(text)) words: list[list[int]] = [] freqs: list[int] = [] for piece, f in counts.items(): b = list(piece.encode("utf-8", errors="surrogatepass")) if len(b) > 1: words.append(b) freqs.append(f) pair_counts: Counter[tuple[int, int]] = Counter() where: dict[tuple[int, int], set[int]] = {} def pairs_of(w: list[int]): for i in range(len(w) - 1): a, b = w[i], w[i + 1] if a not in _DIGIT_BYTES and b not in _DIGIT_BYTES: yield a, b for idx, w in enumerate(words): for p in pairs_of(w): pair_counts[p] += freqs[idx] where.setdefault(p, set()).add(idx) heap = [(-c, p[0], p[1]) for p, c in pair_counts.items()] heapq.heapify(heap) merges: list[tuple[int, int]] = [] n_merges = vocab_size - 256 - len(SPECIAL_TOKENS) while len(merges) < n_merges and heap: neg, a, b = heapq.heappop(heap) pair = (a, b) current = pair_counts.get(pair, 0) if current <= 0 or -neg != current: continue new_id = 256 + len(merges) merges.append(pair) touched: set[tuple[int, int]] = set() for idx in list(where.pop(pair, ())): w, f = words[idx], freqs[idx] old_pairs = list(pairs_of(w)) for p in old_pairs: pair_counts[p] -= f merged, i = [], 0 while i < len(w): if i < len(w) - 1 and w[i] == a and w[i + 1] == b: merged.append(new_id) i += 2 else: merged.append(w[i]) i += 1 words[idx] = merged new_pairs = list(pairs_of(merged)) for p in new_pairs: pair_counts[p] += f where.setdefault(p, set()).add(idx) for p in set(old_pairs) - set(new_pairs): s = where.get(p) if s is not None: s.discard(idx) touched.update(old_pairs) touched.update(new_pairs) pair_counts.pop(pair, None) for p in touched: c = pair_counts.get(p, 0) if c > 0: heapq.heappush(heap, (-c, p[0], p[1])) else: pair_counts.pop(p, None) return BPETokenizer(merges, vocab_size)