Tokenizer abstraction: char + byte-level BPE

[scheuclu]

Summary

Adds a `Tokenizer` base class and two concrete implementations, switchable via CLI.

New `tokenizers/` package

• `base.py` — `Tokenizer` ABC: `encode`, `decode`, `vocab_size`, `train`, `state_dict` / `load_state_dict`
• `char.py` — `CharTokenizer`: one token per unique character (~85–250 vocab on English-ish text). Behavior identical to what was inlined in `gpt.py` before.
• `bpe.py` — `BPETokenizer`: byte-level BPE, minbpe-style. Greedy merges of the most-frequent adjacent pair until target vocab is reached. ~3–4× sequence compression on English text.

`gpt.py` integration

• Replaces the inline char-level setup with `tokenizer = tok.get(args.tokenizer, ...); tokenizer.train(text)`.
• Two new CLI flags: `--tokenizer {char,bpe}` (default `char`) and `--tokenizer-vocab-size N` (default 1024, BPE only).
• Encoded corpus is cached to `{dataset_path}.{tokenizer}_v{vocab_size}.cache.pt` keyed on tokenizer state (stale cache auto-regenerates). BPE encoding under pure Python takes minutes on the 2 GB corpus, so the cache pays off after the first run.
• BPE merge-training runs on a 5 MB prefix of the corpus by default — full-corpus training under pure Python is prohibitive, and the prefix typically converges to near-identical merges.
• Token-distribution TB logging now counts token IDs in the encoded data tensor and displays each token's decoded string, so it works the same way for char and BPE.
• `lm_head.bias` log-unigram init uses `torch.bincount(data)` instead of `Counter(text)` — works for any tokenizer.

Checkpoint format

• Gains `tokenizer_type` and `tokenizer_state` fields.
• Char tokenizers still write the legacy top-level `chars` field so existing `viz_embeddings.py`, `export_onnx.py`, the web frontend, and published checkpoints keep working unchanged.
• `load_model_from_checkpoint` returns `(model, tokenizer, hp)` instead of `(model, chars, hp)`. Falls back to constructing a `CharTokenizer` from `chars` for legacy checkpoints.

Downstream tool updates

• `viz_embeddings.py` is char-tokenizer-only (per-character category coloring doesn't translate to subwords). It now raises a clear error if loaded against a BPE checkpoint.
• `export_onnx.py` writes `tokens` (one decoded string per ID) to `vocab.json` for both tokenizers, plus the legacy `chars` field for char-level so the existing JS frontend works unchanged. BPE prompt encoding in JS is out of scope here — would need a JS BPE encoder.

Test plan

• `uv run pyright` clean
• Legacy `ckpt_default_step_03375.pt` (which only has `chars`, no `tokenizer_state`) loads through the new path
• BPE checkpoint round-trips through state_dict — encoded IDs match before and after save/load
• `export_onnx.py` on the legacy checkpoint produces a vocab.json with `chars` populated
• Training run with `--tokenizer bpe --tokenizer-vocab-size 1024` on Shakespeare completes a few eval intervals
• Compare loss curves: char vs BPE-1024 on the same dataset and profile

🤖 Generated with Claude Code

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 67156d4c-c82c-4e6e-bb6f-21b97dc5fa46

📥 Commits

Reviewing files that changed from the base of the PR and between 4b83ab7 and 3215423.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (11)

• .gitignore
• README.md
• datasets/gutenberg.py
• export_onnx.py
• gpt.py
• pyproject.toml
• tokenizers/__init__.py
• tokenizers/base.py
• tokenizers/bpe.py
• tokenizers/char.py
• viz_embeddings.py

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added support for Byte-Pair Encoding (BPE) tokenizer alongside existing character-level tokenization
  • Implemented parallel dataset downloading for faster training data preparation
  • Added learning rate warmup and adaptive scheduling during model training
  • Increased default training dataset from 3,000 to 10,000 books (~7 GB)

• Documentation

  • Updated README with tokenizer selection, training, and caching information

• Chores

  • Added numba as a runtime dependency

Walkthrough

This PR replaces fixed character-level vocabulary with a configurable tokenizer system supporting character and byte-pair encoding, updates model training to use parallelized corpus encoding and learning-rate scheduling, and adapts all dependent tools to work with tokenizer-aware checkpoints.

Changes

Tokenizer System and Model Integration

Layer / File(s)	Summary
Tokenizer Base Contract and Registry tokenizers/base.py, tokenizers/__init__.py	Introduces abstract Tokenizer base class defining encode/decode, training, and serialization contracts, plus a module-level registry that discovers and factory-instantiates tokenizers by name.
Character-Level Tokenizer tokenizers/char.py	Implements CharTokenizer: learns character vocabulary from corpus, encodes text to character IDs (skipping OOV characters), decodes ID sequences back to text, and persists learned character set via state_dict/load_state_dict.
Byte-Pair Encoding Tokenizer tokenizers/bpe.py	Implements BPETokenizer: learns byte-pair merge rules greedily from a UTF-8 corpus prefix, builds a dense merge-lookup matrix, uses numba-compiled JIT encoding to apply merges, decodes via stored byte sequences, and persists learned merges.
Model API Refactoring gpt.py (imports, Hyperparameters, load_model_from_checkpoint)	Updates Hyperparameters with learning-rate schedule configuration (warmup, ReduceLROnPlateau parameters), rewrites load_model_from_checkpoint() to return (model, tokenizer, hparams) instead of (model, chars, hparams), and handles both new tokenizer-state and legacy chars-only checkpoint formats.
Parallelized Corpus Encoding gpt.py (multiprocessing corpus encoding, training tokenizer initialization)	Adds multiprocessing-based encoding pipeline: per-worker tokenizer initialization, chunked parallel text encoding via _encode_corpus_with_progress(), intelligent caching keyed by dataset/tokenizer-config, and progress reporting with elapsed time and ETA.
Training Loop with Scheduling gpt.py (LM head bias, warmup, scheduler, checkpoints)	Computes LM head bias from tokenizer-encoded token frequencies, applies manual warmup-LR for early iterations, wires ReduceLROnPlateau scheduler for post-warmup LR control, updates TensorBoard logging to emit tokenizer-based token tables, and saves tokenizer_type/tokenizer_state in checkpoints with optional legacy chars field.
CLI Arguments and Inference gpt.py (CLI args, inference decoding)	Adds --tokenizer and --tokenizer-vocab-size CLI arguments, updates inference to load and print tokenizer details, and changes output decoding to use tokenizer.decode().
ONNX Export Metadata export_onnx.py	Updates metadata generation to unpack tokenizer from checkpoint, generates per-token vocabulary list from tokenizer.vocab_size, and conditionally preserves legacy chars field only for char-level tokenizers.
Gutenberg Dataset Parallelization datasets/gutenberg.py	Replaces sequential download with ThreadPoolExecutor for parallel HTTP fetching, introduces _fetch_book() helper for individual book download/normalization, increases default subset from 3000 to 10000 books, adds download_workers=32 configuration, and enhances progress reporting with elapsed time and ETA.
Visualization Tool Updates viz_embeddings.py	Adds cached load_tokenizer() for independent tokenizer loading, updates load_checkpoint() to extract character vocabulary from both chars field and tokenizer state, validates embeddings operate only with CharTokenizer, introduces new "Tokens" tab showing decoded vocabulary, and gates embedding tabs with informational messages when char embeddings unavailable.
Configuration and Documentation pyproject.toml, .gitignore, README.md	Adds numba>=0.65.1 to dependencies and tokenizers to Pyright include, updates .gitignore patterns for input_gutenberg*.txt and *.cache.pt, and adds README sections documenting char/BPE tokenization, corpus caching, checkpoint serialization, and repository layout.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• scheuclu/lukasGPT#15: Prior PR that introduced .tril buffer removal in checkpoints for compatibility, directly related to the stale buffer filtering added in this PR's load_model_from_checkpoint refactoring.

Poem

🐰 The tokenizer bunny hops in with grace,
Swapping chars for tokens all over the place,
BPE merges bytes with a JIT-compiled cheer,
While parallel downloads bring datasets near.
Warmth schedules the learning, checkpoints save state—
A better abstraction makes GPT feel great! ✨

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feature/bpe-tokenizer

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}