"""Compare student v1 (900-utt/900-step distill) vs v2 (full-data/2000-step). Both checkpoints are loaded on top of the same frozen teacher; evaluated on the identical held-out set (actors 19-24 speech, up to 4 files per emotion = 32 files). Metrics per student vs teacher: top-1 agreement, JS divergence, pooled cosine, frame cosine; plus intensity std (dynamics) and mean inference latency. Saves /tmp/work/surgical_out/v1_v2_compare.json """ import glob import json import os 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 sys.path.insert(0, "/tmp/work/s2st_emotion") from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" V1_CKPT = "/tmp/work/v1_student/surgical_emotion_module.pt" V2_CKPT = "/tmp/work/surgical_out/surgical_emotion_module.pt" OUT = "/tmp/work/surgical_out/v1_v2_compare.json" 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) 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(): fun = AutoModel if False else None from funasr import AutoModel fun = AutoModel(model=FUNASR_DIR, trust_remote_code=False, device=DEV, disable_update=True) teacher = fun.model.eval() v1 = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV) v1.adapter.load_state_dict(torch.load(V1_CKPT, map_location=DEV, weights_only=False)["adapter_state"]) v1.eval() v2 = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV) v2.adapter.load_state_dict(torch.load(V2_CKPT, map_location=DEV, weights_only=False)["adapter_state"]) v2.eval() students = {"v1": v1, "v2": v2} 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]) print(f"comparison inputs: {len(bench)} held-out files", flush=True) report = {"n_files": len(bench), "per_student": {}, "files": {}} stats = {k: {"agree": [], "js": [], "pool_cos": [], "frame_cos": [], "ms": [], "inten_std": [], "inten_corr": []} for k in students} for f in bench: wav = load_wav(f).cpu() dur = wav.numel() / 16000 states, lengths = None, None outs = {} for name, enc in students.items(): t0 = time.perf_counter() outs[name] = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV)) torch.cuda.synchronize() stats[name]["ms"].append((time.perf_counter() - t0) * 1000) if states is None: src = F.layer_norm(wav[None], wav[None].shape).to(DEV) r = teacher.extract_features(src, padding_mask=None, mask=False) states, lengths = r["x"].detach(), torch.tensor([r["x"].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 row = {"dur_s": round(dur, 2), "ravdess": EMO_FIELD.get(os.path.basename(f).split("-")[2], "?"), "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"][os.path.basename(f)] = row 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, "mean_ms_per_utt": float(np.mean(s["ms"])), } os.makedirs(os.path.dirname(OUT), exist_ok=True) 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()