| """Download pinned KEV and JevBench evaluation data, without training data.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import httpx |
|
|
| from jev_adapter.protocol import SystemOneRequest |
|
|
| from .data import ( |
| DEFAULT_SUITES, |
| KEV_COMMIT, |
| KEV_RAW, |
| KEV_REPOSITORY, |
| SUITES, |
| parse_partition, |
| population_counts, |
| require_sha256, |
| sha256, |
| source_provenance, |
| ) |
| from .jevbench import JEVBENCH_SUITES, prepare_jevbench |
|
|
|
|
| def read_source( |
| path: str, client: httpx.Client | None, source_root: Path | None |
| ) -> bytes: |
| if source_root is not None: |
| return (source_root / path).read_bytes() |
| if client is None: |
| raise ValueError("an HTTP client or --source-root is required") |
| response = client.get(f"{KEV_RAW}/{path}") |
| response.raise_for_status() |
| return response.content |
|
|
|
|
| def prepare_suite( |
| suite: str, |
| split: str, |
| output: Path, |
| *, |
| client: httpx.Client | None = None, |
| source_root: Path | None = None, |
| allow_test: bool = False, |
| limit: int | None = None, |
| ) -> dict: |
| if suite not in SUITES: |
| raise ValueError(f"unknown suite: {suite}") |
| if split not in ("development", "test"): |
| raise ValueError("only development and test are evaluation partitions") |
| if split == "test" and not allow_test: |
| raise ValueError( |
| "locked test requires --allow-test; use development for iteration" |
| ) |
| if limit is not None and limit < 1: |
| raise ValueError("--limit must be positive") |
| spec = SUITES[suite] |
| manifest_bytes = read_source(f"{spec.path}/manifest.json", client, source_root) |
| require_sha256(manifest_bytes, spec.manifest_sha256, f"{suite}/manifest.json") |
| upstream = json.loads(manifest_bytes) |
| expected_file = upstream["files"][f"{split}.jsonl"] |
| raw_bytes = read_source(f"{spec.path}/{split}.jsonl", client, source_root) |
| require_sha256(raw_bytes, expected_file["sha256"], f"{suite}/{split}.jsonl") |
| records = parse_partition(raw_bytes, suite, split) |
| counts = population_counts(records) |
| if counts["records"] != expected_file["records"]: |
| raise ValueError(f"upstream record count mismatch: {suite}/{split}") |
| if ( |
| "questions" in expected_file |
| and counts["questions"] != expected_file["questions"] |
| ): |
| raise ValueError(f"upstream question count mismatch: {suite}/{split}") |
| if not records: |
| raise ValueError( |
| f"{suite}/{split} is empty; this suite may be development-only" |
| ) |
| selected = records if limit is None else records[:limit] |
| for row in selected: |
| SystemOneRequest.model_validate({"model": "validation-only", **row["record"]}) |
| serialized = ( |
| "".join( |
| json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n" |
| for row in selected |
| ) |
| ).encode("utf-8") |
| directory = output / suite |
| data_path = directory / f"{split}.jsonl" |
| manifest_path = directory / f"{split}.manifest.json" |
| upstream_path = directory / "upstream.manifest.json" |
| manifest = { |
| "schema_version": 1, |
| "suite": suite, |
| "split": split, |
| "description": spec.description, |
| "data_file": data_path.name, |
| "data_sha256": sha256(serialized), |
| "upstream": { |
| "repository": KEV_REPOSITORY, |
| "commit": KEV_COMMIT, |
| "path": f"{spec.path}/{split}.jsonl", |
| "url": f"{KEV_RAW}/{spec.path}/{split}.jsonl", |
| "sha256": sha256(raw_bytes), |
| "manifest_sha256": spec.manifest_sha256, |
| }, |
| "full_partition": counts, |
| "selected": population_counts(selected), |
| "selection": { |
| "method": "full" if len(selected) == len(records) else "prefix", |
| "requested_limit": limit, |
| "is_full_partition": len(selected) == len(records), |
| "note": ( |
| "Full frozen partition." |
| if len(selected) == len(records) |
| else "Smoke subset only; not a full-suite result. Pairs may be incomplete." |
| ), |
| }, |
| "protocol": { |
| "calibration_applied": False, |
| "training_data_downloaded": False, |
| "gold_labels_sent_to_model": False, |
| "headline_variant": "clean", |
| "exclude_from_headline_sources": ["unknowable"], |
| "locked_test": split == "test", |
| "notes": list(spec.notes), |
| "upstream_context_policy": upstream.get("context"), |
| }, |
| "provenance": source_provenance(records, upstream), |
| } |
| manifest_payload = ( |
| json.dumps(manifest, ensure_ascii=False, indent=2) + "\n" |
| ).encode() |
| |
| for path, content in ( |
| (data_path, serialized), |
| (manifest_path, manifest_payload), |
| (upstream_path, manifest_bytes), |
| ): |
| if path.exists() and path.read_bytes() != content: |
| raise FileExistsError( |
| f"refusing to overwrite a different prepared artifact: {path}" |
| ) |
| directory.mkdir(parents=True, exist_ok=True) |
| data_path.write_bytes(serialized) |
| manifest_path.write_bytes(manifest_payload) |
| upstream_path.write_bytes(manifest_bytes) |
| return manifest |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument( |
| "--suite", action="append", choices=tuple(SUITES) + tuple(JEVBENCH_SUITES) |
| ) |
| parser.add_argument( |
| "--split", choices=("development", "test"), default="development" |
| ) |
| parser.add_argument("--allow-test", action="store_true") |
| parser.add_argument( |
| "--limit", type=int, help="deterministic prefix smoke subset per suite" |
| ) |
| parser.add_argument( |
| "--source-root", type=Path, help="offline KEV checkout; same SHA256 checks" |
| ) |
| parser.add_argument( |
| "--jevbench-source-root", type=Path, help="offline pinned JevBench checkout" |
| ) |
| args = parser.parse_args(argv) |
| if args.split == "test" and not args.allow_test: |
| parser.error("--split test requires --allow-test") |
| if args.limit is not None and args.limit < 1: |
| parser.error("--limit must be positive") |
| suites = args.suite or ( |
| DEFAULT_SUITES if args.split == "test" else (*DEFAULT_SUITES, *JEVBENCH_SUITES) |
| ) |
| if args.split == "test" and any(suite in JEVBENCH_SUITES for suite in suites): |
| parser.error("JevBench supplies a public subset, not a locked test split") |
| with httpx.Client(timeout=60, follow_redirects=True) as client: |
| for suite in dict.fromkeys(suites): |
| if suite in JEVBENCH_SUITES: |
| manifest = prepare_jevbench( |
| suite, |
| args.output, |
| client=client, |
| source_root=args.jevbench_source_root, |
| limit=args.limit, |
| ) |
| print( |
| json.dumps( |
| { |
| "suite": suite, |
| "split": "public", |
| "selected": manifest["selected"], |
| "is_full_partition": manifest["selection"][ |
| "is_full_partition" |
| ], |
| } |
| ) |
| ) |
| continue |
| manifest = prepare_suite( |
| suite, |
| args.split, |
| args.output, |
| client=client, |
| source_root=args.source_root, |
| allow_test=args.allow_test, |
| limit=args.limit, |
| ) |
| print( |
| json.dumps( |
| { |
| "suite": suite, |
| "split": args.split, |
| "data": str(args.output / suite / f"{args.split}.jsonl"), |
| "selected": manifest["selected"], |
| "is_full_partition": manifest["selection"]["is_full_partition"], |
| } |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
