"""Compare v1 / v2 / v4 on identical held-out inputs vs the frozen teacher. v1/v2: 12.5 Hz (G=4, old min-max intensity target) — checkpoints from HF. v4: 6.25 Hz (G=8, salience-gated absolute-energy target) — fresh. Bench: 32 held-out RAVDESS files (4 per emotion) + funasr zh clip. Extra gate: ground-truth intensity separation (normal 01 vs strong 02) over ALL 336 emotional held-out files, per student. Also emits a per-sample demo for v4 (audio + plot + json): one STRONG-intensity utterance per emotion + neutral + zh clip. Saves /tmp/work/surgical_out_v4/v1_v2_v4_compare.json """ import glob, json, os, shutil, sys import numpy as np import torch import torch.nn.functional as F import torchaudio import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt sys.path.insert(0, "/tmp/work/s2st_emotion") from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" CKPTS = {"v1": "/tmp/work/v1_student/surgical_emotion_module.pt", "v2": "/tmp/work/v2_student/surgical_emotion_module.pt", "v4": "/tmp/work/surgical_out_v4/surgical_emotion_module.pt"} OUT = "/tmp/work/surgical_out_v4/v1_v2_v4_compare.json" DEMO = "/tmp/work/surgical_out_v4/demo_v4" DEV = "cuda:0" LABELS = ["angry", "disgusted", "fearful", "happy", "neutral", "other", "sad", "surprised", "unk"] EMO_FIELD = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad", "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"} EMO_IDX = [0, 1, 2, 3, 6, 7] GATE_FLOOR = 0.02 with open("/tmp/work/surgical_out_v4/target_meta.json") as f: META = json.load(f) torch.set_grad_enabled(False) os.makedirs(DEMO, exist_ok=True) def load_wav(path, max_s=12.0): wav, sr = torchaudio.load(path) wav = wav.mean(0) if sr != 16000: wav = torchaudio.functional.resample(wav, sr, 16000) wav = wav[: int(max_s * 16000)] return wav / wav.abs().max().clamp(min=1e-4) def rms_raw(wav, T): L = wav.numel(); n_full = (L + 319) // 320 w = torch.zeros(T * 320); w[: min(L, T * 320)] = wav[: min(L, T * 320)] ms = w.view(T, 320).pow(2).mean(1).sqrt() if T > n_full: ms[n_full:] = ms[n_full - 1] if T >= 9: ms = F.conv1d(ms[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0] return ms def js_div(p, q, eps=1e-8): p = p.clamp(min=eps); q = q.clamp(min=eps) m = 0.5 * (p + q) kl = lambda a, b: float((a * (a.log() - b.log())).sum(-1)) return 0.5 * kl(p, m) + 0.5 * kl(q, m) def new_target(probs_t9, wav): p6 = probs_t9[:, EMO_IDX].max(-1).values.clamp(min=GATE_FLOOR).cpu() return p6 * (rms_raw(wav, probs_t9.size(0)) / META["global_ref"]).clamp(0, 1) def demo_plot(name_tag, wav, out, probs_t9, tgt50, rate_hz, extra): fp = out["frame_probs"][0].cpu().numpy() inten = out["frame_intensity"][0, :, 0].cpu() tgt = F.interpolate(tgt50[None, None], size=inten.size(0), mode="area")[0, 0] tp = torch.softmax(probs_t9.mean(1), -1) d = os.path.join(DEMO, name_tag) os.makedirs(d, exist_ok=True) torchaudio.save(os.path.join(d, "audio.wav"), wav[None], 16000) t = np.arange(fp.shape[0]) * (1.0 / rate_hz) fig, ax = plt.subplots(3, 1, figsize=(12, 9), gridspec_kw={"height_ratios": [3, 1.6, 1.2]}) colors = plt.cm.tab10(np.linspace(0, 1, 10)) for k, lab in enumerate(LABELS[:8]): ax[0].plot(t, fp[:, k], label=lab, lw=1.6, color=colors[k]) ax[0].set_title(f"v4 surgical module @ 6.25 Hz — {name_tag}") ax[0].set_ylabel("frame emotion probability") ax[0].legend(fontsize=8, ncol=8, loc="upper center", bbox_to_anchor=(0.5, 1.18)) ax[0].set_ylim(-0.03, 1.05); ax[0].set_xlim(0, t[-1] + 1.0 / rate_hz) ax[1].plot(t, inten.numpy(), color="crimson", lw=2.0, label="v4 intensity") ax[1].plot(t, tgt.numpy(), color="steelblue", lw=1.5, ls="--", label="target") ax[1].set_ylabel("intensity"); ax[1].set_ylim(0, 1.05); ax[1].legend(fontsize=8) y = np.arange(8) ax[2].barh(y + 0.2, tp[:8].cpu().numpy(), height=0.38, color="gray", label="teacher") ax[2].barh(y - 0.2, out["utt_probs"][0][:8].cpu().numpy(), height=0.38, color="crimson", label="v4") ax[2].set_yticks(y, LABELS[:8], fontsize=8); ax[2].set_xlim(0, 1.05) ax[2].invert_yaxis(); ax[2].legend(fontsize=8, loc="lower right") ax[2].set_xlabel("utterance-level probability") ax[1].set_xlabel(f"time (s) — frame = 160 ms | {fp.shape[0]} frames @ 6.25 Hz (1 frame / token)") fig.tight_layout() fig.savefig(os.path.join(d, "analysis.png"), dpi=120, bbox_inches="tight") plt.close(fig) res = {"file": name_tag, "duration_s": round(wav.numel() / 16000, 2), "frame_rate_hz": rate_hz, "n_frames": int(fp.shape[0]), **extra, "v4_top": LABELS[int(out["utt_probs"][0].argmax())], "teacher_top": LABELS[int(tp.argmax())], "v4_probs": {l: round(float(p), 4) for l, p in zip(LABELS, out["utt_probs"][0])}, "teacher_probs": {l: round(float(p), 4) for l, p in zip(LABELS, tp)}, "intensity": {"mean": round(float(inten.mean()), 4), "std": round(float(inten.std()), 4), "min": round(float(inten.min()), 4), "max": round(float(inten.max()), 4)}} with open(os.path.join(d, "result.json"), "w") as fj: json.dump(res, fj, indent=1) def main(): from funasr import AutoModel fun = AutoModel(model=FUNASR_DIR, trust_remote_code=False, device=DEV, disable_update=True) teacher = fun.model.eval() students = {} for name, ck in CKPTS.items(): payload = torch.load(ck, map_location=DEV, weights_only=False) cfg = SurgicalConfig() cfg.group_size = int((payload.get("config") or {}).get("group_size", 4)) enc = SurgicalEmotionEncoder(teacher, cfg).to(DEV) enc.adapter.load_state_dict(payload["adapter_state"]) enc.eval() students[name] = enc print(f"loaded {name} (G={cfg.group_size}, {cfg.frame_rate_hz} Hz)", flush=True) files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True)) speech = [f for f in files if os.path.basename(f).split("-")[1] == "01" and 19 <= int(os.path.basename(f).split("-")[-1].split(".")[0]) <= 24] by_emo = {} for f in speech: by_emo.setdefault(os.path.basename(f).split("-")[2], []).append(f) bench = [] for e in sorted(by_emo): bench.extend(by_emo[e][:4]) bench.append("/tmp/work/emotion2vec_plus_base/example/test.wav") print(f"comparison inputs: {len(bench)} files", flush=True) stats = {k: {"agree": [], "js": [], "pool_cos": [], "frame_cos": [], "int_corr": [], "int_std": []} for k in students} report = {"n_files": len(bench), "per_student": {}, "files": {}} for f in bench: wav = load_wav(f).cpu() outs = {} for name in students: outs[name] = students[name](wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) src = F.layer_norm(wav[None], wav[None].shape).to(DEV) states = teacher.extract_features(src, padding_mask=None, mask=False)["x"] probs_t9 = torch.softmax(teacher.proj(states), -1)[0] tp = torch.softmax(teacher.proj(states.mean(1)), -1)[0] tgt50 = new_target(probs_t9, wav) is_test = f.endswith("example/test.wav") bname = os.path.basename(f) emo = "zh-angry(test)" if is_test else EMO_FIELD.get(bname.split("-")[2], "?") row = {"dur_s": round(wav.numel() / 16000, 2), "corpus": "funasr-zh" if is_test else "RAVDESS-heldout", "teacher_top": LABELS[int(tp.argmax())]} for name, enc in students.items(): o = outs[name] sp = o["utt_probs"][0] a = int(sp.argmax().item() == tp.argmax().item()) j = js_div(sp[None], tp[None]) pc = float(F.cosine_similarity(o["utt_emb"][0], states.mean(1)[0], dim=0)) up = F.interpolate(o["frame_embs"].transpose(1, 2), size=states.size(1), mode="nearest").transpose(1, 2) fc = float(F.cosine_similarity(up[0], states[0], dim=-1).mean()) si = o["frame_intensity"][0, :, 0].cpu() tgt = F.interpolate(tgt50[None, None], size=si.size(0), mode="area")[0, 0] ic = (float(np.corrcoef(si.numpy(), tgt.numpy())[0, 1]) if si.var() > 1e-8 and tgt.var() > 1e-8 else None) stats[name]["agree"].append(a); stats[name]["js"].append(j) stats[name]["pool_cos"].append(pc); stats[name]["frame_cos"].append(fc) stats[name]["int_std"].append(float(si.std())) if ic is not None: stats[name]["int_corr"].append(ic) row[name] = {"top": LABELS[int(sp.argmax())], "p": round(float(sp.max()), 3), "js": round(j, 5), "frame_cos": round(fc, 4), "int_corr": round(ic, 3) if ic is not None else None, "int_mean": round(float(si.mean()), 3)} report["files"][bname] = row if is_test or bname.split("-")[3] == "02": tag = "zh-angry-test" if is_test else f"{emo}_{bname[:20]}" demo_plot(tag, wav, outs["v4"], probs_t9, tgt50, students["v4"].cfg.frame_rate_hz, {"ravdess_label": emo, "intensity_label": "strong" if not is_test else "-"}) # dedicated demo pass: one STRONG utterance per emotion + neutral + zh demo_files = [] for e in sorted(by_emo): if e == "01": continue strong = [f for f in by_emo[e] if os.path.basename(f).split("-")[3] == "02"] demo_files.append(strong[0] if strong else by_emo[e][0]) neutral = by_emo.get("01", []) if neutral: demo_files = [neutral[0]] + demo_files for f in demo_files: wav = load_wav(f).cpu() bname = os.path.basename(f) out = students["v4"](wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) src_ = F.layer_norm(wav[None], wav[None].shape).to(DEV) states = teacher.extract_features(src_, padding_mask=None, mask=False)["x"] probs_t9 = torch.softmax(teacher.proj(states), -1)[0] tgt50 = new_target(probs_t9, wav) emo = EMO_FIELD.get(bname.split("-")[2], "?") ilbl = bname.split("-")[3] demo_plot(f"{emo}_{bname[:20]}", wav, out, probs_t9, tgt50, students["v4"].cfg.frame_rate_hz, {"ravdess_label": emo, "intensity_label": "strong" if ilbl == "02" else "normal"}) # ground-truth intensity separation for ALL students emotional = [f for f in speech if os.path.basename(f).split("-")[2] != "01"] sep = {k: {"01": [], "02": []} for k in students} for f in emotional: wav = load_wav(f) lbl = os.path.basename(f).split("-")[3] for name, enc in students.items(): o = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) sep[name][lbl].append(float(o["frame_intensity"][0, :, 0].mean())) for name in students: a, b = np.array(sep[name]["01"]), np.array(sep[name]["02"]) d = (b.mean() - a.mean()) / np.sqrt((a.var() + b.var()) / 2 + 1e-9) stats[name]["sep"] = {"normal_mean": round(float(a.mean()), 4), "strong_mean": round(float(b.mean()), 4), "diff": round(float(b.mean() - a.mean()), 4), "cohen_d": round(float(d), 3)} print(f"sep {name}: normal {a.mean():.3f} strong {b.mean():.3f} " f"diff {b.mean()-a.mean():+.3f} d={d:.2f}", flush=True) for name in students: s = stats[name] report["per_student"][name] = { "top1_agreement": float(np.mean(s["agree"])), "mean_js": float(np.mean(s["js"])), "mean_pooled_cos": float(np.mean(s["pool_cos"])), "mean_frame_cos": float(np.mean(s["frame_cos"])), "mean_intensity_corr_newtarget": float(np.mean(s["int_corr"])) if s["int_corr"] else None, "mean_intensity_std": float(np.mean(s["int_std"])), "intensity_separation": s["sep"], } with open(OUT, "w") as fp: json.dump(report, fp, indent=1) print(json.dumps(report["per_student"], indent=1), flush=True) print("COMPARE_DONE", flush=True) if __name__ == "__main__": main()