| """Byte-level BPE tokenizer for source code, implemented from scratch. |
| |
| * Base vocabulary: the 256 byte values, so any text (any language, any |
| Unicode) can be encoded with no unknown tokens. |
| * Merges are learned from the training corpus (classic BPE). |
| * Text is first split into chunks with a code-aware regex so merges never |
| cross identifier/operator/whitespace boundaries; runs of indentation |
| become single tokens, which matters a lot for Python. |
| * Special tokens (``<|endoftext|>``, chat markers) are appended after the |
| learned vocabulary and are never produced by merges. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import re |
| from collections import Counter, defaultdict |
| from pathlib import Path |
| from typing import Iterable |
|
|
| EOT = "<|endoftext|>" |
| USER = "<|user|>" |
| ASSISTANT = "<|assistant|>" |
| END = "<|end|>" |
| SPECIAL_TOKENS = (EOT, USER, ASSISTANT, END) |
|
|
| |
| |
| CHUNK_RE = re.compile( |
| r"\n[ \t]*" |
| r"| ?[A-Za-z_][A-Za-z0-9_]*" |
| r"| ?[0-9]{1,3}" |
| r"| ?[^\sA-Za-z0-9_]+" |
| r"|[ \t]+(?=\S)|[ \t]+" |
| r"|\s" |
| ) |
|
|
|
|
| def split_chunks(text: str) -> list[str]: |
| return CHUNK_RE.findall(text) |
|
|
|
|
| class BPETokenizer: |
| def __init__(self, merges: list[tuple[int, int]] | None = None, |
| special_tokens: Iterable[str] = SPECIAL_TOKENS) -> None: |
| self.merges: list[tuple[int, int]] = list(merges or []) |
| self.ranks: dict[tuple[int, int], int] = {pair: i for i, pair in enumerate(self.merges)} |
| |
| self.vocab: list[bytes] = [bytes([i]) for i in range(256)] |
| for a, b in self.merges: |
| self.vocab.append(self.vocab[a] + self.vocab[b]) |
| self.special_tokens = list(special_tokens) |
| base = len(self.vocab) |
| self.special_ids = {tok: base + i for i, tok in enumerate(self.special_tokens)} |
| self.special_by_id = {i: tok for tok, i in self.special_ids.items()} |
| self._special_re = ( |
| re.compile("(" + "|".join(re.escape(t) for t in self.special_tokens) + ")") |
| if self.special_tokens else None |
| ) |
| self._cache: dict[str, list[int]] = {} |
|
|
| |
|
|
| @property |
| def vocab_size(self) -> int: |
| return len(self.vocab) + len(self.special_tokens) |
|
|
| def token_id(self, special: str) -> int: |
| return self.special_ids[special] |
|
|
| @property |
| def eot_id(self) -> int: |
| return self.special_ids[EOT] |
|
|
| |
|
|
| @classmethod |
| def train(cls, texts: Iterable[str], vocab_size: int, *, min_frequency: int = 2, |
| special_tokens: Iterable[str] = SPECIAL_TOKENS, verbose: bool = False) -> "BPETokenizer": |
| """Learn ``vocab_size - 256 - len(special)`` merges from ``texts``.""" |
| special_tokens = list(special_tokens) |
| n_merges = vocab_size - 256 - len(special_tokens) |
| if n_merges < 0: |
| raise ValueError(f"vocab_size must be at least {256 + len(special_tokens)}") |
|
|
| chunk_counts: Counter[str] = Counter() |
| for text in texts: |
| chunk_counts.update(split_chunks(text)) |
|
|
| words: list[list[int]] = [] |
| freqs: list[int] = [] |
| for chunk, count in chunk_counts.items(): |
| data = list(chunk.encode("utf-8")) |
| if len(data) > 1: |
| words.append(data) |
| freqs.append(count) |
|
|
| pair_counts: dict[tuple[int, int], int] = defaultdict(int) |
| where: dict[tuple[int, int], set[int]] = defaultdict(set) |
| for idx, word in enumerate(words): |
| for pair in zip(word, word[1:]): |
| pair_counts[pair] += freqs[idx] |
| where[pair].add(idx) |
|
|
| merges: list[tuple[int, int]] = [] |
| next_id = 256 |
| for step in range(n_merges): |
| if not pair_counts: |
| break |
| best = max(pair_counts.items(), key=lambda kv: (kv[1], -kv[0][0], -kv[0][1])) |
| pair, count = best |
| if count < min_frequency: |
| break |
| merges.append(pair) |
| new_id = next_id |
| next_id += 1 |
| for idx in list(where.get(pair, ())): |
| word = words[idx] |
| freq = freqs[idx] |
| |
| for p in zip(word, word[1:]): |
| pair_counts[p] -= freq |
| if pair_counts[p] <= 0: |
| del pair_counts[p] |
| where[p].discard(idx) |
| merged: list[int] = [] |
| i = 0 |
| while i < len(word): |
| if i + 1 < len(word) and word[i] == pair[0] and word[i + 1] == pair[1]: |
| merged.append(new_id) |
| i += 2 |
| else: |
| merged.append(word[i]) |
| i += 1 |
| words[idx] = merged |
| for p in zip(merged, merged[1:]): |
| pair_counts[p] += freq |
| where[p].add(idx) |
| where.pop(pair, None) |
| if verbose and (step + 1) % 500 == 0: |
| print(f" merges: {step + 1}/{n_merges}") |
| return cls(merges, special_tokens) |
|
|
| |
|
|
| 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) > 1: |
| best_rank, best_pos = None, -1 |
| for i in range(len(ids) - 1): |
| rank = self.ranks.get((ids[i], ids[i + 1])) |
| if rank is not None and (best_rank is None or rank < best_rank): |
| best_rank, best_pos = rank, i |
| if best_rank is None: |
| break |
| ids[best_pos: best_pos + 2] = [256 + best_rank] |
| if len(self._cache) < 200_000: |
| self._cache[chunk] = ids |
| return ids |
|
|
| def encode(self, text: str, *, allow_special: bool = True) -> list[int]: |
| out: list[int] = [] |
| parts = self._special_re.split(text) if (allow_special and self._special_re) else [text] |
| for part in parts: |
| if not part: |
| continue |
| if allow_special and part in self.special_ids: |
| out.append(self.special_ids[part]) |
| continue |
| for chunk in split_chunks(part): |
| out.extend(self._encode_chunk(chunk)) |
| return out |
|
|
| def decode(self, ids: Iterable[int], *, skip_special: bool = False) -> str: |
| buf = bytearray() |
| pieces: list[str] = [] |
| for i in ids: |
| if i in self.special_by_id: |
| if buf: |
| pieces.append(buf.decode("utf-8", errors="replace")) |
| buf = bytearray() |
| if not skip_special: |
| pieces.append(self.special_by_id[i]) |
| elif 0 <= i < len(self.vocab): |
| buf.extend(self.vocab[i]) |
| if buf: |
| pieces.append(buf.decode("utf-8", errors="replace")) |
| return "".join(pieces) |
|
|
| |
|
|
| def to_dict(self) -> dict: |
| return {"type": "ycode-bpe", "version": 1, "merges": self.merges, |
| "special_tokens": self.special_tokens} |
|
|
| def save(self, path: Path) -> None: |
| Path(path).write_text(json.dumps(self.to_dict()), encoding="utf-8") |
|
|
| @classmethod |
| def from_dict(cls, data: dict) -> "BPETokenizer": |
| if data.get("type") != "ycode-bpe": |
| raise ValueError("not a YCode tokenizer file") |
| return cls([tuple(m) for m in data["merges"]], data.get("special_tokens", SPECIAL_TOKENS)) |
|
|
| @classmethod |
| def load(cls, path: Path) -> "BPETokenizer": |
| return cls.from_dict(json.loads(Path(path).read_text(encoding="utf-8"))) |
|
|
