File size: 4,435 Bytes
767a051 | 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 | """Deterministic label-independent train/validation/holdout assignment."""
from __future__ import annotations
import hashlib
from typing import Any, Mapping
from tom_world03.canonical import attach_hash, canonical_bytes, require_hash
from .model import ObservationSet, SPLITS
SPLIT_CERTIFICATE_SCHEMA = "TOM-LEARNER-SPLIT-CERTIFICATE-0.1"
def _allocation(dataset: ObservationSet) -> dict[str, int]:
n = len(dataset.observations)
minima = dict(dataset.split_policy.minimum_counts)
minimum_total = sum(minima.values())
if n < minimum_total:
raise ValueError(
f"observation count {n} is below split minimum total {minimum_total}"
)
remaining = n - minimum_total
ratios = dict(dataset.split_policy.ratios)
total_ratio = sum(ratios.values())
floors: dict[str, int] = {}
remainders: dict[str, int] = {}
for name in SPLITS:
numerator = remaining * ratios[name]
floors[name], remainders[name] = divmod(numerator, total_ratio)
leftover = remaining - sum(floors.values())
# The split-name order is a normative final tie break.
remainder_order = sorted(
SPLITS,
key=lambda name: (-remainders[name], SPLITS.index(name)),
)
extras = {name: 0 for name in SPLITS}
for name in remainder_order[:leftover]:
extras[name] += 1
counts = {
name: minima[name] + floors[name] + extras[name]
for name in SPLITS
}
if sum(counts.values()) != n:
raise AssertionError("split allocation did not consume every observation")
return counts
def assignment_digest(dataset: ObservationSet, observation_id: str) -> str:
"""Digest only identity and policy, never ``t`` or ``y`` values."""
basis = {
"seed_sha256": dataset.seed_sha256,
"observation_set_id": dataset.id,
"split_policy_hash": dataset.split_policy.content_hash,
"salt": dataset.split_policy.salt,
"observation_id": observation_id,
}
return hashlib.sha256(canonical_bytes(basis)).hexdigest()
def deterministic_split(dataset: ObservationSet) -> dict[str, Any]:
counts = _allocation(dataset)
ordered = sorted(
(
assignment_digest(dataset, observation.id),
observation.id,
)
for observation in dataset.observations
)
assignments: list[dict[str, str]] = []
split_ids: dict[str, list[str]] = {name: [] for name in SPLITS}
cursor = 0
for name in SPLITS:
stop = cursor + counts[name]
for digest, observation_id in ordered[cursor:stop]:
assignments.append({
"observation_id": observation_id,
"assignment_digest": digest,
"split": name,
})
split_ids[name].append(observation_id)
cursor = stop
assignments.sort(key=lambda item: item["observation_id"])
for ids in split_ids.values():
ids.sort()
basis = {
"strategy": dataset.split_policy.strategy,
"seed_sha256": dataset.seed_sha256,
"observation_set_id": dataset.id,
"split_policy_hash": dataset.split_policy.content_hash,
"salt": dataset.split_policy.salt,
"observation_ids": sorted(observation.id for observation in dataset.observations),
}
return attach_hash({
"schema": SPLIT_CERTIFICATE_SCHEMA,
"profile": "TOM-LEARNER-0.1",
"observation_set_id": dataset.id,
"observation_set_hash": dataset.content_hash,
"split_policy_hash": dataset.split_policy.content_hash,
"assignment_basis_hash": "sha256:" + hashlib.sha256(canonical_bytes(basis)).hexdigest(),
"assignment_uses_values": False,
"counts": counts,
"assignments": assignments,
"splits": split_ids,
})
def validate_split_certificate(dataset: ObservationSet, certificate: Mapping[str, Any]) -> None:
require_hash(certificate, label="split certificate")
if certificate.get("schema") != SPLIT_CERTIFICATE_SCHEMA:
raise ValueError("unsupported split certificate schema")
if certificate.get("observation_set_hash") != dataset.content_hash:
raise ValueError("split certificate is bound to another observation set")
expected = deterministic_split(dataset)
if canonical_bytes(expected) != canonical_bytes(certificate):
raise ValueError("split certificate does not reproduce deterministically")
|