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