Blame for scripts/make_eval_split.py

"""Create and validate a deterministic, dataset-bound train/validation/test split.

The split is shared by LatticeSpike and faithful baselines. It holds
``train_indices``, optional ``validation_indices``, and ``eval_indices`` (the
untouched test set) as offsets into the exact bound evaluation JSONL.

Identical normalized prompts stay in one partition to prevent question leakage.
Explicit v5.1/v5.2 split groups take precedence when present. Groups are
stratified by source and whether they contain hallucinated tokens.

Example:
    python scripts/make_eval_split.py \
        --input outputs/llama32_1b_uncertain_eval.jsonl \
        --output outputs/llama32_1b_uncertain_split.json \
        --train-fraction 0.75 --validation-fraction 0.15 --test-fraction 0.10 \
        --seed 42
"""

from __future__ import annotations

import argparse
from collections import Counter
import hashlib
import json
import math
from pathlib import Path
import random

import numpy as np
import re


SPLIT_BINDING_SCHEMA_VERSION = 1
SPLIT_SCHEME_PREFIX = "grouped_stratified_holdout test_fraction="
# Corpora that ship their own train/test boundary (RAGTruth) must keep it, or
# their numbers stop being comparable to every published result on them. Such a
# split cannot match the grouped/stratified builder, so it gets its own scheme
# and its own deterministic builder -- the validator still RECOMPUTES it from
# the data rather than trusting the stored indices.
NATIVE_SPLIT_SCHEME = "native_eval_grouped_validation_v1"
# Same builder, but train/validation are restricted to ONE domain while the
# eval set stays whole. That is the train/test/transfer design: fit on one
# distribution, then read in-domain and out-of-domain performance off the same
# evaluation pass by slicing predictions afterwards.
NATIVE_DOMAIN_SPLIT_SCHEME = "native_eval_domain_restricted_train_v1"


def sha256_json(value) -> str:
    encoded = json.dumps(
        value, sort_keys=True, separators=(",", ":"), ensure_ascii=False
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def dataset_binding(rows: list[dict]) -> dict:
    ids: list[str] = []
    row_sha256: list[str] = []
    seen: set[str] = set()
    for index, row in enumerate(rows):
        row_id = row.get("id")
        if not isinstance(row_id, str) or not row_id.strip():
            raise ValueError(f"row {index} requires a non-empty string id")
        if row_id in seen:
            raise ValueError(f"duplicate row id: {row_id!r}")
        seen.add(row_id)
        ids.append(row_id)
        row_sha256.append(sha256_json(row))
    ordered_rows = [
        {"id": row_id, "content_sha256": content_sha256}
        for row_id, content_sha256 in zip(ids, row_sha256)
    ]
    return {
        "schema_version": SPLIT_BINDING_SCHEMA_VERSION,
        "row_count": len(rows),
        "ordered_ids": ids,
        "row_content_sha256": row_sha256,
        "ordered_ids_sha256": sha256_json(ids),
        "ordered_rows_sha256": sha256_json(ordered_rows),
    }


def has_hallucinated(labels) -> bool:
    for label in labels or []:
        if isinstance(label, str):
            if label.strip().upper() in {"1", "H", "HALLUCINATED"}:
                return True
        elif int(label) == 1:
            return True
    return False


def source_name(row: dict) -> str:
    metadata = row.get("metadata") or {}
    return str(metadata.get("source") or row.get("source") or "unknown")


def normalized_prompt(value) -> str:
    return re.sub(r"\s+", " ", str(value or "")).strip().casefold()


def row_split_group(row: dict, index: int) -> str:
    metadata = row.get("metadata")
    extra = metadata.get("extra") if isinstance(metadata, dict) else None
    if not isinstance(extra, dict):
        extra = row.get("extra") if isinstance(row.get("extra"), dict) else None
    v5_1 = extra.get("v5_1") if isinstance(extra, dict) else None
    explicit = v5_1.get("split_group") if isinstance(v5_1, dict) else None
    if explicit:
        return str(explicit)
    normalized = normalized_prompt(row.get("prompt", row.get("question")))
    if normalized:
        return f"prompt:{normalized}"
    record_id = row.get("id")
    if record_id is not None:
        return f"id:{record_id}"
    return f"row:{index}"


def example_group_key(row: dict, index: int) -> str:
    return row_split_group(row, index)


def parse_args(argv=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument(
        "--input",
        type=Path,
        required=True,
        help="Raw auto-labelled or normalized token-evaluation JSONL.",
    )
    parser.add_argument(
        "--output", type=Path, required=True, help="Destination split JSON."
    )
    parser.add_argument(
        "--test-fraction",
        type=float,
        default=0.10,
        help="Fraction of examples held out as the untouched test set.",
    )
    parser.add_argument(
        "--validation-fraction",
        type=float,
        default=0.15,
        help="Fraction of examples held out for checkpoint/threshold selection.",
    )
    parser.add_argument(
        "--train-fraction",
        type=float,
        default=None,
        help="Fraction of examples used for model fitting; defaults to the remainder.",
    )
    parser.add_argument("--seed", type=int, default=42)
    return parser.parse_args(argv)


def _validate_fractions(
    train_fraction: float | None,
    validation_fraction: float,
    test_fraction: float,
) -> float:
    if not 0.0 < test_fraction < 1.0:
        raise ValueError("--test-fraction must be in (0, 1)")
    if validation_fraction < 0.0 or validation_fraction >= 1.0:
        raise ValueError("--validation-fraction must be in [0, 1)")
    if train_fraction is None:
        train_fraction = 1.0 - validation_fraction - test_fraction
    if train_fraction <= 0.0:
        raise ValueError("--train-fraction must be positive")
    total = train_fraction + validation_fraction + test_fraction
    if abs(total - 1.0) > 1e-6:
        raise ValueError("train/validation/test fractions must sum to 1")
    return train_fraction


def _partition_counts(
    n_groups: int,
    train_fraction: float,
    validation_fraction: float,
    test_fraction: float,
) -> dict[str, int]:
    fractions = {
        "train": train_fraction,
        "validation": validation_fraction,
        "test": test_fraction,
    }
    counts = {
        name: int(n_groups * fraction) for name, fraction in fractions.items()
    }
    remainder = n_groups - sum(counts.values())
    order = sorted(
        fractions,
        key=lambda name: (
            n_groups * fractions[name] - counts[name],
            fractions[name],
        ),
        reverse=True,
    )
    for name in order[:remainder]:
        counts[name] += 1
    for name, fraction in fractions.items():
        if fraction > 0.0 and counts[name] == 0:
            donor = max(
                (candidate for candidate in counts if counts[candidate] > 1),
                key=counts.get,
                default=None,
            )
            if donor is None:
                raise ValueError("not enough prompt groups for the requested partitions")
            counts[donor] -= 1
            counts[name] += 1
    return counts


def _choose_strata(
    group_ids: list[str],
    group_source: dict[str, str],
    group_has_h: dict[str, bool],
    holdout_count: int,
) -> tuple[list[str] | None, str]:
    candidates = [
        (
            [
                f"{group_source[group_id]}|has_H={group_has_h[group_id]}"
                for group_id in group_ids
            ],
            "source+has_H",
        ),
        ([group_source[group_id] for group_id in group_ids], "source"),
        ([str(group_has_h[group_id]) for group_id in group_ids], "has_H"),
    ]
    remaining_count = len(group_ids) - holdout_count
    for labels, name in candidates:
        counts = Counter(labels)
        if (
            min(counts.values()) >= 2
            and holdout_count >= len(counts)
            and remaining_count >= len(counts)
        ):
            return labels, name
    return None, "none"


def _split_groups(
    group_ids: list[str],
    holdout_count: int,
    seed: int,
    group_source: dict[str, str],
    group_has_h: dict[str, bool],
) -> tuple[list[str], list[str], str]:
    if holdout_count <= 0:
        return list(group_ids), [], "none"
    strata, stratified_by = _choose_strata(
        group_ids, group_source, group_has_h, holdout_count
    )
    try:
        from sklearn.model_selection import train_test_split

        kept, held_out = train_test_split(
            group_ids,
            test_size=holdout_count,
            random_state=seed,
            shuffle=True,
            stratify=strata,
        )
    except ImportError:  # dependency-light local validation only
        rng = random.Random(seed)
        if strata is None:
            shuffled = list(group_ids)
            rng.shuffle(shuffled)
            held_out = shuffled[:holdout_count]
        else:
            buckets: dict[str, list[str]] = {}
            for group_id, stratum in zip(group_ids, strata):
                buckets.setdefault(str(stratum), []).append(group_id)
            for bucket in buckets.values():
                rng.shuffle(bucket)
            quotas = {
                name: holdout_count * len(bucket) / len(group_ids)
                for name, bucket in buckets.items()
            }
            counts = {name: int(math.floor(value)) for name, value in quotas.items()}
            remaining = holdout_count - sum(counts.values())
            order = sorted(
                buckets,
                key=lambda name: (quotas[name] - counts[name], name),
                reverse=True,
            )
            for name in order[:remaining]:
                counts[name] += 1
            held_out = [
                group_id
                for name in sorted(buckets)
                for group_id in buckets[name][: counts[name]]
            ]
        held_out_set = set(held_out)
        kept = [group_id for group_id in group_ids if group_id not in held_out_set]
    return list(kept), list(held_out), stratified_by


def _repair_row_counts(
    partitions: dict[str, list[str]],
    groups: dict[str, list[int]],
    target_rows: dict[str, int],
    group_source: dict[str, str],
    group_has_h: dict[str, bool],
    fractions: dict[str, float],
) -> None:
    """Move whole prompt groups to hit requested row counts when feasible."""
    row_counts = {
        name: sum(len(groups[group_id]) for group_id in group_ids)
        for name, group_ids in partitions.items()
    }
    stratum_totals = Counter(
        (group_source[group_id], group_has_h[group_id])
        for group_id in groups
        for _ in groups[group_id]
    )
    stratum_counts = {
        name: Counter(
            (group_source[group_id], group_has_h[group_id])
            for group_id in group_ids
            for _ in groups[group_id]
        )
        for name, group_ids in partitions.items()
    }

    while row_counts != target_rows:
        donors = [
            name for name in partitions if row_counts[name] > target_rows[name]
        ]
        receivers = [
            name for name in partitions if row_counts[name] < target_rows[name]
        ]
        candidates = []
        for donor in donors:
            excess = row_counts[donor] - target_rows[donor]
            for receiver in receivers:
                deficit = target_rows[receiver] - row_counts[receiver]
                for group_id in partitions[donor]:
                    size = len(groups[group_id])
                    if size > excess or size > deficit:
                        continue
                    stratum = (group_source[group_id], group_has_h[group_id])
                    stratum_size = size
                    donor_target = stratum_totals[stratum] * fractions[donor]
                    receiver_target = stratum_totals[stratum] * fractions[receiver]
                    before = abs(stratum_counts[donor][stratum] - donor_target) + abs(
                        stratum_counts[receiver][stratum] - receiver_target
                    )
                    after = abs(
                        stratum_counts[donor][stratum]
                        - stratum_size
                        - donor_target
                    ) + abs(
                        stratum_counts[receiver][stratum]
                        + stratum_size
                        - receiver_target
                    )
                    candidates.append((after - before, group_id, donor, receiver))
        if not candidates:
            break
        _, group_id, donor, receiver = min(candidates)
        partitions[donor].remove(group_id)
        partitions[receiver].append(group_id)
        size = len(groups[group_id])
        stratum = (group_source[group_id], group_has_h[group_id])
        row_counts[donor] -= size
        row_counts[receiver] += size
        stratum_counts[donor][stratum] -= size
        stratum_counts[receiver][stratum] += size


def _dotted(row: dict, path: str):
    value = row
    for part in path.split("."):
        if not isinstance(value, dict):
            return None
        value = value.get(part)
    return value


def build_native_split(
    rows: list[dict],
    seed: int,
    validation_fraction: float,
    native_split_field: str,
    eval_value: str,
    grouping_key: str,
    train_domain_field: str | None = None,
    train_domain_value: str | None = None,
) -> dict:
    """Preserve a corpus's own eval split; carve validation out of its train.

    Deterministic in (rows, seed, validation_fraction, field names) alone, so
    `validate_split_summary` can rebuild and compare it exactly the way it does
    for the grouped/stratified builder. Nothing here reads stored indices.
    """
    eval_indices, groups = [], {}
    for index, row in enumerate(rows):
        if str(_dotted(row, native_split_field)) == eval_value:
            # The eval set is never filtered: transfer is measured by slicing
            # these predictions by domain afterwards, not by dropping rows.
            eval_indices.append(index)
            continue
        if train_domain_field is not None:
            if str(_dotted(row, train_domain_field)) != str(train_domain_value):
                continue
        groups.setdefault(str(_dotted(row, grouping_key)), []).append(index)
    if not eval_indices:
        raise ValueError(f"no rows have {native_split_field} == {eval_value!r}")
    if not groups:
        raise ValueError("native split left no rows for train/validation")

    group_ids = sorted(groups)
    generator = np.random.default_rng(seed)
    n_val = max(
        1, min(len(group_ids) - 1, round(len(group_ids) * float(validation_fraction)))
    )
    val_groups = {group_ids[i] for i in generator.permutation(len(group_ids))[:n_val]}

    train_indices, validation_indices = [], []
    for gid, members in groups.items():
        (validation_indices if gid in val_groups else train_indices).extend(members)
    return {
        "split": {
            "train_indices": sorted(train_indices),
            "validation_indices": sorted(validation_indices),
            "eval_indices": sorted(eval_indices),
        }
    }


def build_split(
    rows: list[dict],
    test_fraction: float,
    seed: int,
    validation_fraction: float = 0.0,
    train_fraction: float | None = None,
) -> dict:
    binding = dataset_binding(rows)
    train_fraction = _validate_fractions(
        train_fraction, validation_fraction, test_fraction
    )
    n = len(rows)
    if n < 2:
        raise ValueError("need at least two examples to split")
    groups: dict[str, list[int]] = {}
    for index, row in enumerate(rows):
        groups.setdefault(example_group_key(row, index), []).append(index)
    group_ids = sorted(groups)
    if len(group_ids) < 2:
        raise ValueError("need at least two distinct prompt groups to split")
    group_source = {
        group_id: source_name(rows[indices[0]])
        for group_id, indices in groups.items()
    }
    group_has_h = {
        group_id: any(
            has_hallucinated(
                rows[index].get("labels") or rows[index].get("token_labels")
            )
            for index in indices
        )
        for group_id, indices in groups.items()
    }
    group_counts = _partition_counts(
        len(group_ids), train_fraction, validation_fraction, test_fraction
    )
    train_validation_groups, test_groups, test_strata = _split_groups(
        group_ids,
        group_counts["test"],
        seed,
        group_source,
        group_has_h,
    )
    train_groups, validation_groups, validation_strata = _split_groups(
        train_validation_groups,
        group_counts["validation"],
        seed + 1,
        group_source,
        group_has_h,
    )
    fractions = {
        "train": train_fraction,
        "validation": validation_fraction,
        "test": test_fraction,
    }
    target_rows = _partition_counts(
        n, train_fraction, validation_fraction, test_fraction
    )
    partitions = {
        "train": train_groups,
        "validation": validation_groups,
        "test": test_groups,
    }
    _repair_row_counts(
        partitions,
        groups,
        target_rows,
        group_source,
        group_has_h,
        fractions,
    )
    train_groups = partitions["train"]
    validation_groups = partitions["validation"]
    test_groups = partitions["test"]

    def flatten(selected_groups: list[str]) -> list[int]:
        return sorted(
            index for group_id in selected_groups for index in groups[group_id]
        )

    train_indices = flatten(train_groups)
    validation_indices = flatten(validation_groups)
    eval_indices = flatten(test_groups)
    if not train_indices or not eval_indices:
        raise ValueError("split produced an empty partition; adjust --test-fraction")
    if validation_fraction > 0.0 and not validation_indices:
        raise ValueError(
            "split produced an empty validation partition; adjust --validation-fraction"
        )
    split = {"train_indices": train_indices, "eval_indices": eval_indices}
    if validation_indices:
        split["validation_indices"] = validation_indices
    return {
        "split": split,
        "n": n,
        "n_groups": len(groups),
        "groups": len(groups),
        "seed": seed,
        "test_fraction": float(test_fraction),
        "validation_fraction": float(validation_fraction),
        "train_fraction": float(train_fraction),
        "dataset_binding": binding,
        "grouped_by": "v5_1.split_group or normalized prompt",
        "stratified_by": {
            "test": test_strata,
            "validation": validation_strata if validation_indices else None,
        },
        "fractions": fractions,
        "target_partition_sizes": target_rows,
        "actual_partition_sizes": {
            "train": len(train_indices),
            "validation": len(validation_indices),
            "test": len(eval_indices),
        },
        "scheme": (
            "grouped_stratified_train_validation_test"
            if validation_indices
            else f"grouped_stratified_holdout test_fraction={test_fraction}"
        ),
    }


def split_test_fraction(summary: dict) -> float:
    raw = summary.get("test_fraction")
    if raw is None and isinstance(summary.get("fractions"), dict):
        raw = summary["fractions"].get("test")
    if raw is None:
        scheme = summary.get("scheme")
        if isinstance(scheme, str) and scheme.startswith(SPLIT_SCHEME_PREFIX):
            raw = scheme[len(SPLIT_SCHEME_PREFIX) :]
    try:
        value = float(raw)
    except (TypeError, ValueError) as exc:
        raise ValueError("split is missing a valid test_fraction") from exc
    if not math.isfinite(value) or not 0.0 < value < 1.0:
        raise ValueError("split test_fraction must be in (0, 1)")
    return value


def validate_split_summary(
    summary: dict, input_rows: list[dict]
) -> dict[str, list[int]]:
    split = summary.get("split")
    if not isinstance(split, dict):
        raise ValueError("split summary does not contain a split object")
    for key in ("train_indices", "eval_indices"):
        if not isinstance(split.get(key), list):
            raise ValueError(f"split.{key} must be a list")

    expected_binding = dataset_binding(input_rows)
    actual_binding = summary.get("dataset_binding")
    if not isinstance(actual_binding, dict):
        raise ValueError(
            "split is not bound to its input dataset; rebuild it with "
            "scripts/make_eval_split.py"
        )
    for key, expected in expected_binding.items():
        if actual_binding.get(key) != expected:
            if key == "row_content_sha256":
                raise ValueError("split dataset content is stale")
            if key == "ordered_ids":
                raise ValueError("split dataset IDs are stale or reordered")
            raise ValueError(f"split dataset binding has stale or invalid {key}")
    if summary.get("n") != len(input_rows):
        raise ValueError("split row count does not match its input dataset")

    names = ("train_indices", "validation_indices", "eval_indices")
    parsed: dict[str, list[int]] = {}
    owner: dict[int, str] = {}
    for name in names:
        values = split.get(name, [])
        if not isinstance(values, list):
            raise ValueError(f"split.{name} must be a list")
        if any(not isinstance(value, int) or isinstance(value, bool) for value in values):
            raise ValueError(f"split.{name} must contain only integer indices")
        if len(values) != len(set(values)):
            raise ValueError(f"split.{name} contains duplicate indices")
        for value in values:
            if not 0 <= value < len(input_rows):
                raise ValueError(f"split.{name} contains out-of-range index {value}")
            if value in owner:
                raise ValueError(f"split partitions overlap at index {value}")
            owner[value] = name
        parsed[name] = list(values)
    if set(owner) != set(range(len(input_rows))):
        # A domain-restricted split deliberately omits other domains' training
        # rows. That is the only sanctioned reason to leave a row unassigned,
        # and it is verified rather than assumed: every omitted row must be a
        # non-eval row from a different domain. Any other gap is still an error.
        native = summary.get("native", {}) if isinstance(summary.get("native"), dict) else {}
        field = native.get("train_domain_field")
        value = native.get("train_domain_value")
        if summary.get("scheme") != NATIVE_DOMAIN_SPLIT_SCHEME or field is None:
            raise ValueError("split partitions must cover every input row exactly once")
        for index in set(range(len(input_rows))) - set(owner):
            row = input_rows[index]
            if str(_dotted(row, native["split_field"])) == native["eval_value"]:
                raise ValueError(
                    f"domain-restricted split dropped eval row {index}; the "
                    "evaluation set must stay whole so transfer is measurable"
                )
            if str(_dotted(row, field)) == str(value):
                raise ValueError(
                    f"domain-restricted split dropped row {index} from its own "
                    f"training domain {value!r}"
                )
    if not parsed["train_indices"] or not parsed["eval_indices"]:
        raise ValueError("train and untouched-test partitions must both be non-empty")

    group_owner: dict[str, str] = {}
    for name, values in parsed.items():
        for index in values:
            group = row_split_group(input_rows[index], index)
            prior = group_owner.setdefault(group, name)
            if prior != name:
                raise ValueError(f"duplicate split group {group!r} crosses partitions")

    seed = summary.get("seed")
    if not isinstance(seed, int) or isinstance(seed, bool):
        raise ValueError("split seed must be an integer")
    fractions = summary.get("fractions")
    if isinstance(fractions, dict):
        train_fraction = fractions.get("train")
        validation_fraction = fractions.get("validation", 0.0)
        test_fraction = fractions.get("test")
    else:
        test_fraction = split_test_fraction(summary)
        validation_fraction = summary.get("validation_fraction", 0.0)
        train_fraction = summary.get("train_fraction")
    if summary.get("scheme") in (NATIVE_SPLIT_SCHEME, NATIVE_DOMAIN_SPLIT_SCHEME):
        native = summary.get("native", {})
        expected = build_native_split(
            input_rows,
            seed=seed,
            validation_fraction=float(summary.get("validation_fraction", 0.0)),
            native_split_field=native["split_field"],
            eval_value=native["eval_value"],
            grouping_key=native["grouping_key"],
            train_domain_field=native.get("train_domain_field"),
            train_domain_value=native.get("train_domain_value"),
        )
        for name in names:
            if parsed[name] != expected["split"].get(name, []):
                raise ValueError(
                    "split indices do not match the native-split builder for "
                    f"seed={seed}"
                )
        return {name: values for name, values in parsed.items() if name in split}
    expected = build_split(
        input_rows,
        test_fraction=float(test_fraction),
        seed=seed,
        validation_fraction=float(validation_fraction),
        train_fraction=None if train_fraction is None else float(train_fraction),
    )
    for name in names:
        if parsed[name] != expected["split"].get(name, []):
            raise ValueError(
                "split indices do not match the deterministic grouped/stratified "
                f"builder for seed={seed}"
            )
    return {
        name: values
        for name, values in parsed.items()
        if name in split
    }


def load_bound_split(path: Path, input_rows: list[dict]) -> dict[str, list[int]]:
    summary = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(summary, dict):
        raise ValueError(f"{path} does not contain a JSON object")
    return validate_split_summary(summary, input_rows)


def main(argv=None) -> None:
    args = parse_args(argv)
    # NOT .splitlines(): that also breaks on \u2028/\u2029/\x85 etc., and a
    # generated answer containing one of those (Qwen3 emits NEL, \x85) shears
    # the JSON line mid-string. JSONL is newline-delimited; split on newlines.
    with args.input.open(encoding="utf-8") as handle:
        rows = [json.loads(line) for line in handle if line.strip()]
    split = build_split(
        rows,
        args.test_fraction,
        args.seed,
        validation_fraction=args.validation_fraction,
        train_fraction=args.train_fraction,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(split, indent=2), encoding="utf-8")
    tr = len(split["split"]["train_indices"])
    va = len(split["split"].get("validation_indices", []))
    ev = len(split["split"]["eval_indices"])
    middle = f" / {va} validation" if va else ""
    print(
        f"wrote {args.output}: {tr} train{middle} / {ev} test "
        f"(n={split['n']}, seed={args.seed})"
    )


if __name__ == "__main__":
    main()