Phase 2: token dataset + corpus encoder (code only, not yet run)

- src/dataset.py: memmap uint16 token loader + random-crop batch sampler
  (x, y shifted-by-one, reproducible via generator)
- scripts/encode_corpus.py: stream stories -> encode -> uint16 .bin, one EOT
  token appended per story, flat memory. encode_file() factored out for testing.
- tests/test_dataset.py: batch shapes/shift/bounds/reproducibility +
  encoder round-trip/separators/uint16 on synthetic data (5 tests)

scripts/encode_corpus.py

@@ -0,0 +1,96 @@
+"""Encode the raw corpus into flat uint16 token files.
+
+Each story in the source is delimited by a line containing only
+``<|endoftext|>``. We stream the file one line at a time, encode each story
+with the trained tokenizer, append the end-of-text token id after it, and flush
+the ids to disk in batches — so memory stays flat regardless of corpus size.
+
+Output: data/tokens/train.bin and data/tokens/val.bin.
+
+Usage:
+    python scripts/encode_corpus.py
+"""
+
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+from dataset import TOKEN_DTYPE  # noqa: E402
+from tokenizer import BPETokenizer  # noqa: E402
+
+EOT = "<|endoftext|>"
+TOK_PATH = ROOT / "artifacts" / "tokenizer.json"
+RAW = ROOT / "data" / "raw"
+OUT = ROOT / "data" / "tokens"
+FLUSH_EVERY = 20_000  # stories between disk writes / progress updates
+
+JOBS = [
+    ("TinyStoriesV2-GPT4-train.txt", "train.bin"),
+    ("TinyStoriesV2-GPT4-valid.txt", "val.bin"),
+]
+
+
+def encode_file(tok: BPETokenizer, src: Path, dst: Path, progress: bool = True) -> tuple[int, int]:
+    """Encode one corpus file to a uint16 .bin. Returns (n_stories, n_tokens)."""
+    eot_id = tok.special_tokens[EOT]
+    buf: list[int] = []
+    story: list[str] = []
+    n_stories = n_tokens = 0
+    t0 = time.perf_counter()
+
+    def emit(text: str) -> None:
+        nonlocal n_stories, n_tokens
+        text = text.strip()
+        if not text:
+            return  # skip blank runs between separators
+        ids = tok.encode(text)
+        ids.append(eot_id)
+        buf.extend(ids)
+        n_stories += 1
+        n_tokens += len(ids)
+
+    with open(src, "r", encoding="utf-8", errors="ignore") as fin, open(dst, "wb") as fout:
+        for line in fin:
+            if line.strip() == EOT:
+                emit("".join(story))
+                story = []
+                if n_stories % FLUSH_EVERY == 0 and buf:
+                    np.asarray(buf, dtype=TOKEN_DTYPE).tofile(fout)
+                    buf.clear()
+                    if progress:
+                        rate = n_stories / (time.perf_counter() - t0)
+                        print(f"\r  {dst.name}: {n_stories:,} stories, "
+                              f"{n_tokens:,} tokens ({rate:,.0f}/s)", end="", flush=True)
+            else:
+                story.append(line)
+        emit("".join(story))  # trailing story with no final separator
+        if buf:
+            np.asarray(buf, dtype=TOKEN_DTYPE).tofile(fout)
+    if progress:
+        print()
+    return n_stories, n_tokens
+
+
+def main() -> None:
+    if not TOK_PATH.exists():
+        sys.exit(f"missing {TOK_PATH} — run scripts/train_tokenizer.py first")
+    tok = BPETokenizer.load(TOK_PATH)
+    OUT.mkdir(parents=True, exist_ok=True)
+
+    for src_name, dst_name in JOBS:
+        src, dst = RAW / src_name, OUT / dst_name
+        if not src.exists():
+            sys.exit(f"missing {src} — run scripts/download_data.py first")
+        print(f"encoding {src_name} -> {dst_name}")
+        n_stories, n_tokens = encode_file(tok, src, dst)
+        size_mb = dst.stat().st_size / 1e6
+        print(f"  done: {n_stories:,} stories, {n_tokens:,} tokens, {size_mb:,.0f} MB")
+
+
+if __name__ == "__main__":
+    main()