#!/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()