layer_0.py

"""Layer 0: the binary input.

No parameters. Owns the u8 id streams (train/val) and turns them into
model-ready batches: x (B,T) and y (B,T) int64 tensors, where y is the
same stream shifted by one id (next-token target).

Independence: imports only numpy/torch. Knows nothing about layer 1+.
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np
import torch


@dataclass
class Layer0Config:
    block_size: int = 1024       # T: context length handed to layer 1
    batch_size: int = 32        # B
    chunk_size: int = 65536    # shuffle granularity, in ids
    seed: Optional[int] = 1337


class Layer0:
    """Read-only view of the binary; chunk-shuffle sampling."""

    def __init__(self, cfg, train_path="data/train.bin", val_path="data/val.bin"):
        self.cfg = cfg
        # memmaps: fork-inheritable later, no copies in this process now
        self.train = np.memmap(train_path, dtype=np.uint8, mode="r")
        self.val = np.memmap(val_path, dtype=np.uint8, mode="r")
        self.n_train = int(self.train.size)
        self.n_val = int(self.val.size)
        self.n_chunks = self.n_train // cfg.chunk_size   # tail partial chunk dropped
        self.rng = np.random.default_rng(cfg.seed)

    def sample_batch(self):
        """One training batch: random chunk, random offset, x and y shifted by 1."""
        B, T, C = self.cfg.batch_size, self.cfg.block_size, self.cfg.chunk_size
        xs = np.empty((B, T + 1), dtype=np.uint8)
        for b in range(B):
            c = int(self.rng.integers(0, self.n_chunks))
            o = int(self.rng.integers(0, C - T))          # keeps y inside the chunk
            base = c * C + o
            xs[b] = self.train[base: base + T + 1]
        x = torch.tensor(xs[:, :T], dtype=torch.long)
        y = torch.tensor(xs[:, 1:], dtype=torch.long)
        return x, y

    def val_block(self, i):
        """The i-th contiguous block of the val stream (sequential, no shuffle)."""
        T = self.cfg.block_size
        base = i * T
        assert base + T < self.n_val, f"val block {i} out of range"
        seg = self.val[base: base + T + 1]
        return (torch.tensor(seg[:T], dtype=torch.long),
                torch.tensor(seg[1:], dtype=torch.long))

    @property
    def n_val_blocks(self):
        # blocks i with base+T+1 <= n_val  ->  i <= (n_val - T - 1) // T
        return (self.n_val - self.cfg.block_size - 1) // self.cfg.block_size + 1

    @property
    def steps_per_epoch(self):
        return self.n_train // (self.cfg.batch_size * self.cfg.block_size)

    # ---- stateful truncated BPTT mode: sequential windows over contiguous tapes ----

    @property
    def tape_len(self):
        """Contiguous ids per row, truncated to whole windows."""
        return (self.n_train // self.cfg.batch_size) // self.cfg.block_size * self.cfg.block_size

    @property
    def n_seq_windows(self):
        return self.tape_len // self.cfg.block_size

    def seq_windows(self):
        """Yield (x, y) for each T-window, all B rows advancing in lockstep.

        Row b walks tape b = train[b*tape_len : (b+1)*tape_len] in order.
        Tapes are contiguous in the original stream, so the caller may carry
        RNN state across windows and reset only at sweep end (the one real
        discontinuity). Deterministic: no RNG.
        """
        B, T = self.cfg.batch_size, self.cfg.block_size
        rows = np.arange(B, dtype=np.int64) * self.tape_len
        cols = np.arange(T + 1, dtype=np.int64)
        idx = rows[:, None] + cols[None, :]     # (B, T+1) base indices, built once
        for w in range(self.n_seq_windows):
            off = (idx + w * T).ravel()         # per-window offsets
            xs = self.train[off].reshape(B, T + 1)   # one C-level gather, no python loop
            yield (torch.tensor(xs[:, :T], dtype=torch.long),
                   torch.tensor(xs[:, 1:], dtype=torch.long))


def _self_test():
    PROJ = Path(__file__).resolve().parents[1]
    import sys
    sys.path.insert(0, str(PROJ / "scripts"))
    import tokenizer as tk
    tk.load_vocab()

    cfg = Layer0Config(batch_size=8, block_size=256)
    l0 = Layer0(cfg)
    x, y = l0.sample_batch()

    assert x.shape == (8, 256) and y.shape == (8, 256), "batch shape"
    assert int(x.min()) >= 0 and int(x.max()) <= 108, "ids out of vocab"
    assert torch.equal(y[:, :255], x[:, 1:]), "shift law violated"

    v0, _ = l0.val_block(0)
    assert torch.equal(v0, torch.tensor(l0.val[:256].astype(np.int64))), "val walk"

    def render(row):
        out = bytes(row.numpy().astype(np.uint8)).translate(tk.TABLE_DEC)
        for s, t in tk.SENT_DEC:
            out = out.replace(s, t)
        return tk.display(out).decode()

    print("x[0] (model input):")
    print(render(x[0]))
    print("y[0] (model target):")
    print(render(y[0]))
    print(f"n_train={l0.n_train:,}  n_chunks={l0.n_chunks:,}  "
          f"steps/epoch={l0.steps_per_epoch:,}  n_val_blocks={l0.n_val_blocks:,}")
    print("layer0 self-test: PASS")


if __name__ == "__main__":
    _self_test()