scripts/build_leaderboard_data.py

#!/usr/bin/env python3
"""Build deterministic SearchGen-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]

MODEL_METADATA = {
    "bagel": ("Bagel", "Open"),
    "klein4b": ("Flux.2-Klein-4B", "Open"),
    "klein": ("Flux.2-Klein-9B", "Open"),
    "qwen1": ("Qwen-Image", "Open"),
    "imagen3fast": ("Imagen3-Fast", "Commercial"),
    "qwen2": ("Qwen-Image-2", "Commercial"),
    "qwen_image_2_pro": ("Qwen-Image-2-Pro", "Commercial"),
    "jimeng4d0": ("SeedDream-4.0", "Commercial"),
    "seedream4d5": ("SeedDream-4.5", "Commercial"),
    "xai_image": ("Grok-Imagine-Image", "Commercial"),
    "gemini2d5flash": ("Nano Banana", "Commercial"),
    "gemini3pro": ("Nano Banana Pro", "Commercial"),
    "gpt_image": ("GPT-Image-2", "Commercial"),
}

# Paper Table 1 uses the validated replacement Qwen-Image-2 run. Other uses of
# the legacy `qwen2` ID in ToolGen (for example Table 2) intentionally remain
# separate, so the source-directory mapping is local to this public leaderboard.
MODEL_SOURCE_IDS = {"qwen2": "qwen-image-2.0"}


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 / "augmented_parsed_result_ffjudge_pp.json"
        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_SOURCE_IDS.get(model_id, model_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]:
    if row.get("subset") == "NoSearch":
        return "NoSearch", "NoSearch"
    sample_id = row["sample_id"]
    if "texthard" in sample_id or "text_rendering" in sample_id:
        return "SearchIntensive", "TextualSearch"
    return "SearchIntensive", "VisualSearch"


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 canonical.TABLE1_GENS
        }
        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 = "exclude" if model_id in skip_missing else "zero_fill"
        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

        display_name, model_type = MODEL_METADATA[model_id]
        output.append(
            {
                "model_id": model_id,
                "display_name": 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["overall_10"] or -1), row["display_name"]))


def build_aggregates(records: list[dict[str, Any]], canonical: Any) -> dict[str, Any]:
    model_ids = list(canonical.TABLE1_GENS)
    skip_missing = set(canonical.SKIP_MISSING_GENS)
    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 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 = {
        "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),
    }
    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 canonical.TABLE1_GENS:
        score_list = []
        for record in records:
            score = record["models"][model_id]
            if score["status"] == "scored":
                score_list.append(score["components_raw_0to3"])
            elif model_id not in canonical.SKIP_MISSING_GENS:
                score_list.append(canonical.zero_fill())
        expected_overall = round(canonical.aggregate_components(score_list)["overall"] * 100 / 3, 1)
        actual = full_by_model[model_id]["overall_10"]
        if actual != expected_overall:
            errors.append(f"Overall-10 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 = canonical.EVAL_JSONL
    eval_rows = [json.loads(line) for line in eval_path.read_text().splitlines() if line.strip()]

    records = build_records(eval_rows, canonical)
    aggregates = build_aggregates(records, canonical)
    validate(records, aggregates, canonical)

    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest = {
        "schema_version": "1.0.0",
        "benchmark": "SearchGen-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": "overall_10",
            "source": "augmented_parsed_result_ffjudge_pp.json",
            "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": "zero_fill",
            "missing_policy_exceptions": sorted(canonical.SKIP_MISSING_GENS),
        },
        "partition": {
            "NoSearch": 100,
            "SearchIntensive": 651,
            "VisualSearch": 387,
            "TextualSearch": 264,
        },
                "models": {
            model_id: {
                "display_name": MODEL_METADATA[model_id][0],
                "type": MODEL_METADATA[model_id][1],
                "source_id": MODEL_SOURCE_IDS.get(model_id, model_id),
            }
            for model_id in canonical.TABLE1_GENS
        },
        "artifacts": [
            "prompt_scores.jsonl",
            "leaderboard_overall.json",
            "leaderboard_by_stratum.json",
            "leaderboard_by_domain.json",
            "leaderboard_by_failure_mode.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"])
    print(f"Validated and wrote {len(records)} prompts for {len(canonical.TABLE1_GENS)} models to {output_dir}")


if __name__ == "__main__":
    main()