Prepare examples

#!/usr/bin/env python3
"""Build or verify the Step 2 workbook. Uses only Python's standard library.

Arithmetic here creates/checks labels. This script is not the model's answering
path and must never be used to generate answers for users.
"""

import argparse
import hashlib
import json
from pathlib import Path
import sys


VERSION = "addition-v1"
SEED = 42
MAX_OPERAND = 99
GROUP_COUNTS = {"train": 4040, "validation": 505, "calibration": 252, "test": 253}
OOD_GROUP_COUNT = 250
DEFAULT_OUTPUT = Path(__file__).resolve().parents[1] / "data" / VERSION
FIELDS = {"id", "group_id", "a", "b", "prompt", "answer", "text"}


def group_id(a, b):
    return f"add:{min(a, b)}:{max(a, b)}"


def ordering_key(pair):
    a, b = pair
    message = f"{VERSION}:seed={SEED}:{a}:{b}"
    return hashlib.sha256(message.encode("utf-8")).hexdigest(), a, b


def make_row(a, b):
    prompt = f"{a}+{b}="
    answer = str(a + b)
    return {
        "id": f"add:{a}+{b}",
        "group_id": group_id(a, b),
        "a": a,
        "b": b,
        "prompt": prompt,
        "answer": answer,
        "text": prompt + answer,
    }


def expand_groups(groups):
    pairs = []
    for a, b in groups:
        pairs.append((a, b))
        if a != b:
            pairs.append((b, a))
    return [make_row(a, b) for a, b in sorted(pairs)]


def split_groups():
    groups = sorted(
        ((a, b) for a in range(100) for b in range(a, 100)),
        key=ordering_key,
    )
    result = {}
    start = 0
    for name, count in GROUP_COUNTS.items():
        result[name] = groups[start:start + count]
        start += count
    return result


def ood_groups():
    # At least one operand has three digits; all such prompts are outside 0-99.
    candidates = ((a, b) for a in range(200) for b in range(max(a, 100), 200))
    return sorted(candidates, key=ordering_key)[:OOD_GROUP_COUNT]


def encode_json(value):
    return json.dumps(value, sort_keys=True, ensure_ascii=True)


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def has_carry(a, b):
    carry = 0
    while a or b:
        carry = (a % 10 + b % 10 + carry) // 10
        if carry:
            return True
        a //= 10
        b //= 10
    return False


def summarize(rows):
    return {
        "examples": len(rows),
        "groups": len({row["group_id"] for row in rows}),
        "carry_examples": sum(has_carry(row["a"], row["b"]) for row in rows),
        "no_carry_examples": sum(not has_carry(row["a"], row["b"]) for row in rows),
        "zero_operand_examples": sum(row["a"] == 0 or row["b"] == 0 for row in rows),
        "equal_operand_examples": sum(row["a"] == row["b"] for row in rows),
        "both_single_digit_examples": sum(row["a"] < 10 and row["b"] < 10 for row in rows),
        "answer_lengths": {
            str(length): sum(len(row["answer"]) == length for row in rows)
            for length in (1, 2, 3)
        },
        "longest_text_characters": max(len(row["text"]) for row in rows),
    }


def build_dataset(output):
    output.mkdir(parents=True, exist_ok=True)
    grouped = split_groups()
    grouped["ood"] = ood_groups()
    manifest = {
        "version": VERSION,
        "seed": SEED,
        "operand_range": [0, MAX_OPERAND],
        "group_order": "SHA-256 of addition-v1:seed=42:a:b, with a <= b",
        "splits": {},
    }
    for name, groups in grouped.items():
        rows = expand_groups(groups)
        path = output / f"{name}.jsonl"
        path.write_text("".join(encode_json(row) + "\n" for row in rows), encoding="utf-8")
        manifest["splits"][name] = {
            **summarize(rows),
            "file": path.name,
            "sha256": digest(path),
        }
    (output / "manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    return verify_dataset(output)


def require(condition, message):
    if not condition:
        raise ValueError(message)


def verify_dataset(output):
    manifest = json.loads((output / "manifest.json").read_text(encoding="utf-8"))
    require(manifest["version"] == VERSION, "Unexpected dataset version")
    require(manifest["seed"] == SEED, "Unexpected seed")
    require(manifest["operand_range"] == [0, MAX_OPERAND], "Unexpected operand range")
    require(set(manifest["splits"]) == {*GROUP_COUNTS, "ood"}, "Unexpected split names")
    expected_groups = split_groups()
    expected_groups["ood"] = ood_groups()
    all_ids = set()
    all_groups = set()
    in_range_pairs = set()

    for name in [*GROUP_COUNTS, "ood"]:
        info = manifest["splits"][name]
        require(info["file"] == f"{name}.jsonl", f"{name}: unexpected filename")
        path = output / info["file"]
        require(digest(path) == info["sha256"], f"{name}: file checksum mismatch")
        rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]
        seen_pairs = set()
        groups = set()
        for row in rows:
            require(set(row) == FIELDS, f"{name}: unexpected row fields")
            a, b = row["a"], row["b"]
            require(type(a) is int and type(b) is int, f"{name}: operands must be integers")
            require(0 <= a <= 199 and 0 <= b <= 199, f"{name}: invalid operands")
            require((max(a, b) >= 100) == (name == "ood"), f"{name}: wrong operand range")
            require(row["id"] == f"add:{a}+{b}", f"{name}: wrong example ID")
            require(row["group_id"] == group_id(a, b), f"{name}: wrong group ID")
            require(row["prompt"] == f"{a}+{b}=", f"{name}: wrong prompt format")
            require(row["answer"] == str(a + b), f"{name}: wrong answer")
            require(row["text"] == row["prompt"] + row["answer"], f"{name}: wrong text")
            require(row["id"] not in all_ids, f"{name}: duplicate example")
            all_ids.add(row["id"])
            seen_pairs.add((a, b))
            groups.add(row["group_id"])
            if name != "ood":
                in_range_pairs.add((a, b))
        require(not (groups & all_groups), f"{name}: group leaks across splits")
        all_groups.update(groups)
        require(all((b, a) in seen_pairs for a, b in seen_pairs), f"{name}: missing reversed pair")
        require(groups == {group_id(a, b) for a, b in expected_groups[name]},
                f"{name}: group assignment differs from the fixed split")
        for key, value in summarize(rows).items():
            require(info[key] == value, f"{name}: incorrect manifest statistic {key}")
        require(info["carry_examples"] > 0 and info["no_carry_examples"] > 0,
                f"{name}: missing carry or no-carry examples")
    require(in_range_pairs == {(a, b) for a in range(100) for b in range(100)},
            "The in-range splits must cover all 10,000 ordered pairs exactly once")
    return manifest


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--check", action="store_true", help="Verify saved files without rewriting them")
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="Dataset directory")
    args = parser.parse_args()
    try:
        manifest = verify_dataset(args.output) if args.check else build_dataset(args.output)
    except (OSError, ValueError, KeyError, TypeError) as error:
        print(f"Data check failed: {error}", file=sys.stderr)
        return 1
    print(f"{'Verified' if args.check else 'Built and verified'} {VERSION}: {args.output}")
    for name in [*GROUP_COUNTS, "ood"]:
        info = manifest["splits"][name]
        print(f"  {name}: {info['examples']} examples, {info['groups']} groups, "
              f"{info['carry_examples']} carry examples")
    return 0


if __name__ == "__main__":
    sys.exit(main())