#!/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()