Files
PlantAI-dev/benchmarks/plantgenebench/scripts/summarize.py
T
2026-08-03 18:42:47 +08:00

148 lines
5.1 KiB
Python
Executable File

#!/usr/bin/env python3
"""Summarize scored PlantGeneBench rows."""
from __future__ import annotations
import argparse
import csv
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, List
def load_jsonl(path: Path) -> List[Dict[str, Any]]:
rows: List[Dict[str, Any]] = []
with path.open("r", encoding="utf-8") as handle:
for lineno, line in enumerate(handle, start=1):
line = line.strip()
if not line:
continue
try:
rows.append(json.loads(line))
except json.JSONDecodeError as exc:
raise SystemExit(f"{path}:{lineno}: invalid JSON: {exc}") from exc
return rows
def task_metadata(tasks_path: Path | None) -> Dict[str, Dict[str, Any]]:
if not tasks_path:
return {}
if not tasks_path.exists():
return {}
return {row["task_id"]: row for row in load_jsonl(tasks_path)}
def mean(values: Iterable[float]) -> float:
values = list(values)
if not values:
return 0.0
return sum(values) / len(values)
def write_rows_csv(path: Path, rows: List[Dict[str, Any]], tasks: Dict[str, Dict[str, Any]]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
fields = [
"task_id",
"category",
"difficulty",
"score",
"stage_f1",
"forbidden_ok",
"params",
"dependencies",
"analysis_type",
"errors",
"expected_stages",
"predicted_stages",
]
with path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=fields)
writer.writeheader()
for row in rows:
meta = tasks.get(row["task_id"], {})
components = row.get("components") or {}
writer.writerow({
"task_id": row["task_id"],
"category": meta.get("category", ""),
"difficulty": meta.get("difficulty", ""),
"score": row.get("score", 0.0),
"stage_f1": components.get("stage_f1", 0.0),
"forbidden_ok": components.get("forbidden_ok", 0.0),
"params": components.get("params", 0.0),
"dependencies": components.get("dependencies", 0.0),
"analysis_type": components.get("analysis_type", 0.0),
"errors": ";".join(row.get("errors") or []),
"expected_stages": ",".join(row.get("expected_stages") or []),
"predicted_stages": ",".join(row.get("predicted_stages") or []),
})
def grouped_scores(rows: List[Dict[str, Any]], tasks: Dict[str, Dict[str, Any]], key: str) -> Dict[str, List[float]]:
grouped: Dict[str, List[float]] = defaultdict(list)
for row in rows:
label = str(tasks.get(row["task_id"], {}).get(key) or "unknown")
grouped[label].append(float(row.get("score") or 0.0))
return grouped
def write_summary_md(path: Path, rows: List[Dict[str, Any]], tasks: Dict[str, Dict[str, Any]]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
error_counter: Counter[str] = Counter()
for row in rows:
error_counter.update(row.get("errors") or [])
lines: List[str] = []
lines.append("# PlantGeneBench Summary")
lines.append("")
lines.append(f"- Tasks: {len(rows)}")
lines.append(f"- Mean score: {mean(float(row.get('score') or 0.0) for row in rows):.3f}")
lines.append(f"- Perfect tasks: {sum(1 for row in rows if float(row.get('score') or 0.0) >= 0.999)}")
lines.append("")
for group_key in ("category", "difficulty"):
lines.append(f"## By {group_key}")
lines.append("")
lines.append("| Group | N | Mean score |")
lines.append("|---|---:|---:|")
for label, values in sorted(grouped_scores(rows, tasks, group_key).items()):
lines.append(f"| {label} | {len(values)} | {mean(values):.3f} |")
lines.append("")
lines.append("## Error Types")
lines.append("")
lines.append("| Error | Count |")
lines.append("|---|---:|")
for error, count in error_counter.most_common():
lines.append(f"| {error} | {count} |")
if not error_counter:
lines.append("| none | 0 |")
lines.append("")
path.write_text("\n".join(lines), encoding="utf-8")
def run(args: argparse.Namespace) -> int:
rows = load_jsonl(args.scored)
tasks = task_metadata(args.tasks)
csv_path = Path(f"{args.out_prefix}.rows.csv")
md_path = Path(f"{args.out_prefix}.summary.md")
write_rows_csv(csv_path, rows, tasks)
write_summary_md(md_path, rows, tasks)
print(f"Wrote {csv_path}")
print(f"Wrote {md_path}")
return 0
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--scored", type=Path, required=True)
parser.add_argument("--tasks", type=Path, default=Path("benchmarks/plantgenebench/tasks/plantgenebench_core_seed.jsonl"))
parser.add_argument("--out-prefix", type=Path, required=True)
return parser.parse_args()
if __name__ == "__main__":
raise SystemExit(run(parse_args()))