| """Compare students v1 / v2 / v3 on identical held-out inputs vs the frozen teacher. |
| |
| Checkpoints: v1 & v2 downloaded from HF (amir0907/student, amir0907/student-v2), |
| v3 freshly trained at /tmp/work/surgical_out. |
| Test set: RAVDESS actors 19-24 speech (never in ANY training pool) + the funasr |
| Chinese clip. Also emits a per-sample demo (audio + plot + json) for v3. |
| Saves /tmp/work/surgical_out/v1_v2_v3_compare.json |
| """ |
|
|
| import glob |
| import json |
| import os |
| import shutil |
| import sys |
| import time |
| import warnings |
|
|
| warnings.filterwarnings("ignore") |
| import logging |
|
|
| logging.disable(logging.CRITICAL) |
|
|
| 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", |
| "v3": "/tmp/work/surgical_out/surgical_emotion_module.pt", |
| } |
| OUT = "/tmp/work/surgical_out/v1_v2_v3_compare.json" |
| DEMO = "/tmp/work/surgical_out/demo_v3" |
| 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"} |
|
|
| 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 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 rms_env(wav, T): |
| frames = torch.zeros(T) |
| for i in range(T): |
| seg = wav[i * 320: (i + 1) * 320] |
| frames[i] = float(seg.pow(2).mean().sqrt()) if seg.numel() else (frames[i - 1] if i else 0.0) |
| if T >= 9: |
| frames = F.conv1d(frames[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0] |
| lo, hi = frames.min(), frames.max() |
| return torch.full_like(frames, 0.5) if hi - lo < 1e-6 else (frames - lo) / (hi - lo) |
|
|
|
|
| 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(): |
| enc = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV) |
| enc.adapter.load_state_dict(torch.load(ck, map_location=DEV, weights_only=False)["adapter_state"]) |
| enc.eval() |
| students[name] = enc |
| print(f"loaded {name}", flush=True) |
|
|
| files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True)) |
| files = [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 files: |
| 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) |
|
|
| report = {"n_files": len(bench), "per_student": {}, "files": {}} |
| stats = {k: {"agree": [], "js": [], "pool_cos": [], "frame_cos": [], |
| "inten_std": [], "inten_corr": []} for k in students} |
|
|
| def demo_plot(name_tag, src_path, wav, out, states, lengths, tp, tgt): |
| fp = out["frame_probs"][0].cpu().numpy() |
| inten = out["frame_intensity"][0, :, 0].cpu() |
| T = int(lengths[0]) |
| tgt_ds = F.interpolate(tgt[None, None], size=inten.size(0), mode="linear")[0, 0] |
| 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]) * 0.08 |
| 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"v3 surgical module @ 12.5 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] + 0.08) |
| ax[1].plot(t, inten.numpy(), color="crimson", lw=2.0, label="v3 intensity") |
| ax[1].plot(t, tgt_ds.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="v3") |
| 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 = 80 ms | {fp.shape[0]} frames @ 12.5 Hz") |
| 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), |
| "n_frames": int(fp.shape[0]), |
| "v3_top": LABELS[int(out["utt_probs"][0].argmax())], |
| "teacher_top": LABELS[int(tp.argmax())], |
| "v3_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) |
|
|
| for f in bench: |
| wav = load_wav(f).cpu() |
| dur = wav.numel() / 16000 |
| 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) |
| r = teacher.extract_features(src, padding_mask=None, mask=False) |
| states = r["x"].detach() |
| lengths = torch.tensor([states.size(1)], device=DEV) |
| tp = torch.softmax(teacher.proj(states.mean(1)), -1)[0] |
|
|
| T = int(lengths[0]) |
| tgt = rms_env(wav, T) |
| tmax = torch.softmax(teacher.proj(states[0]), -1)[:, :8].max(-1).values.cpu().clamp(0, 1).sqrt() |
| tgt = tmax * tgt |
|
|
| is_test_wav = f.endswith("example/test.wav") |
| bname = os.path.basename(f) |
| emo = "zh-angry(test)" if is_test_wav else EMO_FIELD.get(bname.split("-")[2], "?") |
| row = {"dur_s": round(dur, 2), "corpus": "funasr-zh" if is_test_wav else "RAVDESS-heldout", |
| "teacher_top": LABELS[int(tp.argmax())]} |
| for name in students: |
| 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_ds = F.interpolate(tgt[None, None], size=si.size(0), mode="linear")[0, 0] |
| ic = float(np.corrcoef(si.numpy(), tgt_ds.numpy())[0, 1]) if si.var() > 1e-6 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]["inten_std"].append(float(si.std())) |
| if ic is not None: |
| stats[name]["inten_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)} |
| report["files"][bname] = row |
|
|
| |
| if is_test_wav or bname.split("-")[3] == "02": |
| tag = ("zh-angry-test" if is_test_wav else f"{emo}_{bname[:20]}") |
| demo_plot(tag, f, wav, outs["v3"], states, lengths, tp, tgt) |
|
|
| 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_std": float(np.mean(s["inten_std"])), |
| "mean_intensity_corr": float(np.mean(s["inten_corr"])) if s["inten_corr"] else None, |
| } |
| 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() |
|
|
