Tokenizer abstraction: char + byte-level BPE

.gitignore

@@ -5,6 +5,7 @@ checkpoints/
 viz/
 input.txt
 input_shakespeare.txt
-input_gutenberg.txt
+input_gutenberg*.txt
+*.cache.pt
 donotcommit.txt
 runs/

README.md

@@ -45,6 +45,26 @@ Downloads verify sha256 and skip files that are already present.
 
 Switch profiles by editing `ACTIVE_PROFILE` in `gpt.py`. Each checkpoint embeds the architecture and vocab so it's self-contained for inference.
 
+### Tokenizer
+
+The training script supports two tokenizers:
+
+```bash
+uv run python gpt.py --tokenizer char                          # default; one token per character
+uv run python gpt.py --tokenizer bpe --tokenizer-vocab-size 1024  # byte-level BPE (minbpe-style)
+```
+
+`char` keeps the original behavior: vocab = unique characters in the corpus (~85–250). `bpe` trains byte-pair merges on a 5 MB prefix of the corpus and gives you ~3–4× sequence compression on English, so `block_size=512` covers ~1.5–2k chars of real context. The encoded corpus is cached to disk (`{dataset_path}.{tokenizer}_v{vocab_size}.cache.pt`) keyed on the tokenizer state, since naive-Python BPE encoding of a 2 GB corpus takes minutes.
+
+Checkpoints store the tokenizer state, so inference reconstructs the exact same tokenizer:
+
+```bash
+uv run python gpt.py --infer checkpoints/ckpt_default_step_04999.pt
+#   tokenizer: bpe (vocab=1024)
+```
+
+Legacy checkpoints (only `chars` saved) load through a compatibility path as if they were `char` checkpoints — nothing breaks.
+
 ## Monitoring with TensorBoard
 
 Each training run writes scalar loss curves and a 200-character text sample at every eval step to `./runs/<timestamp>_<profile>_<dataset>/`. Disable with `--no-tensorboard`, or keep the curves but skip generation with `--no-sample` (worth it on the `large` profile).
@@ -109,6 +129,7 @@ The sidebar lets you scrub across training-step checkpoints and watch the embedd
 ## Layout
 
 - `gpt.py` — the transformer, training loop, inference, and lookahead sampling
+- `tokenizers/` — char and byte-level BPE tokenizer classes (local package; not HuggingFace's)
 - `bigram.py` — the tiny bigram baseline from earlier in the lecture
 - `viz_embeddings.py` — Streamlit embedding viewer
 - `checkpoint_io.py` — shared helpers for checkpoint files (sha256, git SHA tag, HF Hub upload, manifest)

datasets/gutenberg.py

@@ -8,8 +8,10 @@
 
 import os
 import re
+import time
 import urllib.request
 from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
 
 from huggingface_hub import (
     hf_hub_download,  # pyright: ignore[reportUnknownVariableType]
@@ -29,15 +31,30 @@
 PG19_MANIFEST_FILE = "data/train_files.txt"
 
 
+def _fetch_book(rel: str) -> str:
+    """Download one book from GCS, strip the trailing license footer if
+    present, and normalize trailing whitespace. Network-bound; safe to
+    run in many threads concurrently."""
+    with urllib.request.urlopen(PG19_GCS_BASE + rel) as r:
+        raw = r.read().decode("utf-8", errors="replace")
+    m = _END_MARKER.search(raw)
+    if m:
+        raw = raw[: m.start()]
+    return raw.strip() + "\n\n"
+
+
 class Gutenberg(Dataset):
     name = "gutenberg"
     url = PG19_GCS_BASE
-    default_path = "input_gutenberg.txt"
+    default_path = "input_gutenberg_10k.txt"
     description = (
         "DeepMind PG-19 subset (Project Gutenberg books pre-1919). "
-        "First 3000 books, ~2 GB. Set max_books=None to download all 28,602."
+        "First 10000 books, ~7 GB. Set max_books=None to download all 28,602."
     )
-    max_books: int | None = 3000
+    max_books: int | None = 10000
+    # Network-bound; 32 concurrent connections is comfortable against GCS
+    # and gets us roughly an order-of-magnitude speedup over sequential.
+    download_workers: int = 32
 
     def prepare(self, path: str | None = None) -> str:
         path = path or self.default_path
@@ -55,25 +72,32 @@ def prepare(self, path: str | None = None) -> str:
         if self.max_books is not None:
             book_paths = book_paths[: self.max_books]
 
-        print(f"downloading {len(book_paths):,} PG-19 books to {path}")
+        print(
+            f"downloading {len(book_paths):,} PG-19 books to {path} "
+            f"({self.download_workers} workers in parallel)"
+        )
         tmp = f"{path}.partial"
         chars = 0
+        t0 = time.time()
         with open(tmp, "w", encoding="utf-8") as out:
-            for i, rel in enumerate(book_paths):
-                with urllib.request.urlopen(PG19_GCS_BASE + rel) as r:
-                    raw = r.read().decode("utf-8", errors="replace")
-                m = _END_MARKER.search(raw)
-                if m:
-                    raw = raw[: m.start()]
-                block = raw.strip() + "\n\n"
-                out.write(block)
-                chars += len(block)
-                if (i + 1) % 50 == 0 or i + 1 == len(book_paths):
-                    print(
-                        f"\r  {i + 1}/{len(book_paths)} books · {chars / 1e9:.2f} GB",
-                        end="",
-                        flush=True,
-                    )
+            with ThreadPoolExecutor(max_workers=self.download_workers) as ex:
+                # executor.map preserves submission order, so the file is
+                # written in book_paths order even though fetches finish
+                # out of order across worker threads.
+                for i, block in enumerate(ex.map(_fetch_book, book_paths)):
+                    out.write(block)
+                    chars += len(block)
+                    if (i + 1) % 50 == 0 or i + 1 == len(book_paths):
+                        done = (i + 1) / len(book_paths)
+                        elapsed = time.time() - t0
+                        eta = elapsed * (1 - done) / done if done > 0 else 0.0
+                        print(
+                            f"\r  {i + 1}/{len(book_paths)} books · "
+                            f"{chars / 1e9:.2f} GB · "
+                            f"{elapsed:4.0f}s elapsed · ETA {eta:4.0f}s",
+                            end="",
+                            flush=True,
+                        )
         print()
 
         with open(tmp, "r", encoding="utf-8") as f:

export_onnx.py

@@ -42,7 +42,7 @@ def main() -> None:
     args = ap.parse_args()
 
     print(f"loading {args.ckpt}")
-    model, chars, hp = load_model_from_checkpoint(args.ckpt, device="cpu")
+    model, tokenizer, hp = load_model_from_checkpoint(args.ckpt, device="cpu")
     wrapped = InferenceWrapper(model)
     wrapped.eval()
 
@@ -63,15 +63,27 @@ def main() -> None:
     size_mb = os.path.getsize(args.out) / 1e6
     print(f"  wrote {args.out} ({size_mb:.1f} MB)")
 
+    # One decoded string per token ID. For char-level each entry is a single
+    # character; for BPE entries are byte-decoded substrings (possibly empty
+    # / multi-byte). The JS frontend just renders `tokens[id]` after each
+    # sample, so both tokenizers display correctly.
+    tokens = [tokenizer.decode([i]) for i in range(tokenizer.vocab_size)]
     meta: dict[str, object] = {
-        "chars": chars,
+        "tokens": tokens,
+        "tokenizer": tokenizer.name,
         "block_size": hp.block_size,
-        "vocab_size": len(chars),
+        "vocab_size": tokenizer.vocab_size,
         "n_layer": hp.n_layer,
         "n_embd": hp.n_embd,
         "n_head": hp.n_head,
         "checkpoint": os.path.basename(args.ckpt),
     }
+    # Char-level: keep the legacy `chars` field so the existing JS frontend
+    # (which builds its stoi map from it for prompt encoding) keeps working
+    # without changes. BPE prompt encoding would need a JS BPE encoder; out
+    # of scope for this change.
+    if tokenizer.name == "char":
+        meta["chars"] = tokens
     with open(args.vocab_out, "w", encoding="utf-8") as f:
         json.dump(meta, f, indent=2)
     print(f"  wrote {args.vocab_out}")