scripts/export_eval_excel.py

"""Export one completed fair-protocol run to a filterable Excel workbook."""

from __future__ import annotations

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

from openpyxl import Workbook
from openpyxl.styles import Font, PatternFill
from openpyxl.utils import get_column_letter


DATASETS = ("hotpotqa", "2wikimultihopqa", "musique")
PASS_ORDER = {"accuracy": 0, "isolated_latency": 1, "loaded_latency": 2}


def _metadata(run_id: str) -> tuple[str, str, str]:
    for dataset in DATASETS:
        if run_id == f"{run_id.rsplit('_' + dataset, 1)[0]}_{dataset}":
            prefix = run_id[: -(len(dataset) + 1)]
            if prefix.endswith("_accuracy"):
                return dataset, "accuracy", ""
    match = re.search(r"_latency_(isolated|loaded)_(hotpotqa|2wikimultihopqa|musique)_(.+)$", run_id)
    if not match:
        raise ValueError(f"cannot parse run ID: {run_id}")
    mode, dataset, system = match.groups()
    return dataset, f"{mode}_latency", system


def _rows(path: Path) -> list[dict[str, Any]]:
    with path.open(encoding="utf-8") as source:
        return list(csv.DictReader(source))


def _write_sheet(book: Workbook, name: str, rows: list[dict[str, Any]]) -> None:
    sheet = book.create_sheet(name)
    if not rows:
        return
    columns = list(rows[0])
    sheet.append(columns)
    for row in rows:
        sheet.append(
            [
                json.dumps(value, ensure_ascii=False) if isinstance(value, (dict, list, tuple)) else value
                for column in columns
                for value in [row.get(column)]
            ]
        )
    sheet.freeze_panes = "A2"
    sheet.auto_filter.ref = sheet.dimensions
    fill = PatternFill("solid", fgColor="1F4E78")
    for cell in sheet[1]:
        cell.font = Font(color="FFFFFF", bold=True)
        cell.fill = fill
    for index, column in enumerate(columns, start=1):
        values = [str(row.get(column, "")) for row in rows[:100]]
        sheet.column_dimensions[get_column_letter(index)].width = min(42, max(12, len(column) + 2, *(len(v) + 2 for v in values)))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-prefix", required=True)
    parser.add_argument("--artifacts", type=Path, default=Path("artifacts"))
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    report_root = args.artifacts / "reports"
    summary_rows: list[dict[str, Any]] = []
    raw_rows: list[dict[str, Any]] = []
    for csv_path in sorted(report_root.glob(f"{args.run_prefix}*/summary.csv")):
        run_id = csv_path.parent.name
        dataset, evaluation_pass, system_from_id = _metadata(run_id)
        for row in _rows(csv_path):
            row = {"dataset": dataset, "evaluation_pass": evaluation_pass, **row}
            if system_from_id:
                row["system"] = system_from_id
            summary_rows.append(row)
    for results_path in sorted((args.artifacts / "runs").glob(f"{args.run_prefix}*/results.jsonl")):
        run_id = results_path.parent.name
        if run_id.endswith("-warmup"):
            continue
        dataset, evaluation_pass, _ = _metadata(run_id)
        for line in results_path.read_text(encoding="utf-8").splitlines():
            record = json.loads(line)
            raw_rows.append({"dataset": dataset, "evaluation_pass": evaluation_pass, **record})

    summary_rows.sort(key=lambda row: (PASS_ORDER[row["evaluation_pass"]], row["dataset"], row["system"]))
    raw_rows.sort(key=lambda row: (PASS_ORDER[row["evaluation_pass"]], row["dataset"], row["system"], row["example_id"]))
    expected = len(DATASETS) * 9 * 3 * 50
    if len(raw_rows) != expected or any(row.get("status") != "completed" for row in raw_rows):
        raise RuntimeError(f"expected {expected} completed records, found {len(raw_rows)}")

    book = Workbook()
    overview = book.active
    overview.title = "Overview"
    overview.append(["Fair retrieval evaluation", args.run_prefix])
    overview.append(["Benchmarks", ", ".join(DATASETS)])
    overview.append(["Systems", 9])
    overview.append(["Samples per benchmark", 50])
    overview.append(["Completed records", len(raw_rows)])
    overview.append(["Failures", 0])
    overview.append(["Passes", "accuracy; isolated_latency; loaded_latency"])
    overview.column_dimensions["A"].width = 30
    overview.column_dimensions["B"].width = 65
    for cell in overview[1]:
        cell.font = Font(bold=True, color="FFFFFF")
        cell.fill = PatternFill("solid", fgColor="1F4E78")

    _write_sheet(book, "Summary_all", summary_rows)
    _write_sheet(book, "Accuracy", [row for row in summary_rows if row["evaluation_pass"] == "accuracy"])
    _write_sheet(book, "Latency_isolated", [row for row in summary_rows if row["evaluation_pass"] == "isolated_latency"])
    _write_sheet(book, "Latency_loaded", [row for row in summary_rows if row["evaluation_pass"] == "loaded_latency"])
    _write_sheet(book, "Raw_records", raw_rows)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    book.save(args.output)


if __name__ == "__main__":
    main()