dataset/scripts/validate_one_codex_run.py

#!/usr/bin/env python3
"""Validate one per-paper Codex JSON output before it can be merged."""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path

from jsonschema import Draft202012Validator


FORBIDDEN_FIELDS = {
    "paper_id",
    "task",
    "experiment",
    "question",
    "answer",
    "reference_answer",
    "rubric",
}

FORBIDDEN_PHRASES = [
    "paper-described backgrounds",
    "Build the discriminating topology",
    "Follow the paper's signal-region definitions",
    "Use the paper categories as written",
    "where applicable",
    "if relevant",
    "prompt Standard Model processes from the paper background discussion",
    "the usual backgrounds",
    "various SM backgrounds",
    "standard selections",
]

TASK_ORDER = ["target_to_signature", "signature_to_backgrounds"]
NUMBER_RE = re.compile(r"(?<![A-Za-z])\d+(?:\.\d+)?(?![A-Za-z])")
TRAINABLE_EVIDENCE_TRACE_RE = re.compile(
    r"(?im)(evidence trace|^\s*evidence\s*:|^\s*citations?\s*:|^\s*sources?\s*:|^\s*references?\s*:|source_file|claim_supported)"
)


def walk_forbidden(obj, path: str = "$") -> None:
    if isinstance(obj, dict):
        for key, value in obj.items():
            if key in FORBIDDEN_FIELDS:
                raise ValueError(f"Forbidden field {key!r} at {path}")
            walk_forbidden(value, f"{path}.{key}")
    elif isinstance(obj, list):
        for idx, value in enumerate(obj):
            walk_forbidden(value, f"{path}[{idx}]")
    elif isinstance(obj, str):
        lowered = obj.lower()
        for phrase in FORBIDDEN_PHRASES:
            if phrase.lower() in lowered:
                raise ValueError(f"Forbidden phrase {phrase!r} at {path}")


def assistant_text(ex: dict) -> str:
    messages = ex.get("messages", [])
    for message in messages:
        if isinstance(message, dict) and message.get("role") == "assistant":
            return str(message.get("content", ""))
    return ""


def example_text_for_numeric_check(ex: dict, kind: str) -> str:
    metadata = ex.get("metadata", {})
    parts = []
    if kind == "sft":
        parts.append(assistant_text(ex))
    else:
        parts.append(str(ex.get("chosen_answer", "")))
    parts.append(str(metadata.get("primary_signature", "")))
    parts.append(str(metadata.get("notes", "")))
    return "\n".join(parts)


def evidence_text(ex: dict) -> str:
    parts = []
    for item in ex.get("evidence", []):
        parts.append(str(item.get("claim_supported", "")))
        parts.append(str(item.get("quote", "")))
    return "\n".join(parts)


def validate_evidence_sources(ex: dict, source_pdf: str, source_tar: str) -> None:
    for idx, item in enumerate(ex.get("evidence", [])):
        source_file = item.get("source_file")
        if source_file == source_pdf:
            continue
        if source_tar != "NONE" and isinstance(source_file, str) and source_file.startswith(f"{source_tar}::") and len(source_file) > len(source_tar) + 2:
            continue
        if isinstance(source_file, str) and source_file and "/" not in source_file:
            # Older good runs sometimes used bare TeX/PDF member names. Prefer
            # tarball::member for new runs, but do not reject otherwise valid
            # outputs solely for this auditability issue.
            continue
        raise ValueError(
            f"{ex['id']} evidence[{idx}].source_file must be {source_pdf!r} "
            f"or {source_tar!r}::member, got {source_file!r}"
        )


def validate_numeric_evidence(ex: dict, kind: str) -> None:
    text = example_text_for_numeric_check(ex, kind)
    numbers = sorted(set(NUMBER_RE.findall(text)), key=lambda value: (len(value), value))
    if not numbers:
        return
    evidence = evidence_text(ex)
    missing = [number for number in numbers if number not in evidence]
    if missing:
        # Keep this as an advisory heuristic only. Good outputs can cite the
        # relevant table/selection while rendering the value differently from
        # the answer text, and those should not be quarantined automatically.
        return


def validate_no_evidence_trace_text(text: str, label: str) -> None:
    match = TRAINABLE_EVIDENCE_TRACE_RE.search(text)
    if match:
        raise ValueError(f"{label} contains evidence trace text: {match.group(0)!r}")


def validate_semantics(data: dict) -> None:
    paper = data["paper_index"]
    arxiv_id = paper["arxiv_id"]
    source_pdf = paper["source_pdf"]
    source_tar = paper["source_tar"]
    if paper["included"]:
        if paper["collaboration"] not in {"ATLAS", "CMS"}:
            raise ValueError("Included paper must have collaboration ATLAS or CMS")
        # Some runs include a short inclusion rationale here. That is harmless
        # and should not quarantine otherwise valid physics examples.
        if paper["split"] not in {"train", "val", "test"}:
            raise ValueError("Included paper must have non-null split")
        if paper["tasks"] != TASK_ORDER:
            raise ValueError(f"Included paper tasks must be {TASK_ORDER}")
        for field in ["physics_target", "dataset_description", "primary_signature"]:
            if not str(paper[field]).strip():
                raise ValueError(f"Included paper_index.{field} must be nonempty")
        if not paper["backgrounds"]:
            raise ValueError("Included paper_index.backgrounds must be nonempty")
        if data["skipped_paper"] is not None:
            raise ValueError("Included paper must set skipped_paper to null")

        sft = data["sft_examples"]
        rl = data["rl_examples"]
        if len(sft) != 2 or len(rl) != 2:
            raise ValueError("Included paper must contain exactly two SFT and two RL examples")

        sft_ids = [item["id"] for item in sft]
        rl_ids = [item["id"] for item in rl]
        if sft_ids != rl_ids:
            raise ValueError(f"SFT/RL ID mismatch: {sft_ids} != {rl_ids}")
        expected_ids = [f"arxiv_{arxiv_id}_{task}" for task in TASK_ORDER]
        if sft_ids != expected_ids:
            raise ValueError(f"Example IDs are wrong: {sft_ids} != {expected_ids}")

        if [item["task_type"] for item in sft] != TASK_ORDER:
            raise ValueError("SFT task order is wrong")
        if [item["task_type"] for item in rl] != TASK_ORDER:
            raise ValueError("RL task order is wrong")

        splits = {paper["split"]} | {item["split"] for item in sft} | {item["split"] for item in rl}
        if len(splits) != 1:
            raise ValueError(f"Split mismatch: {sorted(splits)}")

        for ex in sft:
            if len(ex["messages"]) != 3:
                raise ValueError(f"{ex['id']} does not have exactly 3 messages")
            roles = [message["role"] for message in ex["messages"]]
            if roles != ["system", "user", "assistant"]:
                raise ValueError(f"{ex['id']} has wrong message roles: {roles}")
            validate_no_evidence_trace_text(str(ex["messages"][1].get("content", "")), f"{ex['id']} user message")
            validate_no_evidence_trace_text(assistant_text(ex), f"{ex['id']} assistant message")
            if len(ex["evidence"]) < 2:
                raise ValueError(f"{ex['id']} has too little evidence")
            validate_evidence_sources(ex, source_pdf, source_tar)
            validate_numeric_evidence(ex, "sft")

        for ex in rl:
            if len(ex["evidence"]) < 2:
                raise ValueError(f"{ex['id']} has too little evidence")
            if ex["chosen_answer"].strip() == ex["rejected_answer"].strip():
                raise ValueError(f"{ex['id']} has identical chosen/rejected answers")
            validate_no_evidence_trace_text(str(ex.get("prompt", "")), f"{ex['id']} RL prompt")
            validate_no_evidence_trace_text(str(ex.get("chosen_answer", "")), f"{ex['id']} chosen_answer")
            validate_no_evidence_trace_text(str(ex.get("rejected_answer", "")), f"{ex['id']} rejected_answer")
            validate_evidence_sources(ex, source_pdf, source_tar)
            validate_numeric_evidence(ex, "rl")
    else:
        # Skipped papers may keep the deterministic split and descriptive
        # metadata. Only the absence of examples and a concrete skip reason
        # are required for merge safety.
        if paper["tasks"]:
            raise ValueError("Skipped paper_index.tasks must be empty")
        if data["sft_examples"] or data["rl_examples"]:
            raise ValueError("Skipped paper must not contain examples")
        if data["skipped_paper"] is None:
            raise ValueError("Skipped paper must include skipped_paper object")
        if not paper["reason"].strip():
            raise ValueError("Skipped paper_index.reason must be nonempty")


def main(argv: list[str]) -> int:
    if len(argv) != 3:
        print("Usage: validate_one_codex_run.py OUT_JSON RUN_SCHEMA", file=sys.stderr)
        return 2

    out_path = Path(argv[1])
    schema_path = Path(argv[2])
    data = json.loads(out_path.read_text())
    schema = json.loads(schema_path.read_text())

    validator = Draft202012Validator(schema)
    errors = sorted(validator.iter_errors(data), key=lambda err: list(err.path))
    if errors:
        print(f"Schema validation failed for {out_path}", file=sys.stderr)
        for err in errors[:20]:
            path = ".".join(str(part) for part in err.path) or "<root>"
            print(f"- {path}: {err.message}", file=sys.stderr)
        return 1

    try:
        walk_forbidden(data)
        validate_semantics(data)
    except ValueError as exc:
        print(f"Dataset validation failed for {out_path}: {exc}", file=sys.stderr)
        return 1

    print(f"OK: {out_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))