| """Precompute v4 caches — same GPU-heavy pipeline as precompute_fast.py, but the |
| cache stores the RAW (unnormalized) 50 Hz RMS envelope instead of the old |
| min-max "conf-gated" intensity target. The v4 intensity target is derived in |
| train_v4 from (probs, rms_raw) with a global reference level. |
| |
| Cache item: {states (T,768) f16, probs (T,9) f16, rms (T,) f16, file} |
| Usage: python3 precompute_v4.py <device> <shard_idx> <n_shards> |
| """ |
| import os, sys, threading, queue, warnings |
| warnings.filterwarnings("ignore") |
| import logging; logging.disable(logging.CRITICAL) |
|
|
| import torch, torchaudio |
| import torch.nn.functional as F |
|
|
| sys.path.insert(0, "/tmp/work/s2st_emotion") |
| from precompute_shard_v3 import gather_files, MAX_SEC |
| from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder |
|
|
| DEV = sys.argv[1] if len(sys.argv) > 1 else "cuda:0" |
| SHARD = int(sys.argv[2]) if len(sys.argv) > 2 else 0 |
| N_SHARDS = int(sys.argv[3]) if len(sys.argv) > 3 else 4 |
| N_LOAD = 4 |
| OUT_DIR = "/tmp/work/teacher_cache_v4" |
| os.makedirs(OUT_DIR, exist_ok=True) |
|
|
| resamplers = {} |
|
|
| def to_16k(wav, sr): |
| if sr == 16000: |
| return wav |
| if sr not in resamplers: |
| resamplers[sr] = torchaudio.transforms.Resample(sr, 16000).to(DEV) |
| return resamplers[sr](wav.to(DEV)).cpu() |
|
|
| def load_one(f): |
| try: |
| wav, sr = torchaudio.load(f) |
| wav = wav.mean(0) |
| if sr != 16000: |
| wav = to_16k(wav, sr) |
| except Exception: |
| return None |
| if wav.numel() < 8000: |
| return None |
| if wav.numel() > MAX_SEC * 16000: |
| wav = wav[: int(MAX_SEC * 16000)] |
| wav = wav / wav.abs().max().clamp(min=1e-4) |
| return wav |
|
|
| def rms_raw_gpu(wav_dev, T_f): |
| """RAW 50 Hz RMS envelope (9-tap box smoothed), NO min-max normalization.""" |
| L = wav_dev.numel() |
| n_full = (L + 319) // 320 |
| w = torch.zeros(T_f * 320, device=wav_dev.device, dtype=wav_dev.dtype) |
| w[: min(L, T_f * 320)] = wav_dev[: min(L, T_f * 320)] |
| ms = w.view(T_f, 320).pow(2).mean(1).sqrt() |
| if T_f > n_full: |
| ms[n_full:] = ms[n_full - 1] |
| if T_f >= 9: |
| ms = F.conv1d(ms[None, None], torch.ones(9, device=DEV)[None, None] / 9.0, |
| padding=4)[0, 0] |
| return ms |
|
|
| def main(): |
| files, counts = gather_files() |
| mine = files[SHARD::N_SHARDS] |
| print(f"[v4 shard {SHARD}/{N_SHARDS} on {DEV}] total={len(files)} mine={len(mine)}", flush=True) |
|
|
| partial_path = f"{OUT_DIR}/shard_{SHARD}_partial.pt" |
| cache = [] |
| if os.path.exists(partial_path): |
| cache = torch.load(partial_path, map_location="cpu", weights_only=False) |
| print(f"[v4 shard {SHARD}] resuming with {len(cache)} done", flush=True) |
| mine = mine[len(cache):] |
|
|
| from funasr import AutoModel |
| fun = AutoModel(model="/tmp/work/emotion2vec_plus_base", trust_remote_code=False, |
| device=DEV, disable_update=True) |
| enc = SurgicalEmotionEncoder(fun.model, SurgicalConfig()).to(DEV) |
| enc.eval() |
| torch.set_grad_enabled(False) |
|
|
| q = queue.Queue(maxsize=32) |
| def loader(paths, q): |
| for p in paths: |
| q.put((p, load_one(p))) |
| q.put(None) |
| th = threading.Thread(target=loader, args=(mine, q), daemon=True) |
| th.start() |
|
|
| skipped = 0 |
| done = len(cache) |
| while True: |
| item = q.get() |
| if item is None: |
| break |
| path, wav = item |
| if wav is None: |
| skipped += 1 |
| continue |
| wg = wav.to(DEV) |
| src = F.layer_norm(wg[None], wg[None].shape) |
| st = enc.teacher.extract_features(src, padding_mask=None, mask=False)["x"][0].float() |
| Ti = st.size(0) |
| probs = torch.softmax(enc.teacher.proj(st), dim=-1) |
| rms = rms_raw_gpu(wg, Ti) |
| cache.append({ |
| "states": st.half().cpu(), |
| "probs": probs.half().cpu(), |
| "rms": rms.half().cpu(), |
| "file": os.path.basename(path), |
| }) |
| done = len(cache) |
| if done % 500 == 0: |
| print(f"[v4 shard {SHARD}] {done}/{len(files)}", flush=True) |
| if done % 2000 == 0: |
| torch.save(cache, partial_path) |
|
|
| torch.save(cache, f"{OUT_DIR}/shard_{SHARD}.pt") |
| if os.path.exists(partial_path): |
| os.remove(partial_path) |
| print(f"[v4 shard {SHARD}] saved {len(cache)} utts (skipped {skipped})", flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
