Phase 1: from-scratch byte-level BPE tokenizer
- 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
This commit is contained in:
@@ -0,0 +1,188 @@
|
||||
"""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"])
|
||||
Reference in New Issue
Block a user