"""Distillation training v4 — fixed intensity target + 6.25 Hz (G=8, 1 frame per dots.tts token). Changes vs v3: - group_size 8: emotion frame rate 12.5 -> 6.25 Hz (exactly 1 frame per 6.25 Hz LLM token) - intensity target redesigned: tgt(t) = p_emo(t) * clip(rms_raw(t)/GLOBAL_REF, 0, 1) p_emo = max teacher prob over the 6 emotional classes (neutral/other/unk EXCLUDED, so a confidently-neutral clip gets LOW intensity — the v3 gate saturated at ~0.99 because it included neutral and was sqrt-ed) rms_raw: raw 50 Hz RMS envelope, normalized by a GLOBAL reference level (median over the pool of per-utterance P90 RMS) so cross-utterance loudness differences survive (v3 min-max erased them) - intensity loss: MSE + Pearson-correlation + log-std matching (anti mean-collapse; v3 collapsed to a constant ~0.31 because plain MSE to a noisy target rewards the mean) Cache: /tmp/work/teacher_cache_v4/shard_*.pt (states, probs, rms_raw) Same held-out gate: RAVDESS actors 19-24. 4000 AdamW steps, batch 8, fp16, cuda:0. """ import glob, json, math, os, random, sys, time, warnings warnings.filterwarnings("ignore") import logging; logging.disable(logging.CRITICAL) import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, "/tmp/work/s2st_emotion") from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder FUNASR_DIR = "/tmp/work/emotion2vec_plus_base" CACHE_DIR = "/tmp/work/teacher_cache_v4" OUT_DIR = "/tmp/work/surgical_out_v4" BATCH, LR, STEPS, WARMUP, SEED = 8, 3e-4, 4000, 200, 42 VAL_FRAC, FP16, DEV = 0.05, True, "cuda:0" # labels: angry0 disgusted1 fearful2 happy3 neutral4 other5 sad6 surprised7 unk8 EMO_IDX = [0, 1, 2, 3, 6, 7] # exclude neutral / other / unk GATE_FLOOR = 0.02 # keep tiny salience from zeroing the whole curve def load_cache(): shards = sorted(glob.glob(os.path.join(CACHE_DIR, "shard_*.pt"))) assert shards, "no v4 teacher shards found" cache = [] for s in shards: cache.extend(torch.load(s, map_location="cpu", weights_only=False)) return cache def make_batch(pool): items = random.sample(pool, min(BATCH, len(pool))) T = max(it["states"].size(0) for it in items) B = len(items) states = torch.zeros(B, T, 768); probs = torch.zeros(B, T, 9); rms = torch.zeros(B, T) lengths = torch.zeros(B, dtype=torch.long) for j, it in enumerate(items): t = it["states"].size(0) states[j, :t] = it["states"].float(); probs[j, :t] = it["probs"].float() rms[j, :t] = it["rms"].float(); lengths[j] = t return states.to(DEV), probs.to(DEV), rms.to(DEV), lengths.to(DEV) def masked_pearson(a, b, m): n = m.sum(1).clamp(min=1) ma = (a * m).sum(1) / n; mb = (b * m).sum(1) / n da = (a - ma[:, None]) * m; db = (b - mb[:, None]) * m cov = (da * db).sum(1) va = (da * da).sum(1); vb = (db * db).sum(1) corr = cov / (va * vb).clamp(min=1e-10).sqrt() return corr, va, vb, n def intensity_target(tprobs, trms, global_ref): p6 = tprobs[..., EMO_IDX].max(-1).values # (B,T) emotional salience env = (trms / global_ref).clamp(0.0, 1.0) # absolute-level envelope return torch.clamp(p6, min=GATE_FLOOR) * env # (B,T) @ 50 Hz def compute_loss(enc, batch, global_ref): states, tprobs, trms, lengths = batch out = enc.adapter(states, lengths) embs, flen = out["frame_embs"], out["frame_lengths"] T, T_out = states.size(1), embs.size(1) up = F.interpolate(embs.transpose(1, 2), size=T, mode="nearest").transpose(1, 2) m = (torch.arange(T, device=states.device)[None, :] < lengths[:, None]).unsqueeze(-1) cos = F.cosine_similarity(up, states, dim=-1).unsqueeze(-1) l_frame = ((1 - cos) * m).sum() / m.sum().clamp(min=1) l_mse = F.mse_loss(F.layer_norm(up, (768,)), F.layer_norm(states, (768,)), reduction="none") l_mse = (l_mse * m).sum() / m.sum().clamp(min=1) om = (torch.arange(T_out, device=states.device)[None, :] < flen[:, None]).unsqueeze(-1) utt_s = (embs * om).sum(1) / flen.unsqueeze(-1).float() utt_probs_s = torch.softmax(enc.adapter.emotion_head(utt_s.float()), dim=-1) utt_t = (states * m).sum(1) / lengths.unsqueeze(-1).float() l_utt_cos = (1 - F.cosine_similarity(utt_s, utt_t, dim=-1)).mean() utt_probs_t = torch.softmax(enc.teacher.proj(utt_t), dim=-1) l_utt_kl = F.kl_div(torch.log(utt_probs_s.clamp(min=1e-8)), utt_probs_t, reduction="batchmean") tprobs_g = F.interpolate(tprobs.transpose(1, 2), size=T_out, mode="area").transpose(1, 2) tprobs_g = tprobs_g / tprobs_g.sum(-1, keepdim=True).clamp(min=1e-8) l_frm = F.kl_div(torch.log(out["frame_probs"].clamp(min=1e-8)), tprobs_g, reduction="none").sum(-1) l_frm = (l_frm * om.squeeze(-1)).sum() / om.sum().clamp(min=1) # ---- v4 intensity: salience-gated absolute-energy target, corr+std losses ---- tgt50 = intensity_target(tprobs, trms, global_ref) # (B,T) fp32 tgt = F.interpolate(tgt50[:, None, :], size=T_out, mode="area")[:, 0, :] omf = om.squeeze(-1).float() si = out["frame_intensity"][:, :, 0].float() l_int_mse = (((si - tgt) ** 2) * omf).sum() / omf.sum().clamp(min=1) corr, va, vb, n = masked_pearson(si, tgt, omf) ok = (va > 1e-8) & (vb > 1e-8) if ok.any(): std_s = (va[ok] / n[ok]).clamp(min=1e-12).sqrt() std_t = (vb[ok] / n[ok]).clamp(min=1e-12).sqrt() l_corr = (1 - corr[ok]).mean() l_var = ((std_s / std_t).clamp(1e-4, 1e4).log() ** 2).mean() else: l_corr = si.sum() * 0.0; l_var = si.sum() * 0.0 l_int = l_int_mse + 1.0 * l_corr + 0.5 * l_var loss = (1.0 * l_frame + 0.5 * l_mse + 0.5 * l_utt_cos + 1.0 * l_utt_kl + 1.0 * l_frm + 1.0 * l_int) return loss, {"frame": l_frame.item(), "mse": l_mse.item(), "utt_cos": l_utt_cos.item(), "utt_kl": l_utt_kl.item(), "frm_kl": l_frm.item(), "int_mse": l_int_mse.item(), "int_corr": (1 - l_corr.item()) if ok.any() else float("nan"), "int_var": l_var.item()} def main(): random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) os.makedirs(OUT_DIR, exist_ok=True) log = open(os.path.join(OUT_DIR, "train_v4.log"), "a", buffering=1) def P(*a): print(*a, flush=True); print(*a, file=log) P(f"=== distill v4 start {time.strftime('%H:%M:%S')} device={DEV} G=8 (6.25 Hz) ===") cache = load_cache() P(f"teacher cache: {len(cache)} utts from {len(glob.glob(CACHE_DIR + '/shard_*.pt'))} shards") random.shuffle(cache) n_val = max(16, int(len(cache) * VAL_FRAC)) val, train = cache[:n_val], cache[n_val:] P(f"train={len(train)} val={len(val)}") p90s = [torch.quantile(it["rms"].float(), 0.90).item() for it in cache if it["rms"].numel()] global_ref = float(np.median(p90s)) P(f"GLOBAL_REF (median P90 rms): {global_ref:.6f}") with open(os.path.join(OUT_DIR, "target_meta.json"), "w") as f: json.dump({"global_ref": global_ref, "emo_idx": EMO_IDX, "gate_floor": GATE_FLOOR, "formula": "tgt = clamp(max(p[emo6]),floor) * clip(rms/GLOBAL_REF,0,1)", "version": 4, "group_size": 8, "frame_rate_hz": 6.25}, f, indent=1) enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, device=DEV) enc.cfg.group_size = 8 enc.adapter = type(enc.adapter)(enc.cfg).to(DEV) adapter = enc.adapter opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: min((s + 1) / WARMUP, 0.5 * (1 + math.cos(math.pi * s / STEPS)))) scaler = torch.amp.GradScaler("cuda", enabled=FP16) best_val = float("inf") for step in range(STEPS): adapter.train() with torch.autocast("cuda", dtype=torch.float16, enabled=FP16): loss, parts = compute_loss(enc, make_batch(train), global_ref) opt.zero_grad(set_to_none=True) scaler.scale(loss).backward() scaler.unscale_(opt) torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) scaler.step(opt); scaler.update(); sched.step() if step % 100 == 0: P(f"step {step:4d} loss {loss.item():.4f} " + " ".join(f"{k}={v:.3f}" for k, v in parts.items())) if step % 400 == 0 or step == STEPS - 1: adapter.eval() vlosses, vcorr = [], [] with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=FP16): for _ in range(4): _, vp = compute_loss(enc, make_batch(val), global_ref) vlosses.append(sum(v for k, v in vp.items() if k != "int_corr")) vcorr.append(vp["int_corr"]) vloss = float(np.mean(vlosses)) P(f" >> val step {step}: {vloss:.4f} val_int_corr={np.nanmean(vcorr):.3f}") if vloss < best_val: best_val = vloss enc.save_surgical(os.path.join(OUT_DIR, "surgical_emotion_module.pt")) P(" >> saved best") enc.save_surgical(os.path.join(OUT_DIR, "surgical_emotion_module_last.pt")) P(f"=== done, best_val={best_val:.4f} ===") if __name__ == "__main__": main()