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)

src/dataset.py

@@ -0,0 +1,47 @@
+"""Memory-mapped token dataset and random-crop batch sampler.
+
+Tokens live on disk as a flat array of uint16 (our vocab is < 65536, so two
+bytes each). We memory-map the file so the multi-hundred-MB token array is
+never loaded into RAM as Python objects — the OS pages in only the slices each
+batch actually touches.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+# Vocab < 65536, so every token id fits in an unsigned 16-bit integer.
+TOKEN_DTYPE = np.uint16
+
+
+def load_tokens(path: str | Path) -> np.memmap:
+    """Memory-map a .bin token file as a read-only 1-D array (never loaded whole)."""
+    return np.memmap(path, dtype=TOKEN_DTYPE, mode="r")
+
+
+def get_batch(
+    data: np.ndarray,
+    batch_size: int,
+    context_length: int,
+    device: str,
+    generator: torch.Generator | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sample a batch of random crops from the token stream.
+
+    Returns (x, y), each shape (batch_size, context_length), dtype long.
+    y is x shifted left by one: the target at position t is the token that
+    should follow x[:t+1]. A crop of context_length + 1 tokens gives both.
+    """
+    # Highest start index that still leaves room for context_length + 1 tokens.
+    hi = len(data) - context_length - 1
+    starts = torch.randint(0, hi, (batch_size,), generator=generator).tolist()
+    x = torch.empty((batch_size, context_length), dtype=torch.long)
+    y = torch.empty((batch_size, context_length), dtype=torch.long)
+    for row, s in enumerate(starts):
+        crop = data[s : s + context_length + 1].astype(np.int64)
+        x[row] = torch.from_numpy(crop[:-1])
+        y[row] = torch.from_numpy(crop[1:])
+    return x.to(device), y.to(device)