| """ |
| 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 |
|
|
| |
|
|
| 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 |
|
|
| 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()) |
| 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 |
|
|
|
|
| |
| _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) |
|
|
|
|
| |
| 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 |
| |
| 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) |
| |
| 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") |
|
|
|
|
| |
| 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() |
|
|
