| """Precompute frozen-teacher caches for a shard of training files. |
| |
| Usage: python3 precompute_shard.py <device> <shard_idx> <n_shards> |
| Files are enumerated deterministically (train actors 1-18, speech + song); |
| this process handles every n_shards-th file starting at shard_idx. |
| Saves /tmp/work/teacher_cache/shard_<idx>.pt |
| """ |
|
|
| import glob |
| import os |
| import sys |
| import warnings |
|
|
| warnings.filterwarnings("ignore") |
| import logging |
|
|
| logging.disable(logging.CRITICAL) |
|
|
| import torch |
| import torchaudio |
|
|
| sys.path.insert(0, "/tmp/work/s2st_emotion") |
| from funasr import AutoModel |
| 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 2 |
| MAX_SEC = 12.0 |
|
|
| OUT_DIR = "/tmp/work/teacher_cache" |
| os.makedirs(OUT_DIR, exist_ok=True) |
|
|
|
|
| def gather_files(): |
| files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True)) |
| out = [] |
| for f in files: |
| b = os.path.basename(f) |
| parts = b.split("-") |
| try: |
| actor = int(parts[-1].split(".")[0]) |
| channel = parts[1] |
| except Exception: |
| continue |
| if actor <= 18 and channel in ("01", "02"): |
| out.append(f) |
| return out |
|
|
|
|
| def rms_envelope_50hz(wav, T_frames): |
| frames = torch.zeros(T_frames) |
| for i in range(T_frames): |
| 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_frames >= 9: |
| frames = torch.nn.functional.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(): |
| files = gather_files() |
| mine = files[SHARD::N_SHARDS] |
| print(f"[shard {SHARD}/{N_SHARDS} on {DEV}] total files={len(files)}, mine={len(mine)}", flush=True) |
|
|
| 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) |
|
|
| cache, skipped = [], 0 |
| for i, f in enumerate(mine): |
| try: |
| wav, sr = torchaudio.load(f) |
| wav = wav.mean(0) |
| if sr != 16000: |
| wav = torchaudio.functional.resample(wav, sr, 16000) |
| except Exception: |
| skipped += 1 |
| continue |
| if wav.numel() < 8000: |
| skipped += 1 |
| continue |
| if wav.numel() > MAX_SEC * 16000: |
| wav = wav[: int(MAX_SEC * 16000)] |
| wav = wav / wav.abs().max().clamp(min=1e-4) |
| states, lengths = enc.extract_teacher_states(wav[None].to(DEV), |
| torch.tensor([wav.numel()], device=DEV)) |
| probs, inten = enc.teacher_frame_probs(states, lengths) |
| T = int(lengths[0].item()) |
| rms = rms_envelope_50hz(wav, T) |
| inten = (inten.squeeze(0)[:, 0].cpu().clamp(min=0, max=1).sqrt() * rms).unsqueeze(-1) |
| cache.append({ |
| "states": states.squeeze(0).half().cpu(), |
| "probs": probs.squeeze(0).half().cpu(), |
| "inten": inten.half().cpu(), |
| "file": os.path.basename(f), |
| }) |
| if (i + 1) % 200 == 0: |
| print(f"[shard {SHARD}] {i+1}/{len(mine)}", flush=True) |
|
|
| torch.save(cache, f"{OUT_DIR}/shard_{SHARD}.pt") |
| print(f"[shard {SHARD}] saved {len(cache)} utts (skipped {skipped}) -> shard_{SHARD}.pt", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
