bench/emonet/eval_emonet.py

#!/usr/bin/env python3
"""Evaluate the 61 VoiceClap attribute-regression heads on EmoNet Voice Bench.

Protocol (matches arXiv:2506.09827 Table 5): each row has ONE target emotion and
a human intensity in {0,1,2} averaged over 2-4 experts, mapped to 0/5/10.
We take the single head named by the row's emotion and score the clip with it.
Primary metrics are the scale-free ones (Pearson r, Spearman rho) because our
heads emit the Empathic-Insight 0-7 scale, not 0-10.
"""
import io, json, os, sys, time
import numpy as np
import pyarrow.parquet as pq
import soundfile as sf
import torch
import torch.nn.functional as F

BENCH = os.path.dirname(os.path.abspath(__file__))
NB = "/e/data1/datasets/playground/mmlaion/schuhmann1/dramabox"
sys.path.insert(0, f"{NB}/headspub/repo")
from voiceclap_heads import AttributeScorer  # noqa: E402

SR, MAXS = 16000, 480000
OUT = f"{BENCH}/results"
os.makedirs(OUT, exist_ok=True)


def decode(b):
    with sf.SoundFile(io.BytesIO(b)) as f:
        sr = f.samplerate
        x = f.read(frames=int(30.0 * sr), dtype="float32", always_2d=True)
    x = x.mean(1) if x.shape[1] > 1 else x[:, 0]
    if len(x) < 400:
        return None
    if sr != SR:
        import torchaudio
        x = torchaudio.functional.resample(torch.from_numpy(np.ascontiguousarray(x)), sr, SR).numpy()
    return x[:MAXS]


def pearson(a, b):
    if a.std() < 1e-9 or b.std() < 1e-9:
        return float("nan")
    return float(np.corrcoef(a, b)[0, 1])


def spearman(a, b):
    return pearson(np.argsort(np.argsort(a)).astype(float), np.argsort(np.argsort(b)).astype(float))


def main():
    lab = pq.read_table(f"{BENCH}/labels.parquet").to_pydict()
    n = len(lab["clip_id"])
    print(f"{n} benchmark rows", flush=True)

    scorer = AttributeScorer(heads_path=f"{NB}/attrdistill/heads/heads.pt")
    print("device", scorer.device, flush=True)

    preds = np.full(n, np.nan, dtype=np.float64)
    all61 = np.full((n, 61), np.nan, dtype=np.float32)
    dims61 = scorer.dims

    # group row indices by source parquet so each file is opened once
    by_file = {}
    for i in range(n):
        by_file.setdefault(lab["parquet_file"][i], []).append(i)

    t0 = time.time()
    done = 0
    for fpath, idxs in by_file.items():
        tbl = pq.read_table(os.path.join(BENCH, fpath), columns=["audioId"])
        col = tbl.column("audioId")
        CH = 128
        for s in range(0, len(idxs), CH):
            chunk = idxs[s:s + CH]
            wavs, keep = [], []
            for i in chunk:
                rec = col[lab["row_index"][i]].as_py()
                w = decode(rec["bytes"])
                if w is None:
                    continue
                wavs.append(w)
                keep.append(i)
            if not wavs:
                continue
            emb = scorer.embed(wavs, batch_size=32)
            sc = scorer.score_embeddings(emb)
            M = np.stack([sc[d] for d in dims61], 1)
            for j, i in enumerate(keep):
                all61[i] = M[j]
                preds[i] = sc[lab["head_name"][i]][j]
            done += len(chunk)
            if done % 1280 < CH:
                print(f"  {done}/{n}  {time.time()-t0:.0f}s", flush=True)

    gold = np.asarray(lab["gold_0_10"], dtype=np.float64)
    head = np.asarray(lab["head_name"])
    ok = np.isfinite(preds)
    print(f"scored {ok.sum()}/{n}", flush=True)

    def block(mask, name):
        per = {}
        for c in sorted(set(head[mask])):
            m = mask & (head == c)
            if m.sum() < 20:
                continue
            p, g = preds[m], gold[m]
            # oracle affine (upper bound on MAE/RMSE, fit ON the benchmark - reported as such)
            A = np.stack([p, np.ones_like(p)], 1)
            coef, *_ = np.linalg.lstsq(A, g, rcond=None)
            pc = A @ coef
            per[c] = {
                "n": int(m.sum()), "r": pearson(p, g), "rho": spearman(p, g),
                "mae_fixed": float(np.abs(np.clip(p / 7.0 * 10.0, 0, 10) - g).mean()),
                "rmse_fixed": float(np.sqrt(((np.clip(p / 7.0 * 10.0, 0, 10) - g) ** 2).mean())),
                "mae_oracle": float(np.abs(pc - g).mean()),
                "rmse_oracle": float(np.sqrt(((pc - g) ** 2).mean())),
                "pred_mean": float(p.mean()), "gold_mean": float(g.mean()),
            }
        agg = {k: float(np.nanmean([v[k] for v in per.values()]))
               for k in ["r", "rho", "mae_fixed", "rmse_fixed", "mae_oracle", "rmse_oracle"]}
        agg["n_classes"] = len(per)
        agg["n_rows"] = int(mask.sum())
        pooled_p, pooled_g = preds[mask], gold[mask]
        agg["pooled_r"] = pearson(pooled_p, pooled_g)
        agg["pooled_rho"] = spearman(pooled_p, pooled_g)
        print(f"\n=== {name}: rows={agg['n_rows']} classes={agg['n_classes']} "
              f"macro r={agg['r']:.3f} rho={agg['rho']:.3f} "
              f"MAE(fixed)={agg['mae_fixed']:.2f} MAE(oracle)={agg['mae_oracle']:.2f}", flush=True)
        return {"per_class": per, "macro": agg}

    res = {"all": block(ok, "ALL 12,600 rows")}
    una = np.asarray(lab["unanimous"], dtype=bool)
    res["unanimous"] = block(ok & una, "unanimous-annotator subset")

    # presence detection: absent-by-all (gold==0) vs present-by-all
    pa = np.asarray(lab["present_all"], dtype=bool)
    absent = ok & (gold == 0.0)
    present = ok & pa
    aucs = {}
    for c in sorted(set(head[ok])):
        a = preds[absent & (head == c)]
        b = preds[present & (head == c)]
        if len(a) < 10 or len(b) < 10:
            continue
        allv = np.concatenate([a, b])
        r = np.argsort(np.argsort(allv)).astype(float) + 1
        n1 = len(b)
        auc = (r[len(a):].sum() - n1 * (n1 + 1) / 2) / (len(a) * n1)
        aucs[c] = {"auc": float(auc), "n_absent": len(a), "n_present": n1}
    res["presence_auc"] = {"per_class": aucs,
                           "macro_auc": float(np.mean([v["auc"] for v in aucs.values()])) if aucs else None,
                           "n_classes": len(aucs)}
    print(f"\npresence-detection macro ROC-AUC over {len(aucs)} classes: "
          f"{res['presence_auc']['macro_auc']:.3f}", flush=True)

    json.dump(res, open(f"{OUT}/emonet_metrics.json", "w"), indent=1)
    np.save(f"{OUT}/emonet_pred_target.npy", preds)
    np.save(f"{OUT}/emonet_pred_all61.f16.npy", all61.astype(np.float16))
    json.dump(dims61, open(f"{OUT}/emonet_dims61.json", "w"))
    print("WROTE", OUT, flush=True)


if __name__ == "__main__":
    main()