Blame for scripts/build_leaderboard_data.py

#!/usr/bin/env python3
"""Build deterministic AgentGen-Bench prompt and aggregate data artifacts."""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable


OVERALL_9_COMPONENTS = [
    "checklist",
    "rubric_adaptive",
    "prompt_faithfulness",
    "image_quality",
    "text_rendering",
    "ai_naturalness",
    "composition_and_aesthetics",
    "physical_plausibility",
    "visual_reference_evaluation",
]
OVERALL_9_EXCLUDED_COMPONENT = "text_reference_evaluation"
DISPLAY_COMPONENTS = [*OVERALL_9_COMPONENTS, OVERALL_9_EXCLUDED_COMPONENT]
CONFIG_PATH = Path(__file__).resolve().parents[3] / "release_config.json"
CONFIG = json.loads(CONFIG_PATH.read_text())
MODELS = CONFIG["models"]
MODEL_BY_ID = {model["id"]: model for model in MODELS}
MODEL_IDS = [model["id"] for model in MODELS]
PRIMARY_FIELD = CONFIG["primary_metric"].lower().replace("overall", "overall_")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--toolgen-root",
        type=Path,
        default=Path(os.environ["TOOLGEN_ROOT"]) if "TOOLGEN_ROOT" in os.environ else None,
        help="ToolGen checkout containing paper_materials and final_20k_release_v2",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "public" / "data",
    )
    parser.add_argument(
        "--generated-at",
        help="ISO timestamp for reproducible rebuilds (defaults to SOURCE_DATE_EPOCH or now)",
    )
    args = parser.parse_args()
    if args.toolgen_root is None:
        parser.error("--toolgen-root or TOOLGEN_ROOT is required")
    return args


def load_canonical_module(toolgen_root: Path):
    candidates = [
        toolgen_root / "neurips_paper_materials" / "recompute_tables.py",
        toolgen_root / "paper_materials" / "recompute_tables.py",
    ]
    module_path = next((path for path in candidates if path.is_file()), None)
    if module_path is None:
        raise FileNotFoundError(f"Canonical scorer not found in: {', '.join(map(str, candidates))}")
    spec = importlib.util.spec_from_file_location("searchgen_recompute_tables", module_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Unable to import canonical scorer: {module_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def source_commit(toolgen_root: Path) -> str | None:
    try:
        return subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=toolgen_root,
            check=True,
            capture_output=True,
            text=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return None


def generated_at(value: str | None) -> str:
    if value:
        return value
    epoch = os.environ.get("SOURCE_DATE_EPOCH")
    if epoch:
        return datetime.fromtimestamp(int(epoch), tz=timezone.utc).isoformat()
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def mean_present(components: dict[str, float | None], keys: list[str]) -> float:
    values = [components[key] for key in keys if components.get(key) is not None]
    return sum(values) / len(values) if values else 0.0


def locate_result(row_dir: Path, model_id: str) -> tuple[str, Path | None]:
    candidate_dirs = [row_dir / "none" / f"{model_id}_generator", row_dir / "none" / model_id]
    existing_dir = False
    for model_dir in candidate_dirs:
        if model_dir.is_dir():
            existing_dir = True
        result_path = model_dir / CONFIG["dataset"]["result_filename"]
        if result_path.is_file():
            return "present", result_path
    return ("missing_evaluation" if existing_dir else "missing_generation"), None


def load_score(row_dir: Path, model_id: str, canonical: Any) -> dict[str, Any]:
    source_id = MODEL_BY_ID[model_id]["source_id"]
    status, result_path = locate_result(row_dir, source_id)
    if result_path is None:
        return {"status": status, "lane": "none", "components_raw_0to3": None}
    try:
        payload = json.loads(result_path.read_text())
        parsed = payload.get("parsed", {})
        if not parsed:
            raise ValueError("missing parsed result")
        components = canonical.extract_10comp(parsed)
    except (OSError, json.JSONDecodeError, TypeError, ValueError):
        return {"status": "invalid_evaluation", "lane": "none", "components_raw_0to3": None}

    return {
        "status": "scored",
        "lane": "none",
        "components_raw_0to3": components,
        "overall_10_raw": mean_present(components, list(canonical.COMPONENT_KEYS)),
        "overall_9_raw": mean_present(components, OVERALL_9_COMPONENTS),
    }


def classify_prompt(row: dict[str, Any]) -> tuple[str, str]:
    subsets = set(row.get(CONFIG["evaluation_slices"]["subset_field"]) or [])
    if "NoSearch" in subsets:
        return "NoSearch", "NoSearch"
    if subsets & {"texthard", "text_rendering_v2"}:
        return "SearchIntensive", "TextualSearch"
    return "SearchIntensive", "VisualSearch"


def select_evaluation_slice(rows: list[dict[str, Any]], slice_name: str) -> list[dict[str, Any]]:
    slice_config = CONFIG["evaluation_slices"]["named"][slice_name]
    field = CONFIG["evaluation_slices"]["subset_field"]
    tags = set(slice_config["include_tags"])
    selected = []
    seen = set()
    for row in rows:
        row_tags = set(row.get(field) or [])
        matches = not tags or (
            bool(row_tags & tags) if slice_config["match"] == "any" else tags <= row_tags
        )
        key = row[slice_config["deduplicate_by"]]
        if matches and key not in seen:
            selected.append(row)
            seen.add(key)
    if len(selected) != slice_config["expected_rows"]:
        raise ValueError(
            f"{slice_name} expected {slice_config['expected_rows']} rows, found {len(selected)}"
        )
    return selected


def select_record_slice(records: list[dict[str, Any]], slice_name: str) -> list[dict[str, Any]]:
    slice_config = CONFIG["evaluation_slices"]["named"][slice_name]
    tags = set(slice_config["include_tags"])
    selected = []
    seen = set()
    for record in records:
        row_tags = set(record.get("original_subset") or [])
        matches = not tags or (
            bool(row_tags & tags) if slice_config["match"] == "any" else tags <= row_tags
        )
        key = record[slice_config["deduplicate_by"]]
        if matches and key not in seen:
            selected.append(record)
            seen.add(key)
    if len(selected) != slice_config["expected_rows"]:
        raise ValueError(
            f"{slice_name} expected {slice_config['expected_rows']} records, found {len(selected)}"
        )
    return selected


def build_records(eval_rows: list[dict[str, Any]], canonical: Any) -> list[dict[str, Any]]:
    records = []
    for index, row in enumerate(eval_rows):
        stratum, search_type = classify_prompt(row)
        row_dir = canonical.RELEASE_ROOT / row["release_row"]
        models = {
            model_id: load_score(row_dir, model_id, canonical)
            for model_id in MODEL_IDS
        }
        records.append(
            {
                "sample_id": row["sample_id"],
                "prompt_index": index,
                "release_row": row["release_row"],
                "original_subset": row.get("subset"),
                "stratum": stratum,
                "search_type": search_type,
                "domains": sorted(set(row.get("domains", []))),
                "failure_modes": sorted(set(row.get("failure_modes", []))),
                "difficulty": row.get("difficulty"),
                "language": row.get("language"),
                "generation_task_type": row.get("generation_task_type"),
                "is_miniset": bool(row.get("is_miniset")),
                "models": models,
            }
        )
    return records


def aggregate_group(
    records: list[dict[str, Any]],
    model_ids: list[str],
    skip_missing: set[str],
) -> list[dict[str, Any]]:
    output = []
    for model_id in model_ids:
        scored = [r["models"][model_id] for r in records if r["models"][model_id]["status"] == "scored"]
        missing_policy = CONFIG["missing_scores"]["exceptions"].get(
            model_id, CONFIG["missing_scores"]["default"]
        )
        n_total = len(records)
        n_scored = len(scored)
        n_included = n_scored if missing_policy == "exclude" else n_total

        overall_9_values = [s["overall_9_raw"] for s in scored]
        overall_10_values = [s["overall_10_raw"] for s in scored]
        if missing_policy == "zero_fill":
            overall_9_values.extend([0.0] * (n_total - n_scored))
            overall_10_values.extend([0.0] * (n_total - n_scored))

        component_scores = {}
        component_counts = {}
        for component in DISPLAY_COMPONENTS:
            values = [
                s["components_raw_0to3"][component]
                for s in scored
                if s["components_raw_0to3"].get(component) is not None
            ]
            if missing_policy == "zero_fill":
                values.extend([0.0] * (n_total - n_scored))
            component_counts[component] = len(values)
            component_scores[component] = round((sum(values) / len(values)) * 100 / 3, 1) if values else None

        model = MODEL_BY_ID[model_id]
        output.append(
            {
                "model_id": model_id,
                "display_name": model["display_name"],
                "type": model["type"],
                "n_total": n_total,
                "n_scored": n_scored,
                "n_included": n_included,
                "coverage": round(n_scored / n_total, 4) if n_total else 0.0,
                "missing_policy": missing_policy,
                "overall_9": round((sum(overall_9_values) / len(overall_9_values)) * 100 / 3, 1)
                if overall_9_values
                else None,
                "overall_10": round((sum(overall_10_values) / len(overall_10_values)) * 100 / 3, 1)
                if overall_10_values
                else None,
                "components": component_scores,
                "component_counts": component_counts,
            }
        )
    return sorted(output, key=lambda row: (-(row[PRIMARY_FIELD] or -1), row["display_name"]))


def build_aggregates(records: list[dict[str, Any]], canonical: Any) -> dict[str, Any]:
    model_ids = MODEL_IDS
    skip_missing: set[str] = set()
    selectors: dict[str, Callable[[dict[str, Any]], bool]] = {
        "All": lambda _: True,
        "NoSearch": lambda row: row["stratum"] == "NoSearch",
        "SearchIntensive": lambda row: row["stratum"] == "SearchIntensive",
        "VisualSearch": lambda row: row["search_type"] == "VisualSearch",
        "TextualSearch": lambda row: row["search_type"] == "TextualSearch",
    }
    strata = {
        name: aggregate_group([r for r in records if selector(r)], model_ids, skip_missing)
        for name, selector in selectors.items()
    }
    domains = sorted({tag for row in records for tag in row["domains"]})
    failure_modes = sorted({tag for row in records for tag in row["failure_modes"]})
    by_domain = {
        tag: aggregate_group([r for r in records if tag in r["domains"]], model_ids, skip_missing)
        for tag in domains
    }
    by_failure_mode = {
        tag: aggregate_group([r for r in records if tag in r["failure_modes"]], model_ids, skip_missing)
        for tag in failure_modes
    }
    return {
        "overall": strata["All"],
        "strata": strata,
        "domains": by_domain,
        "failure_modes": by_failure_mode,
    }


def partition_counts(records: list[dict[str, Any]]) -> dict[str, int]:
    return {
        "NoSearch": sum(r["stratum"] == "NoSearch" for r in records),
        "SearchIntensive": sum(r["stratum"] == "SearchIntensive" for r in records),
        "VisualSearch": sum(r["search_type"] == "VisualSearch" for r in records),
        "TextualSearch": sum(r["search_type"] == "TextualSearch" for r in records),
    }


def validate(records: list[dict[str, Any]], aggregates: dict[str, Any], canonical: Any) -> None:
    errors = []
    sample_ids = [r["sample_id"] for r in records]
    if len(records) != 751:
        errors.append(f"expected 751 prompts, found {len(records)}")
    if len(set(sample_ids)) != len(sample_ids):
        errors.append("sample_id values are not unique")

    counts = partition_counts(records)
    expected = {"NoSearch": 100, "SearchIntensive": 651, "VisualSearch": 387, "TextualSearch": 264}
    if counts != expected:
        errors.append(f"partition mismatch: {counts} != {expected}")

    for row in records:
        if len(row["domains"]) != len(set(row["domains"])) or len(row["failure_modes"]) != len(set(row["failure_modes"])):
            errors.append(f"duplicate tag in {row['sample_id']}")
        for model_id, score in row["models"].items():
            components = score.get("components_raw_0to3")
            if score["status"] != "scored":
                continue
            for key, value in components.items():
                if value is not None and not 0 <= value <= 3:
                    errors.append(f"out-of-range score {row['sample_id']} {model_id} {key}={value}")
            expected_9 = mean_present(components, OVERALL_9_COMPONENTS)
            expected_10 = mean_present(components, list(canonical.COMPONENT_KEYS))
            if abs(score["overall_9_raw"] - expected_9) > 1e-12 or abs(score["overall_10_raw"] - expected_10) > 1e-12:
                errors.append(f"overall recomputation mismatch for {row['sample_id']} {model_id}")

    # Ensure the exported paper metric matches the canonical aggregation helper.
    full_by_model = {row["model_id"]: row for row in aggregates["overall"]}
    for model_id in MODEL_IDS:
        values = [
            record["models"][model_id][f"{PRIMARY_FIELD}_raw"]
            for record in records
            if record["models"][model_id]["status"] == "scored"
        ]
        expected_overall = round(sum(values) / len(values) * 100 / 3, 1)
        actual = full_by_model[model_id][PRIMARY_FIELD]
        if actual != expected_overall:
            errors.append(f"{CONFIG['primary_metric']} mismatch for {model_id}: {actual} != {expected_overall}")

    if errors:
        raise ValueError("Data validation failed:\n- " + "\n- ".join(errors[:50]))


def write_json(path: Path, payload: Any) -> None:
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True) + "\n")


def main() -> None:
    args = parse_args()
    toolgen_root = args.toolgen_root.resolve()
    canonical = load_canonical_module(toolgen_root)
    eval_path = Path(CONFIG["dataset"]["path"])
    all_eval_rows = [json.loads(line) for line in eval_path.read_text().splitlines() if line.strip()]
    for slice_name in CONFIG["evaluation_slices"]["named"]:
        select_evaluation_slice(all_eval_rows, slice_name)
    output_slice = CONFIG["outputs"]["interactive_leaderboard"]
    eval_rows = select_evaluation_slice(all_eval_rows, output_slice)

    records = build_records(eval_rows, canonical)
    aggregates = build_aggregates(records, canonical)
    validate(records, aggregates, canonical)
    published_slices = CONFIG["outputs"]["interactive_leaderboard_slices"]
    slice_records = {
        slice_name: select_record_slice(records, slice_name)
        for slice_name in published_slices
    }
    slice_aggregates = {
        slice_name: build_aggregates(selected, canonical)
        for slice_name, selected in slice_records.items()
    }
    slice_partitions = {
        slice_name: partition_counts(selected)
        for slice_name, selected in slice_records.items()
    }

    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest = {
        "schema_version": "1.1.0",
        "benchmark": "AgentGen-Bench",
        "generated_at": generated_at(args.generated_at),
        "dataset": {
            "filename": eval_path.name,
            "sha256": sha256_file(eval_path),
            "n_prompts": len(records),
        },
        "source_commit": source_commit(toolgen_root),
        "scoring": {
            "primary_metric": PRIMARY_FIELD,
            "metric_label": CONFIG["primary_metric"],
            "source": CONFIG["dataset"]["result_filename"],
            "scale_raw": [0, 3],
            "scale_public": [0, 100],
            "public_rounding_decimals": 1,
            "overall_10_components": list(canonical.COMPONENT_KEYS),
            "overall_9_components": OVERALL_9_COMPONENTS,
            "overall_9_excluded_component": OVERALL_9_EXCLUDED_COMPONENT,
            "missing_policy_default": CONFIG["missing_scores"]["default"],
            "missing_policy_exceptions": CONFIG["missing_scores"]["exceptions"],
        },
        "partition": {
            "NoSearch": 100,
            "SearchIntensive": 651,
            "VisualSearch": 387,
            "TextualSearch": 264,
        },
        "slice_partitions": slice_partitions,
                "models": {
            model_id: {
                "display_name": MODEL_BY_ID[model_id]["display_name"],
                "type": MODEL_BY_ID[model_id]["type"],
                "source_id": MODEL_BY_ID[model_id]["source_id"],
                "aliases": MODEL_BY_ID[model_id]["aliases"],
            }
            for model_id in MODEL_IDS
        },
        "evaluation_slices": CONFIG["evaluation_slices"],
        "output_slice": CONFIG["outputs"]["interactive_leaderboard"],
        "published_slices": published_slices,
        "artifacts": [
            "prompt_scores.jsonl",
            "leaderboard_overall.json",
            "leaderboard_by_stratum.json",
            "leaderboard_by_domain.json",
            "leaderboard_by_failure_mode.json",
            "leaderboard_by_slice.json",
        ],
    }

    with (output_dir / "prompt_scores.jsonl").open("w") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
    write_json(output_dir / "manifest.json", manifest)
    write_json(output_dir / "leaderboard_overall.json", aggregates["overall"])
    write_json(output_dir / "leaderboard_by_stratum.json", aggregates["strata"])
    write_json(output_dir / "leaderboard_by_domain.json", aggregates["domains"])
    write_json(output_dir / "leaderboard_by_failure_mode.json", aggregates["failure_modes"])
    write_json(output_dir / "leaderboard_by_slice.json", slice_aggregates)
    print(
        f"Validated and wrote {len(records)} prompts, {len(MODEL_IDS)} models, "
        f"and {len(slice_aggregates)} leaderboard slices to {output_dir}"
    )


if __name__ == "__main__":
    main()