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
71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
"""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()
|