WalkXR-AI / eval / aggregate_rubrics.py
aggregate_rubrics.py
Raw
import csv, argparse
from pathlib import Path
from collections import defaultdict
from statistics import mean, pstdev

NUM_FIELDS = ["persona_adherence","empathy","flow","helpfulness","grounding","safety","clarity"]

def to_float(x):
    try:
        return float(x)
    except Exception:
        return None

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--eval_id", required=True)
    ap.add_argument("--in_dir", default="rubrics/responses")
    ap.add_argument("--out_dir", default=None)
    args = ap.parse_args()

    base = Path(args.in_dir) / args.eval_id
    out_dir = Path(args.out_dir) if args.out_dir else Path("eval/runs")/args.eval_id/"summary"
    out_dir.mkdir(parents=True, exist_ok=True)

    rows = []
    for p in base.rglob("rubric_*.csv"):
        with open(p, newline="", encoding="utf-8") as f:
            for r in csv.DictReader(f):
                rows.append(r)
    # also read the shared rubric_v1.csv if people appended there
    shared = Path("rubrics/rubric_v1.csv")
    if shared.exists():
        with open(shared, newline="", encoding="utf-8") as f:
            for r in csv.DictReader(f):
                if r.get("scenario_id"):
                    rows.append(r)

    by_scn = defaultdict(lambda: {k: [] for k in NUM_FIELDS})
    counts = defaultdict(int)
    for r in rows:
        scn = (r.get("scenario_id") or "").strip()
        if not scn:
            continue
        counts[scn] += 1
        for k in NUM_FIELDS:
            v = to_float(r.get(k))
            if v is not None:
                by_scn[scn][k].append(v)

    # write scenario-level CSV
    out_csv = out_dir / "manual_scores_by_scenario.csv"
    with open(out_csv, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        header = ["scenario_id","n_ratings"]
        for k in NUM_FIELDS:
            header += [f"mean_{k}", f"stdev_{k}"]
        w.writerow(header)
        for scn in sorted(by_scn.keys()):
            row = [scn, counts[scn]]
            for k in NUM_FIELDS:
                vals = by_scn[scn][k]
                m = round(mean(vals), 3) if vals else ""
                s = round(pstdev(vals), 3) if len(vals) > 1 else ""
                row += [m, s]
            w.writerow(row)

    print(f"✔ wrote {out_csv}")

if __name__ == "__main__":
    main()