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