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