File size: 5,101 Bytes
b615ad6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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()
|