src/pretraining/data.py

"""
Memmap-backed token dataset for next-token prediction pretraining.

Why memmap?
===========
A pretraining corpus is one long sequence of token ids. The naive
approach (load the whole thing into a Python list and slice) wastes
memory and is slow to copy onto the GPU. The standard pattern, used
by nanoGPT and basically every educational pretraining repo, is:

1. Tokenize once at preparation time and write the result to a binary
   file as raw uint16 little-endian (for vocab <= 65535).
2. At training time, ``np.memmap`` the file. The OS handles paging,
   so the dataset behaves as if fully in RAM but only the slices you
   actually touch are read from disk.
3. For each training step, sample ``batch_size`` random offsets into
   the memmap and read ``block_size + 1`` tokens at each one. Split
   into ``x = tokens[:-1]`` (input) and ``y = tokens[1:]`` (target)
   for next-token prediction.

This is fast (sequential reads), trivial to implement (~30 lines),
scales to corpora that don't fit in RAM, and has zero
preprocessing-at-train-time overhead.

uint16 vs int32
===============
uint16 covers vocab sizes up to 65,535, which is ample for our 8k
custom BPE. If you ever need a larger vocab, switch to uint32 — but
the file size doubles, and you almost certainly don't need a larger
vocab for from-scratch tiny pretraining (the whole point is keeping
the embedding small).
"""

from pathlib import Path
from typing import Protocol, Tuple, Union, runtime_checkable

import numpy as np
import torch


@runtime_checkable
class TokenSource(Protocol):
    """Anything the trainer can pull next-token-prediction batches from.

    Implementing this is the *only* thing required to plug a new corpus
    into the pretraining loop. Concretely, ``get_batch`` must return a
    pair of ``torch.long`` tensors of shape ``(batch_size, block_size)``
    such that ``y[:, t] == x[:, t+1]`` for some implicit underlying token
    stream — i.e. the standard shifted-by-one next-token target.

    The trainer treats sources as opaque: no iteration order, no epoch
    boundaries, no shuffling state. It just calls ``get_batch`` once per
    train step and once per eval step. If you need a streaming HF
    dataset, a multi-file shard, or anything fancier, hide it behind
    this method.

    The optional ``generator`` argument is used to make train batch
    sampling reproducible across runs; eval calls pass ``None`` and
    rely on the global RNG, which is fine because eval is a noisy
    average over many batches anyway.
    """

    def get_batch(
        self,
        batch_size: int,
        device: Union[str, torch.device] = "cpu",
        generator: Union[torch.Generator, None] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        ...


# Token dtype for the on-disk binary. Must match what
# scripts/prepare_tinystories.py writes.
TOKEN_DTYPE = np.uint16


class TokenDataset:
    """
    Random-access view over a memmapped token binary.

    Not a ``torch.utils.data.Dataset`` — for from-scratch pretraining
    we don't need iteration order, multi-worker shuffling, or any of
    the DataLoader machinery. ``get_batch`` directly samples random
    offsets, which is simpler and faster.

    Args:
        bin_path: Path to a uint16 binary written by the preparation
            script.
        block_size: Number of *input* tokens per sample. Each sample
            actually reads ``block_size + 1`` tokens so the target can
            be the same window shifted by one.
    """

    def __init__(self, bin_path: Union[str, Path], block_size: int):
        self.bin_path = Path(bin_path)
        self.block_size = block_size

        if not self.bin_path.exists():
            raise FileNotFoundError(
                f"Token binary not found at {self.bin_path}. Run "
                f"scripts/prepare_tinystories.py first."
            )

        # mode='r' = read-only memmap. The OS will lazily page in
        # whatever bytes we touch.
        self.tokens = np.memmap(self.bin_path, dtype=TOKEN_DTYPE, mode="r")

        if len(self.tokens) < block_size + 1:
            raise ValueError(
                f"Token binary at {self.bin_path} has only {len(self.tokens)} "
                f"tokens, but block_size+1 = {block_size + 1} are required "
                f"for a single sample."
            )

    def __len__(self) -> int:
        """Number of valid starting positions for a (block_size+1)-token window."""
        return len(self.tokens) - self.block_size

    def get_batch(
        self,
        batch_size: int,
        device: Union[str, torch.device] = "cpu",
        generator: Union[torch.Generator, None] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Sample a random batch of (input, target) sequences.

        Args:
            batch_size: Number of independent sequences in the batch.
            device: Device to copy the tensors to. For MPS / CUDA the
                copy is asynchronous and overlaps with the next forward.
            generator: Optional ``torch.Generator`` for reproducible
                sampling. If ``None``, uses the default RNG.

        Returns:
            ``(x, y)`` where each is shape ``(batch_size, block_size)``
            of dtype ``torch.long``. ``y`` is ``x`` shifted by one
            position to the right (the standard next-token target).
        """
        max_start = len(self) - 1  # inclusive max for randint
        # torch.randint is faster than np.random.randint for our needs
        # and integrates with torch.Generator for reproducibility.
        starts = torch.randint(
            low=0,
            high=max_start + 1,
            size=(batch_size,),
            generator=generator,
        )

        # Build the batch on CPU as int64 (the dtype embeddings expect),
        # then ship to the target device in one transfer.
        x = torch.empty((batch_size, self.block_size), dtype=torch.long)
        y = torch.empty((batch_size, self.block_size), dtype=torch.long)
        for i, start in enumerate(starts.tolist()):
            window = self.tokens[start : start + self.block_size + 1]
            # .astype(np.int64) realizes the slice into RAM (the memmap
            # slice is a uint16 view); from_numpy then shares the buffer.
            window = window.astype(np.int64, copy=False)
            x[i] = torch.from_numpy(window[:-1])
            y[i] = torch.from_numpy(window[1:])

        if str(device) != "cpu":
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
        return x, y


def write_tokens_bin(tokens: np.ndarray, output_path: Union[str, Path]) -> None:
    """
    Persist a 1-D token id array to disk in the format ``TokenDataset``
    expects.

    Args:
        tokens: 1-D numpy array of token ids. Will be cast to
            ``TOKEN_DTYPE`` — caller is responsible for ensuring all
            ids fit (i.e. ``tokens.max() < 65536``).
        output_path: Destination ``.bin`` file. Parent directory is
            created if needed.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if tokens.ndim != 1:
        raise ValueError(f"Expected 1-D token array, got shape {tokens.shape}")
    if tokens.size > 0 and tokens.max() >= np.iinfo(TOKEN_DTYPE).max + 1:
        raise ValueError(
            f"Token id {int(tokens.max())} does not fit in {TOKEN_DTYPE}. "
            f"Use a smaller vocabulary or switch to uint32."
        )

    tokens.astype(TOKEN_DTYPE, copy=False).tofile(output_path)