| """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: |
| |
| 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) |
|
|
