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:
@@ -130,6 +130,7 @@ experiment-llm/
|
||||
├── requirements.txt
|
||||
├── configs/ # one file per experiment (v1, v2, ...)
|
||||
├── scripts/ # one-shot entry points: download, encode, plots
|
||||
├── artifacts/ # trained tokenizer.json (small, committed)
|
||||
├── src/
|
||||
│ ├── tokenizer.py # BPE: train / encode / decode
|
||||
│ ├── dataset.py # memmapped token files, batch sampling
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Train the BPE tokenizer on a sample of the corpus.
|
||||
|
||||
Pure-Python BPE training is superlinear, so we train on a sample rather than
|
||||
the full 2 GB. TinyStories' restricted vocabulary means a ~50 MB sample yields
|
||||
essentially the same merges as the whole corpus.
|
||||
|
||||
Usage:
|
||||
python scripts/train_tokenizer.py # defaults: 50 MB, 4096 vocab
|
||||
python scripts/train_tokenizer.py --sample-mb 25 --vocab-size 4096
|
||||
|
||||
Writes the tokenizer to artifacts/tokenizer.json and prints the compression
|
||||
ratio (bytes per token) measured on held-out text.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
|
||||
from tokenizer import BPETokenizer # noqa: E402
|
||||
|
||||
EOT = "<|endoftext|>"
|
||||
TRAIN_TXT = ROOT / "data" / "raw" / "TinyStoriesV2-GPT4-train.txt"
|
||||
VALID_TXT = ROOT / "data" / "raw" / "TinyStoriesV2-GPT4-valid.txt"
|
||||
OUT = ROOT / "artifacts" / "tokenizer.json"
|
||||
|
||||
|
||||
def read_prefix(path: Path, mb: float) -> str:
|
||||
"""Read the first `mb` megabytes of a file as UTF-8 (ignoring a split char)."""
|
||||
with open(path, "rb") as f:
|
||||
data = f.read(int(mb * 1_000_000))
|
||||
return data.decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--sample-mb", type=float, default=50.0)
|
||||
ap.add_argument("--vocab-size", type=int, default=4096)
|
||||
args = ap.parse_args()
|
||||
|
||||
if not TRAIN_TXT.exists():
|
||||
sys.exit(f"missing {TRAIN_TXT} — run scripts/download_data.py first")
|
||||
|
||||
print(f"reading {args.sample_mb:g} MB sample from {TRAIN_TXT.name}")
|
||||
sample = read_prefix(TRAIN_TXT, args.sample_mb)
|
||||
|
||||
print(f"training BPE: target vocab {args.vocab_size}, special token {EOT!r}")
|
||||
t0 = time.perf_counter()
|
||||
tok = BPETokenizer.train(sample, vocab_size=args.vocab_size, special_tokens=[EOT])
|
||||
dt = time.perf_counter() - t0
|
||||
print(f"trained in {dt:.1f}s — {len(tok.merges)} merges, "
|
||||
f"vocab_size {tok.vocab_size}")
|
||||
|
||||
OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
tok.save(OUT)
|
||||
print(f"saved tokenizer to {OUT}")
|
||||
|
||||
# Compression ratio on held-out validation text (not used for training).
|
||||
held = read_prefix(VALID_TXT, 5.0) if VALID_TXT.exists() else sample[: 5_000_000]
|
||||
ids = tok.encode(held)
|
||||
n_bytes = len(held.encode("utf-8"))
|
||||
print(f"\ncompression on held-out text: {n_bytes:,} bytes -> {len(ids):,} tokens"
|
||||
f" ({n_bytes / len(ids):.2f} bytes/token)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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"])
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Tokenizer tests. Run directly: `.venv/bin/python tests/test_tokenizer.py`
|
||||
|
||||
No pytest dependency — plain asserts and a tiny runner, to keep the project's
|
||||
dependencies limited to torch + numpy.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
||||
|
||||
from tokenizer import BPETokenizer, pre_tokenize # noqa: E402
|
||||
|
||||
EOT = "<|endoftext|>"
|
||||
|
||||
# A small but varied training corpus.
|
||||
CORPUS = (
|
||||
"Once upon a time, there was a little cat. The cat liked to play.\n"
|
||||
"The cat ran and ran. Then the cat found a red ball!\n" + EOT + "\n"
|
||||
"Tom and Lily went to the park. They were very happy.\n"
|
||||
) * 50
|
||||
|
||||
|
||||
def _fresh_tokenizer() -> BPETokenizer:
|
||||
return BPETokenizer.train(CORPUS, vocab_size=400, special_tokens=[EOT])
|
||||
|
||||
|
||||
def test_pretokenize_is_lossless():
|
||||
for text in ["Hello, world!", " spaces\tand\nnewlines ", "a", "", "!!!??"]:
|
||||
assert "".join(pre_tokenize(text)) == text
|
||||
|
||||
|
||||
def test_roundtrip_on_training_text():
|
||||
tok = _fresh_tokenizer()
|
||||
for line in CORPUS.split("\n"):
|
||||
assert tok.decode(tok.encode(line)) == line
|
||||
|
||||
|
||||
def test_roundtrip_on_unseen_text():
|
||||
tok = _fresh_tokenizer()
|
||||
unseen = "A brave dog jumped over the fence; zebras watched quietly."
|
||||
assert tok.decode(tok.encode(unseen)) == unseen
|
||||
|
||||
|
||||
def test_special_token_is_single_id():
|
||||
tok = _fresh_tokenizer()
|
||||
ids = tok.encode(f"Hello {EOT} world")
|
||||
eot_id = tok.special_tokens[EOT]
|
||||
assert ids.count(eot_id) == 1
|
||||
# The special token id must not appear from encoding ordinary text.
|
||||
assert eot_id not in tok.encode("Hello world")
|
||||
assert tok.decode(tok.encode(f"a{EOT}b")) == f"a{EOT}b"
|
||||
|
||||
|
||||
def test_unicode_roundtrip():
|
||||
tok = _fresh_tokenizer()
|
||||
# Multibyte characters must survive being split across byte-level tokens.
|
||||
text = "café naïve \U0001f600"
|
||||
assert tok.decode(tok.encode(text)) == text
|
||||
|
||||
|
||||
def test_vocab_size_and_layout():
|
||||
tok = _fresh_tokenizer()
|
||||
# vocab_size is a target ceiling: a tiny, repetitive corpus saturates
|
||||
# (every word merges to a single token) before reaching it, so we may get
|
||||
# fewer merges. The derived vocab dict must always match the reported size.
|
||||
assert tok.vocab_size <= 400
|
||||
assert len(tok.vocab) == tok.vocab_size
|
||||
# First 256 ids are the raw bytes.
|
||||
assert tok.vocab[65] == b"A"
|
||||
|
||||
|
||||
def test_save_load_roundtrip():
|
||||
tok = _fresh_tokenizer()
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
path = Path(d) / "tok.json"
|
||||
tok.save(path)
|
||||
loaded = BPETokenizer.load(path)
|
||||
sample = f"The cat and Tom. {EOT}"
|
||||
assert loaded.encode(sample) == tok.encode(sample)
|
||||
assert loaded.vocab_size == tok.vocab_size
|
||||
|
||||
|
||||
def test_determinism():
|
||||
a = _fresh_tokenizer()
|
||||
b = _fresh_tokenizer()
|
||||
assert a.merges == b.merges
|
||||
|
||||
|
||||
def main() -> None:
|
||||
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
||||
for t in tests:
|
||||
t()
|
||||
print(f"ok {t.__name__}")
|
||||
print(f"\n{len(tests)} passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user