File size: 5,163 Bytes
ec1ecb9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """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()
|