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