#!/usr/bin/env python3 """Freeze a source-disjoint holdout and a quality-ranked capability selection. This selector never reads benchmark rows, labels, predictions, task IDs, or question templates. It joins an already sealed non-benchmark teacher release back to its generation evidence, removes nominal capability rows whose answer was available too early, and reserves a policy holdout before selecting SFT rows. A final training selection additionally requires a complete student screen over the structural candidate set. """ from __future__ import annotations import argparse import collections import hashlib import json from pathlib import Path from typing import Any, Iterable from viss_core.data_v0.canonical import ( canonical_json_line, file_sha256, read_jsonl, stable_id, ) from viss_core.data_v0.semantic_release import validate_semantic_release SELECTION_FORMAT = "iclr_video.quality_capability_selection.v1" HOLDOUT_FORMAT = "iclr_video.nonbenchmark_policy_holdout.v1" STUDENT_SCREEN_FORMAT = "iclr_video.nonbenchmark_student_screen.v1" PROFILE_KIND = "nonbenchmark_capability_profile" PARENT_KIND = "teacher_generated_nonbenchmark_visual_mcq" PROFILES = ( "backward_epm", "backward_asi", "rec_counting", "sequence_crr", "sequence_ssr", ) DEFAULT_HOLDOUT = {profile: 60 for profile in PROFILES} DEFAULT_TRAIN = {"backward_epm": 120, "backward_asi": 180} class QualityCapabilitySelectionError(ValueError): pass def _target(row: dict[str, Any]) -> str: values = [ str(item.get("text") or "").strip() for message in row.get("messages") or [] if message.get("role") == "assistant" and message.get("loss") is True for item in message.get("content") or [] if item.get("type") == "text" ] if len(values) != 1 or values[0] not in set("ABCD"): raise QualityCapabilitySelectionError("row lacks one A-D target") return values[0] def _single_parent(row: dict[str, Any], kind: str) -> str: values = [ str(parent.get("id") or "") for parent in row.get("parents") or [] if parent.get("kind") == kind ] if len(values) != 1 or not values[0]: raise QualityCapabilitySelectionError(f"row lacks one {kind} parent") return values[0] def _rank(seed: str, profile: str, answer: str, uid: str) -> str: return hashlib.sha256( f"quality-capability-v1:{seed}:{profile}:{answer}:{uid}".encode() ).hexdigest() def _balanced_select( rows: Iterable[dict[str, Any]], *, quota: int, seed: str, profile: str, ) -> list[dict[str, Any]]: if quota <= 0 or quota % 4: raise QualityCapabilitySelectionError("balanced quota must be positive / 4") by_answer: dict[str, list[dict[str, Any]]] = collections.defaultdict(list) for row in rows: by_answer[str(row["answer"])].append(row) per_answer = quota // 4 selected = [] for answer in "ABCD": values = sorted( by_answer.get(answer, []), key=lambda row: (_rank(seed, profile, answer, str(row["uid"])), str(row["uid"])), ) if len(values) < per_answer: raise QualityCapabilitySelectionError( f"{profile}/{answer} has {len(values)} rows; requires {per_answer}" ) selected.extend(values[:per_answer]) return sorted(selected, key=lambda row: str(row["uid"])) def _image_paths(request: dict[str, Any]) -> list[str]: return [ str(image) for group in request.get("groups") or [] for image in group.get("images") or [] ] def _structural_candidate( request: dict[str, Any], generation: dict[str, Any] ) -> tuple[bool, dict[str, Any]]: profile = str(request.get("target_capability") or "") groups = [ tuple(str(image) for image in group.get("images") or []) for group in request.get("groups") or [] ] flattened = [image for group in groups for image in group] evidence_groups = generation.get("evidence_groups") or [] first_sufficient = generation.get("first_sufficient_group") minimum_sufficient = { "backward_epm": 1, "backward_asi": 2, "rec_counting": 1, "sequence_crr": 1, "sequence_ssr": 1, }.get(profile) diagnostics = { "first_sufficient_group": first_sufficient, "evidence_groups": len(evidence_groups), "observation_groups": len(groups), "distinct_observation_groups": len(set(groups)), "image_positions": len(flattened), "unique_images": len(set(flattened)), "duplicate_image_positions": len(flattened) - len(set(flattened)), } eligible = ( minimum_sufficient is not None and generation.get("decision") == "accept" and generation.get("confidence") == "high" and isinstance(first_sufficient, int) and first_sufficient >= minimum_sufficient and len(evidence_groups) >= 2 and len(set(groups)) >= 2 and len(set(flattened)) >= 6 ) return bool(eligible), diagnostics def _student_keep(row: dict[str, Any], audit: dict[str, Any]) -> bool: if audit.get("error") is not None: return False target = str(row["answer"]) if str(audit.get("target") or "") != target: raise QualityCapabilitySelectionError("student audit target diverges") prediction = str(audit.get("prediction") or "") margin = float(audit.get("visible_option_margin")) if not 0.0 <= margin <= 1.0: raise QualityCapabilitySelectionError("student audit margin is outside [0,1]") return prediction != target or margin <= 0.25 def _parse_quota(value: str) -> tuple[str, int]: try: profile, raw = value.rsplit("=", 1) count = int(raw) except (TypeError, ValueError) as exc: raise argparse.ArgumentTypeError("quota must be PROFILE=COUNT") from exc if profile not in PROFILES or count <= 0 or count % 4: raise argparse.ArgumentTypeError("quota must use a known profile and be / 4") return profile, count def _excluded_parent_ids(releases: list[Path]) -> set[str]: excluded: set[str] = set() for raw in releases: root = raw.resolve(strict=True) validate_semantic_release(root, verify_media_hashes=False) for row in read_jsonl( root / "trainer/sft.jsonl", verify_content_hashes=True, require_canonical=True, ): excluded.update( str(parent["id"]) for parent in row.get("parents") or [] if parent.get("kind") == PARENT_KIND and parent.get("id") ) return excluded def build(args: argparse.Namespace) -> tuple[dict[str, Any], dict[str, Any]]: release = args.release.resolve(strict=True) pack = args.pack.resolve(strict=True) generation_path = args.generation.resolve(strict=True) release_manifest = validate_semantic_release(release, verify_media_hashes=False) pack_manifest_path = pack / "manifest.json" pack_manifest = json.loads(pack_manifest_path.read_text(encoding="utf-8")) if ( (release_manifest.get("inputs") or {}).get("pack_manifest_sha256") != file_sha256(pack_manifest_path) or (release_manifest.get("inputs") or {}).get("generation_sha256") != file_sha256(generation_path) ): raise QualityCapabilitySelectionError("release is not bound to pack/generation") requests = read_jsonl(pack / "requests.jsonl", require_canonical=True) generation_rows = read_jsonl(generation_path, require_canonical=True) trainer_rows = read_jsonl( release / "trainer/sft.jsonl", verify_content_hashes=True, require_canonical=True, ) request_by_uid = {str(row["uid"]): row for row in requests} generation_by_uid = {str(row["uid"]): row for row in generation_rows} trainer_by_uid = {str(row["uid"]): row for row in trainer_rows} if len(request_by_uid) != len(requests) or len(generation_by_uid) != len(generation_rows): raise QualityCapabilitySelectionError("request/generation UIDs collide") excluded = _excluded_parent_ids(list(args.exclude_release or [])) joined = [] for request_uid, request in request_by_uid.items(): trainer_uid = stable_id("benchmarkdistillsft", request_uid, digest_length=24) trainer = trainer_by_uid.get(trainer_uid) generation = generation_by_uid.get(request_uid) if trainer is None: continue if generation is None: raise QualityCapabilitySelectionError("filtered row lacks generation record") parent_id = _single_parent(trainer, PARENT_KIND) profile = _single_parent(trainer, PROFILE_KIND) if profile != str(request.get("target_capability") or ""): raise QualityCapabilitySelectionError("request/trainer profile diverges") if parent_id in excluded: continue joined.append( { "uid": trainer_uid, "request_uid": request_uid, "parent_id": parent_id, "profile": profile, "answer": _target(trainer), "source_video_basename": str(request["source_video_basename"]), "request": request, "generation": generation, } ) holdout_quotas = dict(args.holdout_quota or DEFAULT_HOLDOUT) train_quotas = dict(args.train_quota or DEFAULT_TRAIN) holdout_rows = [] for profile, quota in sorted(holdout_quotas.items()): holdout_rows.extend( _balanced_select( (row for row in joined if row["profile"] == profile), quota=quota, seed=args.seed + ":holdout", profile=profile, ) ) holdout_uids = {str(row["uid"]) for row in holdout_rows} holdout_videos = {str(row["source_video_basename"]) for row in holdout_rows} if len(holdout_uids) != len(holdout_rows) or len(holdout_videos) != len(holdout_rows): raise QualityCapabilitySelectionError("policy holdout is not source-video unique") structural = [] for row in joined: if row["uid"] in holdout_uids: continue eligible, diagnostics = _structural_candidate( row["request"], row["generation"] ) if eligible: structural.append({**row, "diagnostics": diagnostics}) audit_by_uid: dict[str, dict[str, Any]] | None = None audit_record = None if args.student_audit is not None: audit_path = args.student_audit.resolve(strict=True) audit_value = json.loads(audit_path.read_text(encoding="utf-8")) if ( audit_value.get("format") != STUDENT_SCREEN_FORMAT or audit_value.get("benchmark_rows_or_labels_used") is not False or audit_value.get("source_release_manifest_sha256") != file_sha256(release / "manifest.json") ): raise QualityCapabilitySelectionError("student screen contract diverges") audit_rows = list(audit_value.get("rows") or []) audit_by_uid = {str(row.get("uid") or ""): row for row in audit_rows} if set(audit_by_uid) != {str(row["uid"]) for row in structural}: raise QualityCapabilitySelectionError("student screen coverage diverges") audit_record = {"path": str(audit_path), "sha256": file_sha256(audit_path)} selected_rows = [] if audit_by_uid is not None: hard = [row for row in structural if _student_keep(row, audit_by_uid[row["uid"]])] for profile, quota in sorted(train_quotas.items()): selected_rows.extend( _balanced_select( (row for row in hard if row["profile"] == profile), quota=quota, seed=args.seed + ":train", profile=profile, ) ) def public_row(row: dict[str, Any]) -> dict[str, Any]: return { key: row[key] for key in ( "uid", "request_uid", "parent_id", "profile", "answer", "source_video_basename", ) } | ({"diagnostics": row["diagnostics"]} if "diagnostics" in row else {}) holdout = { "format": HOLDOUT_FORMAT, "status": "sealed_before_training_selection", "source_release": str(release), "source_release_manifest_sha256": file_sha256(release / "manifest.json"), "selection_seed": args.seed + ":holdout", "quotas": holdout_quotas, "rows": [public_row(row) for row in sorted(holdout_rows, key=lambda row: row["uid"])], "parent_ids": sorted(row["parent_id"] for row in holdout_rows), "contracts": { "source_video_disjoint_from_sft_selection": True, "may_enter_sft_or_rl": False, "benchmark_rows_or_labels_used": False, }, } selection = { "format": SELECTION_FORMAT, "status": ( "sealed_nonbenchmark_quality_selection" if audit_by_uid is not None else "awaiting_complete_student_screen" ), "training_allowed": audit_by_uid is not None, "source_release": str(release), "source_release_manifest_sha256": file_sha256(release / "manifest.json"), "pack_manifest_sha256": file_sha256(pack_manifest_path), "generation_sha256": file_sha256(generation_path), "excluded_parent_ids": len(excluded), "holdout_parent_ids_sha256": hashlib.sha256( canonical_json_line(sorted(holdout["parent_ids"])) ).hexdigest(), "structural_candidates": [ public_row(row) for row in sorted(structural, key=lambda row: row["uid"]) ], "structural_counts_by_profile": dict( sorted(collections.Counter(row["profile"] for row in structural).items()) ), "student_screen": audit_record, "train_quotas": train_quotas, "selected_rows": [ public_row(row) for row in sorted(selected_rows, key=lambda row: row["uid"]) ], "parent_ids": sorted(row["parent_id"] for row in selected_rows), "transformation": "deduplicate_identical_images_keep_first_chronological_occurrence", "contracts": { "benchmark_rows_labels_predictions_or_tasks_used": False, "counterfactual_rows": 0, "holdout_overlap": 0, "already_trained_parent_overlap": 0, "teacher_confidence_high": True, "two_order_option_consensus_inherited": True, }, } return holdout, selection def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--release", type=Path, required=True) parser.add_argument("--pack", type=Path, required=True) parser.add_argument("--generation", type=Path, required=True) parser.add_argument("--exclude-release", type=Path, action="append", default=[]) parser.add_argument("--student-audit", type=Path) parser.add_argument("--holdout-quota", type=_parse_quota, action="append") parser.add_argument("--train-quota", type=_parse_quota, action="append") parser.add_argument("--seed", default="quality-capability-v1") parser.add_argument("--holdout-output", type=Path, required=True) parser.add_argument("--selection-output", type=Path, required=True) args = parser.parse_args() if args.holdout_quota: args.holdout_quota = dict(args.holdout_quota) if args.train_quota: args.train_quota = dict(args.train_quota) for path in (args.holdout_output, args.selection_output): if path.exists(): raise FileExistsError(path) holdout, selection = build(args) args.holdout_output.parent.mkdir(parents=True, exist_ok=True) args.selection_output.parent.mkdir(parents=True, exist_ok=True) args.holdout_output.write_bytes(canonical_json_line(holdout)) args.selection_output.write_bytes(canonical_json_line(selection)) print( json.dumps( { "holdout_rows": len(holdout["rows"]), "structural_counts_by_profile": selection["structural_counts_by_profile"], "selected_rows": len(selection["selected_rows"]), "training_allowed": selection["training_allowed"], }, indent=2, sort_keys=True, ) ) if __name__ == "__main__": main()