rzen
2c5fc29592
- src/tokenizer.py: byte-level BPE with GPT-2-style word pre-tokenization, protected <|endoftext|> special token, save/load as inspectable JSON. Simple recount-per-merge trainer over unique words. - tests/test_tokenizer.py: round-trip, special-token integrity, unicode, lossless pre-tokenization, save/load, determinism (8 tests, no pytest dep) - scripts/train_tokenizer.py: train on a corpus sample, report bytes/token
189 lines
7.4 KiB
Python
189 lines
7.4 KiB
Python
"""Byte-level Byte-Pair Encoding (BPE) tokenizer, implemented from scratch.
|
|
|
|
Design choices (optimised for readability over speed):
|
|
|
|
- **Byte-level.** The base vocabulary is the 256 possible byte values, so any
|
|
UTF-8 text is representable and there are no "unknown token" holes.
|
|
- **Word pre-tokenization.** Text is first split into word-ish chunks with a
|
|
regex (GPT-2 style, keeping a word's leading space attached: " the" is one
|
|
chunk). BPE merges only ever happen *within* a chunk, never across word or
|
|
whitespace boundaries — this is what keeps the learned tokens sensible.
|
|
- **Special tokens.** Strings like ``<|endoftext|>`` are reserved: they get a
|
|
fixed id, are never split by BPE, and decode back to their literal text.
|
|
|
|
ID layout: 0..255 raw bytes | 256..256+M-1 merges | then one id per special.
|
|
|
|
The trainer recounts pair frequencies from scratch on every merge. That is the
|
|
simple, legible algorithm; it operates over *unique* pre-tokenized words (not
|
|
the raw byte stream), which keeps it fast enough to train a 4k vocab on a
|
|
~50 MB sample in a few minutes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
# Tiles the input completely: a word (optionally led by one space), or a run of
|
|
# punctuation/symbols (optionally led by one space), or a run of whitespace, or
|
|
# — as a catch-all so no character is ever dropped — any single character.
|
|
# re.DOTALL lets the final "." also match newlines.
|
|
_SPLIT_PATTERN = re.compile(r" ?\w+| ?[^\w\s]+|\s+|.", re.DOTALL)
|
|
|
|
|
|
def pre_tokenize(text: str) -> list[str]:
|
|
"""Split text into word-ish chunks. Concatenating the result rebuilds text."""
|
|
return _SPLIT_PATTERN.findall(text)
|
|
|
|
|
|
def _merge(ids: list[int], pair: tuple[int, int], new_id: int) -> list[int]:
|
|
"""Replace every non-overlapping occurrence of `pair` in `ids` with `new_id`."""
|
|
out: list[int] = []
|
|
i = 0
|
|
while i < len(ids):
|
|
if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
|
|
out.append(new_id)
|
|
i += 2
|
|
else:
|
|
out.append(ids[i])
|
|
i += 1
|
|
return out
|
|
|
|
|
|
class BPETokenizer:
|
|
def __init__(
|
|
self,
|
|
merges: list[tuple[int, int]],
|
|
special_tokens: dict[str, int],
|
|
) -> None:
|
|
self.merges = merges
|
|
self.special_tokens = special_tokens
|
|
# rank = merge priority (lower merges first); id of a merged pair = 256 + rank
|
|
self.ranks = {pair: rank for rank, pair in enumerate(merges)}
|
|
self.vocab = self._build_vocab(merges, special_tokens)
|
|
|
|
@staticmethod
|
|
def _build_vocab(
|
|
merges: list[tuple[int, int]],
|
|
special_tokens: dict[str, int],
|
|
) -> dict[int, bytes]:
|
|
"""Map every id to the byte string it expands to (used for decoding)."""
|
|
vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
|
|
for rank, (a, b) in enumerate(merges):
|
|
vocab[256 + rank] = vocab[a] + vocab[b]
|
|
for text, idx in special_tokens.items():
|
|
vocab[idx] = text.encode("utf-8")
|
|
return vocab
|
|
|
|
@property
|
|
def vocab_size(self) -> int:
|
|
return 256 + len(self.merges) + len(self.special_tokens)
|
|
|
|
# ------------------------------------------------------------------ training
|
|
@classmethod
|
|
def train(
|
|
cls,
|
|
text: str,
|
|
vocab_size: int,
|
|
special_tokens: list[str] | None = None,
|
|
) -> "BPETokenizer":
|
|
special_tokens = special_tokens or []
|
|
num_merges = vocab_size - 256 - len(special_tokens)
|
|
if num_merges < 0:
|
|
raise ValueError("vocab_size too small for 256 bytes + special tokens")
|
|
|
|
# Count unique pre-tokenized words. Special-token strings are stripped
|
|
# first so BPE never learns to merge pieces of "<|endoftext|>".
|
|
word_freqs: Counter[bytes] = Counter()
|
|
for segment, is_special in cls._split_on_specials(text, special_tokens):
|
|
if is_special:
|
|
continue
|
|
for chunk in pre_tokenize(segment):
|
|
word_freqs[chunk.encode("utf-8")] += 1
|
|
|
|
# Each unique word is a list of byte ids that we progressively merge.
|
|
words: dict[bytes, list[int]] = {w: list(w) for w in word_freqs}
|
|
|
|
merges: list[tuple[int, int]] = []
|
|
for _ in range(num_merges):
|
|
pair_counts: Counter[tuple[int, int]] = Counter()
|
|
for w, freq in word_freqs.items():
|
|
ids = words[w]
|
|
for pair in zip(ids, ids[1:]):
|
|
pair_counts[pair] += freq
|
|
if not pair_counts:
|
|
break # nothing left to merge
|
|
# Highest count wins; ties broken by pair value for reproducibility.
|
|
best = max(pair_counts.items(), key=lambda kv: (kv[1], kv[0]))[0]
|
|
new_id = 256 + len(merges)
|
|
merges.append(best)
|
|
for w in words:
|
|
words[w] = _merge(words[w], best, new_id)
|
|
|
|
specials = {tok: 256 + len(merges) + i for i, tok in enumerate(special_tokens)}
|
|
return cls(merges, specials)
|
|
|
|
# ------------------------------------------------------------------ encoding
|
|
@staticmethod
|
|
def _split_on_specials(text: str, specials: list[str]):
|
|
"""Yield (segment, is_special), splitting on any special-token string."""
|
|
if not specials:
|
|
if text:
|
|
yield text, False
|
|
return
|
|
pattern = "(" + "|".join(re.escape(s) for s in specials) + ")"
|
|
for part in re.split(pattern, text):
|
|
if part == "":
|
|
continue
|
|
yield part, part in specials
|
|
|
|
def _encode_chunk(self, chunk: str) -> list[int]:
|
|
"""Apply learned merges to one pre-tokenized word, best (lowest) rank first."""
|
|
ids = list(chunk.encode("utf-8"))
|
|
while len(ids) >= 2:
|
|
# Find the present pair with the lowest merge rank.
|
|
best_rank: int | None = None
|
|
best_i = -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_i = rank, i
|
|
if best_rank is None:
|
|
break # no learned merge applies
|
|
ids = _merge(ids, self.merges[best_rank], 256 + best_rank)
|
|
return ids
|
|
|
|
def encode(self, text: str) -> list[int]:
|
|
ids: list[int] = []
|
|
for segment, is_special in self._split_on_specials(
|
|
text, list(self.special_tokens)
|
|
):
|
|
if is_special:
|
|
ids.append(self.special_tokens[segment])
|
|
else:
|
|
for chunk in pre_tokenize(segment):
|
|
ids.extend(self._encode_chunk(chunk))
|
|
return ids
|
|
|
|
def decode(self, ids: list[int]) -> str:
|
|
data = b"".join(self.vocab[i] for i in ids)
|
|
return data.decode("utf-8", errors="replace")
|
|
|
|
# ------------------------------------------------------------------ persistence
|
|
def save(self, path: str | Path) -> None:
|
|
"""Save as human-inspectable JSON (merges + specials; vocab is derived)."""
|
|
path = Path(path)
|
|
blob = {
|
|
"special_tokens": self.special_tokens,
|
|
"merges": [[a, b] for a, b in self.merges],
|
|
}
|
|
path.write_text(json.dumps(blob))
|
|
|
|
@classmethod
|
|
def load(cls, path: str | Path) -> "BPETokenizer":
|
|
blob = json.loads(Path(path).read_text())
|
|
merges = [(a, b) for a, b in blob["merges"]]
|
|
return cls(merges, blob["special_tokens"])
|