Cormac Shannon
29e323b2fc
nanoGPT gets a qk_norm flag (off = upstream behaviour), a seed option and a resolved-config dump. run.py logs each run's git commit, config, data checksums, environment and metrics, and stops diverged runs. Claude-Session: https://claude.ai/code/session_01S8KN7NuUCm8nXLmy6eDuVv
102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Download TinyStories (Eldan & Li 2023, V2 / GPT-4 split) and tokenise it for nanoGPT.
|
|
|
|
Usage (from 06-experiments/):
|
|
.venv/bin/python data/prepare_tinystories.py # full: ~2.2 GB train text
|
|
.venv/bin/python data/prepare_tinystories.py --spike # small: valid file only, split 90/10
|
|
|
|
Writes data/<name>/{train,val}.bin (uint16 GPT-2 BPE tokens, nanoGPT format) and
|
|
data/<name>/prep.json (source URLs, SHA-256 of the raw files, token counts), so every
|
|
run can record exactly which data it used. Raw text and .bin files are git-ignored.
|
|
"""
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sys
|
|
import urllib.request
|
|
|
|
import numpy as np
|
|
import tiktoken
|
|
|
|
BASE = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/"
|
|
TRAIN, VALID = "TinyStoriesV2-GPT4-train.txt", "TinyStoriesV2-GPT4-valid.txt"
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
SEP = "<|endoftext|>"
|
|
|
|
|
|
def download(fname, raw_dir):
|
|
path = os.path.join(raw_dir, fname)
|
|
if not os.path.exists(path):
|
|
print(f"downloading {fname} ...")
|
|
tmp = path + ".part"
|
|
urllib.request.urlretrieve(BASE + fname, tmp)
|
|
os.rename(tmp, path)
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
return path, h.hexdigest()
|
|
|
|
|
|
def stories(path):
|
|
"""Yield one story at a time (stories are separated by <|endoftext|>)."""
|
|
buf = []
|
|
with open(path, encoding="utf-8") as f:
|
|
for line in f:
|
|
if SEP in line:
|
|
head, _, tail = line.partition(SEP)
|
|
buf.append(head)
|
|
s = "".join(buf).strip()
|
|
if s:
|
|
yield s
|
|
buf = [tail]
|
|
else:
|
|
buf.append(line)
|
|
s = "".join(buf).strip()
|
|
if s:
|
|
yield s
|
|
|
|
|
|
def write_bin(texts, out_path, enc):
|
|
eot = enc.eot_token
|
|
ids = []
|
|
for t in texts:
|
|
ids.extend(enc.encode_ordinary(t))
|
|
ids.append(eot)
|
|
arr = np.array(ids, dtype=np.uint16)
|
|
arr.tofile(out_path)
|
|
return len(arr)
|
|
|
|
|
|
def main():
|
|
spike = "--spike" in sys.argv
|
|
name = "tinystories_spike" if spike else "tinystories"
|
|
out_dir = os.path.join(HERE, name)
|
|
raw_dir = os.path.join(HERE, "raw")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
os.makedirs(raw_dir, exist_ok=True)
|
|
enc = tiktoken.get_encoding("gpt2")
|
|
meta = {"dataset": "TinyStories V2 (GPT-4)", "source": BASE, "tokenizer": "tiktoken gpt2", "files": {}}
|
|
|
|
vpath, vsha = download(VALID, raw_dir)
|
|
meta["files"][VALID] = vsha
|
|
if spike:
|
|
allv = list(stories(vpath))
|
|
cut = int(len(allv) * 0.9)
|
|
meta["split"] = f"valid file split 90/10 by story order: {cut} train / {len(allv) - cut} val stories"
|
|
n_tr = write_bin(allv[:cut], os.path.join(out_dir, "train.bin"), enc)
|
|
n_va = write_bin(allv[cut:], os.path.join(out_dir, "val.bin"), enc)
|
|
else:
|
|
tpath, tsha = download(TRAIN, raw_dir)
|
|
meta["files"][TRAIN] = tsha
|
|
meta["split"] = "official train / valid files"
|
|
n_tr = write_bin(stories(tpath), os.path.join(out_dir, "train.bin"), enc)
|
|
n_va = write_bin(stories(vpath), os.path.join(out_dir, "val.bin"), enc)
|
|
meta["tokens"] = {"train": n_tr, "val": n_va}
|
|
json.dump(meta, open(os.path.join(out_dir, "prep.json"), "w"), indent=2)
|
|
print(json.dumps(meta, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|