| """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 "-"}) |
|
|
| |
| 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"}) |
|
|
| |
| 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() |
|
|
