Biomanus/experiments/ablation/scripts/summarize_ablation_results.py

#!/usr/bin/env python3
"""Summarize Biomanus ablation runs for BioAgentBench and LAB-Bench."""

from __future__ import annotations

import argparse
import csv
import json
import re
from pathlib import Path
from statistics import mean
from typing import Any


SCRIPT_PATH = Path(__file__).resolve()
ABLATION_ROOT = SCRIPT_PATH.parent.parent
REPO_ROOT = ABLATION_ROOT.parent.parent

VARIANTS = {
    "biomanus": "full BioManus",
    "mcp_flat": "MCP + flat retrieval",
    "mcp_metadata": "MCP + metadata retrieval",
    "minus_graph": "MCP infra only / minus graph",
    "minus_mcp": "no MCP",
    "minus_mcp_graph": "no MCP + no graph",
}


def load_json(path: Path, default: Any = None) -> Any:
    if not path.exists():
        return default
    return json.loads(path.read_text(encoding="utf-8"))


def token_count(text: str) -> int:
    if not text:
        return 0
    try:
        import tiktoken

        return len(tiktoken.get_encoding("cl100k_base").encode(text))
    except Exception:
        return max(1, len(re.findall(r"\S+", text)))


def latest_run_dir(runs_root: Path, task_id: str) -> Path | None:
    pattern = re.compile(rf"^{re.escape(task_id)}_\d{{8}}_\d{{6}}$")
    candidates = sorted(path for path in runs_root.iterdir() if path.is_dir() and pattern.match(path.name))
    return candidates[-1] if candidates else None


def selected_tool_names(plan: dict[str, Any]) -> list[str]:
    names = plan.get("selected_resource_names", {}).get("tools", [])
    out: list[str] = []
    for item in names or []:
        if isinstance(item, str):
            out.append(item)
        elif isinstance(item, dict) and item.get("name"):
            out.append(str(item["name"]))
    if out:
        return out
    tools = plan.get("selected_resources", {}).get("tools", [])
    return [str(tool["name"]) for tool in tools if isinstance(tool, dict) and tool.get("name")]


def load_gold_counts(path: Path) -> dict[str, int]:
    payload = load_json(path, default={}) or {}
    counts = {}
    for task_id, item in payload.items():
        if isinstance(item, dict):
            counts[task_id] = len(item.get("gold_tools", []) or []) + len(item.get("gold_servers", []) or [])
    return counts


def artifact_match_score(eval_item: dict[str, Any]) -> float | None:
    scores: list[float] = []
    for artifact in eval_item.get("artifacts", []) or []:
        if not isinstance(artifact, dict):
            continue
        metrics = artifact.get("metrics", {})
        if not isinstance(metrics, dict):
            continue
        for key in ("match_f1", "f1", "key_f1"):
            value = metrics.get(key)
            if isinstance(value, (int, float)) and not isinstance(value, bool):
                scores.append(float(value))
                break
    return mean(scores) if scores else None


def result_match_score(eval_item: dict[str, Any]) -> float | None:
    artifact_score = artifact_match_score(eval_item)
    if artifact_score is not None:
        return artifact_score
    eval_result = eval_item.get("evaluation_results", {})
    if not isinstance(eval_result, dict):
        return None
    for key in ("results_match_score", "results_match"):
        value = eval_result.get(key)
        if isinstance(value, bool):
            return 1.0 if value else 0.0
        if isinstance(value, (int, float)):
            return float(value)
    return None


def bioagent_summary(runs_root: Path, evaluation_json: Path, gold_counts: dict[str, int]) -> dict[str, Any]:
    evaluation = load_json(evaluation_json, default={}) or {}
    eval_rows = {
        item.get("task_id"): item
        for item in evaluation.get("results", [])
        if isinstance(item, dict) and item.get("task_id")
    }
    task_ids = sorted(set(eval_rows) | set(gold_counts))
    if not task_ids:
        task_ids = sorted(
            {
                re.sub(r"_\d{8}_\d{6}$", "", path.name)
                for path in runs_root.iterdir()
                if path.is_dir() and re.search(r"_\d{8}_\d{6}$", path.name)
            }
        )

    match_scores = []
    selected_counts = []
    context_counts = []
    planning_latencies = []
    selection_rates = []
    gold_items = []
    overheads = []

    for task_id in task_ids:
        eval_item = eval_rows.get(task_id) or {}
        score = result_match_score(eval_item)
        if score is not None:
            match_scores.append(score)

        run_dir = Path(eval_item.get("run_dir") or "") if eval_item.get("run_dir") else latest_run_dir(runs_root, task_id)
        if not run_dir or not run_dir.exists():
            continue
        plan = load_json(run_dir / "retrieval_plan.json", default={}) or {}
        summary = load_json(run_dir / "run_summary.json", default={}) or {}

        selected = len(selected_tool_names(plan))
        selected_counts.append(selected)
        gold_items.append(gold_counts.get(task_id, 0))

        context_text = plan.get("planning_context_text") or json.dumps(plan.get("query_context", {}), ensure_ascii=False)
        context_counts.append(token_count(context_text))

        planning_latency = plan.get("planning_latency_seconds") or summary.get("planning_latency_seconds")
        total_runtime = plan.get("total_runtime_seconds") or summary.get("total_runtime_seconds")
        if isinstance(planning_latency, (int, float)):
            planning_latencies.append(float(planning_latency))
        if isinstance(planning_latency, (int, float)) and isinstance(total_runtime, (int, float)) and total_runtime:
            overheads.append(float(planning_latency) / float(total_runtime))

        registered = plan.get("registered_tool_count")
        if isinstance(registered, int) and registered > 0:
            selection_rates.append(selected / registered)

    return {
        "results_match": mean(match_scores) if match_scores else "",
        "selected_tools": mean(selected_counts) if selected_counts else "",
        "overhead_planning_ratio": mean(overheads) if overheads else "",
        "gold_items": mean(gold_items) if gold_items else "",
        "context_tokens": mean(context_counts) if context_counts else "",
        "planning_latency": mean(planning_latencies) if planning_latencies else "",
        "selection_rate": mean(selection_rates) if selection_rates else "",
    }


def labbench_accuracy(jsonl_path: Path) -> str:
    if not jsonl_path.exists():
        return ""
    total = 0
    correct = 0
    with jsonl_path.open("r", encoding="utf-8", errors="replace") as handle:
        for line in handle:
            if not line.strip():
                continue
            try:
                item = json.loads(line)
            except json.JSONDecodeError:
                continue
            if "answer" not in item or "agent_answer" not in item:
                continue
            total += 1
            correct += str(item["answer"]).strip() == str(item["agent_answer"]).strip()
    return f"{correct / total:.3f} ({correct}/{total})" if total else ""


def fmt(value: Any) -> Any:
    if isinstance(value, float):
        return round(value, 6)
    return value


def main() -> int:
    parser = argparse.ArgumentParser(description="Summarize Biomanus ablation experiment outputs.")
    parser.add_argument("--root", type=Path, default=ABLATION_ROOT)
    parser.add_argument(
        "--gold",
        type=Path,
        default=REPO_ROOT.parent / "Biomni" / "experiments" / "bioagent_bench" / "gold_tools.json",
    )
    parser.add_argument("--out-csv", type=Path, default=None)
    parser.add_argument("--out-json", type=Path, default=None)
    args = parser.parse_args()

    out_csv = args.out_csv or args.root / "results" / "ablation_summary.csv"
    out_json = args.out_json or args.root / "results" / "ablation_summary.json"
    gold_counts = load_gold_counts(args.gold)
    rows = []
    for variant, label in VARIANTS.items():
        variant_root = args.root / "results" / variant
        bioagent = bioagent_summary(
            variant_root / "bioagentbench",
            variant_root / "bioagentbench_evaluation.json",
            gold_counts,
        )
        row = {
            "Agent System": label,
            "results_match": bioagent["results_match"],
            "Selected Tools": fmt(bioagent["selected_tools"]),
            "Overhead/planning占整个流程": fmt(bioagent["overhead_planning_ratio"]),
            "Gold Items": fmt(bioagent["gold_items"]),
            "Context Tokens": fmt(bioagent["context_tokens"]),
            "Planning Latency": fmt(bioagent["planning_latency"]),
            "Selection Rate": fmt(bioagent["selection_rate"]),
            "LAB-Bench:DbQA": labbench_accuracy(variant_root / "labbench_DbQA.jsonl"),
            "LAB-Bench:SeqQA": labbench_accuracy(variant_root / "labbench_SeqQA.jsonl"),
        }
        rows.append(row)

    out_csv.parent.mkdir(parents=True, exist_ok=True)
    fields = list(rows[0])
    with out_csv.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)
    out_json.write_text(json.dumps({"rows": rows}, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Saved CSV: {out_csv}")
    print(f"Saved JSON: {out_json}")
    print(json.dumps({"rows": rows}, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())