File size: 8,149 Bytes
5ef3875 | 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 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | """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)
# Code-aware pre-tokenization: newline+indent runs, identifiers (with one
# leading space), short digit runs, operator runs, other whitespace.
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)}
# id -> bytes for every non-special token
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]] = {}
# ------------------------------------------------------------- props
@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]
# ---------------------------------------------------------- training
@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]
# Remove this word's old pair contributions.
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)
# ---------------------------------------------------------- encoding
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)
# ------------------------------------------------------- persistence
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")))
|