| """Byte-level data pipeline (TinyStories -> uint8 memmaps). |
| |
| TinyStories is stored as UTF-8 bytes (no BPE): ``V = 256``. Stories are |
| separated by ``0xFF`` which can never occur inside valid UTF-8, so it is a safe |
| end-of-text marker. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
| from typing import Any, Iterator, Mapping, Optional, Tuple |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
| EOT = 255 |
|
|
|
|
| |
| |
| |
| class ByteDataset(Dataset): |
| """Random-access token blocks backed by a memmap. |
| |
| ``dtype=np.uint8`` is the byte-level corpus (V=256); ``np.uint16`` is the |
| BPE corpus (V<=65535). The class name is historical. |
| """ |
|
|
| def __init__(self, path: str | Path, seq_len: int, |
| dtype: Any = np.uint8) -> None: |
| self.path = Path(path) |
| if not self.path.exists(): |
| raise FileNotFoundError( |
| f"{self.path} not found -- run `python scripts/prepare_data.py` first") |
| self.seq_len = int(seq_len) |
| self.dtype = np.dtype(dtype) |
| self.data = np.memmap(self.path, dtype=self.dtype, mode="r") |
| if len(self.data) < self.seq_len + 1: |
| raise ValueError(f"{self.path} is too small ({len(self.data)} tokens)") |
|
|
| def __len__(self) -> int: |
| return max(1, (len(self.data) - 1) // self.seq_len) |
|
|
| def __getitem__(self, i: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| start = (i * self.seq_len) % (len(self.data) - self.seq_len - 1) |
| chunk = np.asarray(self.data[start:start + self.seq_len + 1], dtype=np.int64) |
| x = torch.from_numpy(chunk[:-1]) |
| y = torch.from_numpy(chunk[1:]) |
| return x, y |
|
|
| def get_batch(self, batch_size: int, device: torch.device | str = "cpu", |
| generator: Optional[torch.Generator] = None |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| n = len(self.data) |
| ix = torch.randint(0, n - self.seq_len - 1, (batch_size,), generator=generator) |
| x = torch.stack([torch.from_numpy( |
| np.asarray(self.data[i:i + self.seq_len], dtype=np.int64)) for i in ix]) |
| y = torch.stack([torch.from_numpy( |
| np.asarray(self.data[i + 1:i + 1 + self.seq_len], dtype=np.int64)) for i in ix]) |
| return x.to(device, non_blocking=True), y.to(device, non_blocking=True) |
|
|
| def iter_blocks(self, batch_size: int, device: torch.device | str = "cpu" |
| ) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]: |
| """Deterministic sequential pass (used for validation).""" |
| n_blocks = len(self) |
| for start in range(0, n_blocks, batch_size): |
| idxs = range(start, min(start + batch_size, n_blocks)) |
| xs, ys = [], [] |
| for b in idxs: |
| x, y = self[b] |
| xs.append(x) |
| ys.append(y) |
| yield torch.stack(xs).to(device), torch.stack(ys).to(device) |
|
|
|
|
| def make_synthetic_bin(path: str | Path, n_bytes: int = 1 << 20, |
| seed: int = 0) -> Path: |
| """Random bytes with a learnable first-order structure (smoke tests).""" |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| rng = np.random.default_rng(seed) |
| alphabet = np.arange(32, 127, dtype=np.uint8) |
| data = rng.choice(alphabet, size=n_bytes) |
| |
| for i in range(1, n_bytes): |
| if rng.random() < 0.6: |
| data[i] = min(255, data[i - 1] + 1) |
| data.tofile(path) |
| return path |
|
|
|
|
| |
| |
| |
| SEP = b"\xff" |
|
|
|
|
| def dataset_paths(cfg: Mapping[str, Any], data_dir: str | Path | None = None |
| ) -> dict: |
| """Resolve corpus paths + dtype from the tokenizer choice. |
| |
| ``bpe`` -> ``data/{train,val}.bpe.bin`` as ``uint16`` |
| ``byte`` -> ``data/{train,val}.bin`` as ``uint8`` |
| """ |
| dc = cfg["data"] |
| base = Path(data_dir or dc.get("data_dir", "data")) |
| kind = str(dc.get("tokenizer", "byte")).lower() |
| if kind in ("byte", "bytes", "byte_level"): |
| return {"train": base / dc.get("train_bin", "train.bin"), |
| "val": base / dc.get("val_bin", "val.bin"), |
| "dtype": np.uint8, "tokenizer": "byte"} |
| if kind == "bpe": |
| return {"train": base / dc.get("train_bin", "train.bpe.bin"), |
| "val": base / dc.get("val_bin", "val.bpe.bin"), |
| "dtype": np.uint16, "tokenizer": "bpe"} |
| raise ValueError(f"unknown tokenizer {kind!r}") |
|
|
|
|
| def open_dataset(cfg: Mapping[str, Any], split: str = "train", |
| data_dir: str | Path | None = None, |
| seq_len: int | None = None) -> "ByteDataset": |
| paths = dataset_paths(cfg, data_dir) |
| n = int(seq_len if seq_len is not None else cfg["data"]["seq_len"]) |
| return ByteDataset(paths[split], n, dtype=paths["dtype"]) |
|
|
|
|
| def _encode_texts(texts, eot: int = EOT, budget: Optional[int] = None) -> bytes: |
| sep = bytes([eot]) |
| out = bytearray() |
| for text in texts: |
| out += text.encode("utf-8", errors="ignore") |
| out += sep |
| if budget is not None and len(out) >= budget: |
| break |
| return bytes(out) |
|
|
|
|
| def prepare_tinystories(cfg: Mapping[str, Any]) -> dict: |
| """Download/encode TinyStories into ``{data_dir}/train.bin`` and ``val.bin``.""" |
| from datasets import load_dataset |
|
|
| dc = cfg["data"] |
| data_dir = Path(dc.get("data_dir", "data")) |
| data_dir.mkdir(parents=True, exist_ok=True) |
| eot = int(dc.get("eot_id", EOT)) |
| max_stories = dc.get("max_stories", 200_000) |
| max_val = dc.get("max_val_stories", 10_000) |
|
|
| train_path = data_dir / "train.bin" |
| val_path = data_dir / "val.bin" |
| meta: dict = {"dataset": dc["dataset"], "eot": eot} |
|
|
| if train_path.exists() and val_path.exists() and not dc.get("force", False): |
| meta["train_bytes"] = train_path.stat().st_size |
| meta["val_bytes"] = val_path.stat().st_size |
| meta["cached"] = True |
| return meta |
|
|
| ds = load_dataset(dc["dataset"]) |
| tr_key = "train" if "train" in ds else list(ds.keys())[0] |
| va_key = "validation" if "validation" in ds else ( |
| "valid" if "valid" in ds else None) |
| del ds |
|
|
| n_tr = int(max_stories) if max_stories else None |
| tr_split = f"{tr_key}[:{n_tr}]" if n_tr else tr_key |
| print(f"[data] encoding train split {tr_split} ...") |
| tr_ds = load_dataset(dc["dataset"], split=tr_split) |
| tr_bytes = _encode_texts(tr_ds["text"], eot) |
| train_path.write_bytes(tr_bytes) |
| del tr_ds |
|
|
| if va_key is not None: |
| n_va = int(max_val) if max_val else None |
| va_split = f"{va_key}[:{n_va}]" if n_va else va_key |
| print(f"[data] encoding validation split {va_split} ...") |
| va_ds = load_dataset(dc["dataset"], split=va_split) |
| va_bytes = _encode_texts(va_ds["text"], eot) |
| del va_ds |
| else: |
| n = len(tr_bytes) |
| split_at = int(n * (1.0 - float(dc.get("val_split", 0.05)))) |
| va_bytes = tr_bytes[split_at:] |
| tr_bytes = tr_bytes[:split_at] |
| train_path.write_bytes(tr_bytes) |
| val_path.write_bytes(va_bytes) |
|
|
| meta["train_bytes"] = len(tr_bytes) |
| meta["val_bytes"] = len(va_bytes) |
| meta["train_stories"] = n_tr |
| meta["cached"] = False |
| return meta |
|
|
