data.py

"""
Stage 1 — DATA (lecture: Data).

What it does:
  1. Fetch TinyStories (simple children's stories — small but clean, a good fit for small models).
  2. Clean & dedup (a mini version of the Common Crawl filtering pipeline from the lecture).
  3. Train the BPE tokenizer.
  4. Encode every story -> uint16 array on disk (memmap), separated by EOS tokens.
  5. Hold out separate stories (no overlap with pretraining) for the SFT dataset later.

Run:  python data.py
"""
from __future__ import annotations

import hashlib
import json
import os
import random
import time
from multiprocessing import Pool

import numpy as np

from config import get_config
from tokenizer import BPETokenizer

# ---------------------------------------------------------------- sources

def _synthetic_stories(n: int, seed: int) -> list[str]:
    """Template stories for the offline smoke test. Not for real training."""
    rng = random.Random(seed)
    names = ["Lily", "Tom", "Mia", "Ben", "Sara", "Max", "Anna", "Leo", "Zoe", "Sam"]
    animals = ["cat", "dog", "bird", "frog", "bunny", "duck", "fox", "bear"]
    things = ["ball", "kite", "box", "hat", "boat", "flower", "cake", "book", "drum", "shell"]
    places = ["park", "garden", "forest", "beach", "house", "school", "river", "farm"]
    feels = ["happy", "sad", "scared", "excited", "tired", "proud", "curious", "angry"]
    acts = ["played", "jumped", "danced", "looked", "ran", "laughed", "searched", "climbed"]
    out = []
    for _ in range(n):
        a, b = rng.sample(names, 2)
        an, th, pl, fe, ac = (rng.choice(x) for x in (animals, things, places, feels, acts))
        th2 = rng.choice(things)
        s = [
            f"Once upon a time, there was a little {an} named {a}.",
            f"{a} lived near a big {pl}.",
            f"One day, {a} found a {th} in the {pl}.",
            f"{a} felt {fe} and {ac} with the {th}.",
            f"Then {b} came and saw the {th}.",
            f'"Can I play too?" asked {b}.',
            f'"Yes!" said {a}. They shared the {th} and a {th2}.',
            f"At the end of the day, {a} and {b} were good friends.",
        ]
        rng.shuffle(s[1:7])
        out.append(" ".join(s))
    return out


def load_stories(cfg, total: int) -> list[str]:
    if cfg.data.synthetic:
        return _synthetic_stories(total, seed=0)
    from datasets import load_dataset  # imported here: not required for the smoke test

    ds = load_dataset(cfg.data.dataset, split="train", streaming=True)
    out = []
    for row in ds:
        out.append(row["text"])
        if len(out) >= total:
            break
    return out


def clean_and_dedup(stories: list[str]) -> list[str]:
    """Heuristic filtering: drop too short/long stories, excessive non-ASCII, and exact duplicates."""
    seen, out = set(), []
    for s in stories:
        s = " ".join(s.split())  # normalize whitespace
        if not (100 <= len(s) <= 4000):
            continue
        if sum(c.isascii() for c in s) / len(s) < 0.98:
            continue
        h = hashlib.md5(s.lower().encode()).hexdigest()
        if h in seen:
            continue
        seen.add(h)
        out.append(s)
    return out


# --------------------------------------------------------------- encoding
_TOK: BPETokenizer | None = None


def _init_worker(path: str) -> None:
    global _TOK
    _TOK = BPETokenizer.load(path)


def _encode_batch(stories: list[str]) -> list[int]:
    ids: list[int] = []
    for s in stories:
        ids.extend(_TOK.encode(s, allow_special=False))
        ids.append(_TOK.eos_id)
    return ids


def encode_to_bin(stories: list[str], tok_path: str, out_path: str, workers: int) -> int:
    chunks = [stories[i:i + 2000] for i in range(0, len(stories), 2000)]
    if workers > 1 and len(chunks) > 1:
        with Pool(workers, initializer=_init_worker, initargs=(tok_path,)) as pool:
            parts = pool.map(_encode_batch, chunks)
    else:
        _init_worker(tok_path)
        parts = [_encode_batch(c) for c in chunks]
    arr = np.fromiter((t for p in parts for t in p), dtype=np.uint16)
    arr.tofile(out_path)
    return len(arr)


# ------------------------------------------------------------------- main
def main() -> None:
    cfg = get_config()
    d = cfg.data
    os.makedirs(d.data_dir, exist_ok=True)
    n_sft = cfg.sft.n_examples + cfg.sft.n_val
    # request more than needed because some are dropped by cleaning + the SFT length filter
    total_needed = d.n_train_stories + d.n_val_stories + int(n_sft * 2.5)

    t0 = time.time()
    raw = load_stories(cfg, int(total_needed * 1.05))
    stories = clean_and_dedup(raw)
    print(f"Stories: {len(raw):,} raw -> {len(stories):,} after cleaning/dedup ({time.time() - t0:.0f}s)")

    train = stories[: d.n_train_stories]
    val = stories[d.n_train_stories: d.n_train_stories + d.n_val_stories]
    sft_pool = stories[d.n_train_stories + d.n_val_stories:]

    tok_path = os.path.join(d.data_dir, "tokenizer.json")
    if not os.path.exists(tok_path):
        text, n = [], 0
        for s in train:
            text.append(s)
            n += len(s)
            if n >= d.tokenizer_train_chars:
                break
        t0 = time.time()
        tok = BPETokenizer.train("\n".join(text), d.vocab_size)
        tok.save(tok_path)
        print(f"Tokenizer trained in {time.time() - t0:.0f}s")
    tok = BPETokenizer.load(tok_path)
    # A small corpus can run out of pairs before reaching the target vocab -> spare embedding rows go unused
    assert tok.vocab_size <= cfg.model.vocab_size, (
        f"tokenizer vocab ({tok.vocab_size}) > model.vocab_size ({cfg.model.vocab_size})"
    )

    sample = train[0][:120]
    ids = tok.encode(sample)
    print(f"Example: {len(sample)} chars -> {len(ids)} tokens | roundtrip OK: {tok.decode(ids) == sample}")

    workers = max(1, (os.cpu_count() or 2) - 1)
    for split, items in (("train", train), ("val", val)):
        t0 = time.time()
        n = encode_to_bin(items, tok_path, os.path.join(d.data_dir, f"{split}.bin"), workers)
        print(f"{split}.bin: {n:,} tokens ({time.time() - t0:.0f}s)")

    with open(os.path.join(d.data_dir, "sft_stories.jsonl"), "w") as f:
        for s in sft_pool:
            f.write(json.dumps({"text": s}) + "\n")
    print(f"sft_stories.jsonl: {len(sft_pool):,} held-out stories for SFT/DPO")


# --------------------------------------------------------------- batching
class TokenDataset:
    """Sample random context_len windows from the token file (memmap, RAM-friendly)."""

    def __init__(self, path: str, context_len: int):
        self.data = np.memmap(path, dtype=np.uint16, mode="r")
        self.T = context_len

    def __len__(self) -> int:
        return len(self.data)

    def batch(self, batch_size: int, device: str, rng: np.random.Generator):
        import torch

        ix = rng.integers(0, len(self.data) - self.T - 1, size=batch_size)
        x = np.stack([self.data[i: i + self.T] for i in ix]).astype(np.int64)
        y = np.stack([self.data[i + 1: i + 1 + self.T] for i in ix]).astype(np.int64)
        x, y = torch.from_numpy(x), torch.from_numpy(y)
        if device == "cuda":
            return x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
        return x.to(device), y.to(device)


if __name__ == "__main__":
    main()