File size: 6,067 Bytes
3c85459 | 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 135 136 137 138 139 140 141 142 | """Prepare frozen TinyStories streams for both training arms."""
import hashlib
from itertools import chain
import json
from pathlib import Path
import tempfile
from urllib.request import Request, urlopen
import numpy as np
import sentencepiece as spm
ROOT = Path(__file__).resolve().parent
CONFIG_PATH = ROOT / "protocol.json"
CONFIG = json.loads(CONFIG_PATH.read_bytes())
BASELINE = ROOT.with_name("stories260k-baseline")
REFERENCE = json.loads((BASELINE / "protocol.json").read_bytes())
TOKENIZER = BASELINE / "assets/tok512.model"
VALIDATION = BASELINE / "assets/TinyStories-valid.txt"
DELIMITER = b"<|endoftext|>"
DATA = ROOT / "data"
TOKENS = DATA / "train.u16"
MANIFEST = DATA / "manifest.json"
BLOCK = 1024 * 1024
def sha256(path):
with path.open("rb") as stream:
return hashlib.file_digest(stream, "sha256").hexdigest()
def baseline_asset(path):
expected = REFERENCE["assets"][path.relative_to(BASELINE).as_posix()]["sha256"]
if sha256(path) != expected:
raise ValueError(f"Baseline input hash differs for {path.name}")
return expected
def atomic_json(path, value):
temp = path.with_suffix(".tmp")
temp.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n", encoding="utf-8")
temp.replace(path)
def stories(stream, digest=None):
pending = b""
while chunk := stream.read(BLOCK):
if digest is not None:
digest.update(chunk)
records = (pending + chunk).split(DELIMITER)
pending = records.pop()
for record in records:
yield record.strip()
def corpus_tokens(config):
return config.get("data_tokens", config["training_tokens"])
def verify(config=CONFIG, directory=DATA):
manifest = json.loads((directory / MANIFEST.name).read_bytes())
tokens = directory / TOKENS.name
if (manifest["training_source"] != config["training_source"]
or manifest["requested_bytes"] != config["training_source_limit_bytes"]
or manifest.get("source_sha256") != config.get("training_source_sha256")
or manifest["tokenizer_sha256"] != baseline_asset(TOKENIZER)
or manifest["validation_sha256"] != baseline_asset(VALIDATION)
or manifest["token_count"] != corpus_tokens(config) + 1
or tokens.stat().st_size != 2 * manifest["token_count"]
or sha256(tokens) != manifest["tokens_sha256"]):
raise ValueError("Prepared training data does not match its protocol and manifest")
return manifest
def prepare(config=CONFIG, directory=DATA):
manifest_path = directory / MANIFEST.name
if manifest_path.exists():
return verify(config, directory)
tokenizer_sha = baseline_asset(TOKENIZER)
validation_sha = baseline_asset(VALIDATION)
tokenizer = spm.SentencePieceProcessor(model_file=str(TOKENIZER))
if tokenizer.vocab_size() != REFERENCE["architecture"]["vocab_size"]:
raise ValueError("Unexpected vocabulary size")
excluded = {hashlib.sha256(row.strip()).digest() for row in VALIDATION.read_bytes().split(DELIMITER)}
source_sha = config.get("training_source_sha256")
source_digest = hashlib.sha256() if source_sha else None
needed = corpus_tokens(config) + 1
limit = config["training_source_limit_bytes"]
headers = {} if source_sha else {"Range": f"bytes=0-{limit - 1}"}
request = Request(config["training_source"], headers=headers)
directory.mkdir(exist_ok=True)
with tempfile.NamedTemporaryFile(dir=directory, delete=False) as output:
temp = Path(output.name)
try:
written = kept = skipped = 0
with urlopen(request, timeout=60) as stream, temp.open("wb") as output:
if source_sha:
if stream.status != 200 or stream.headers.get("Content-Length") != str(limit):
raise ValueError("Expected the entire pinned training source")
elif stream.status != 206 or not stream.headers.get("Content-Range", "").startswith(f"bytes 0-{limit - 1}/"):
raise ValueError("Training source did not honor the bounded range request")
batch = []
for story in chain(stories(stream, source_digest), (None,)):
if story:
if hashlib.sha256(story).digest() in excluded:
skipped += 1
else:
batch.append(story.decode("utf-8"))
if len(batch) < 256 and story is not None:
continue
for ids in tokenizer.encode(batch, add_bos=True, add_eos=True):
if not source_sha:
ids = ids[:needed - written]
output.write(np.asarray(ids, dtype="<u2").tobytes())
written += len(ids)
kept += 1
if not source_sha and written == needed:
break
if not source_sha and written == needed:
break
batch.clear()
if source_sha:
if source_digest.hexdigest() != source_sha:
raise ValueError("Full training source hash differs from the protocol")
context = config["context"]
written = (written - 1) // context * context + 1
output.truncate(written * 2)
if written != needed:
raise ValueError(f"Training source yielded {written:,} tokens, expected {needed:,}")
manifest = {"training_source": config["training_source"], "requested_bytes": limit,
"tokenizer_sha256": tokenizer_sha, "validation_sha256": validation_sha,
"tokens_sha256": sha256(temp), "token_count": written, "stories_used": kept,
"exact_validation_duplicates_skipped": skipped}
if source_sha:
manifest["source_sha256"] = source_sha
temp.replace(directory / TOKENS.name)
atomic_json(manifest_path, manifest)
return manifest
finally:
temp.unlink(missing_ok=True)
|