| """Fine-tune Meta omniASR CTC (default omniASR_CTC_1B_v2, Apache-2.0) on Nigerian speech (+ optional synthetic code-switch). |
| Full fine-tune (CNN feature extractor frozen), bf16 autocast, AdamW, warmup + linear decay, duration-bucketed batches. |
| The model stays a plain CTC wav2vec2 (same architecture/tokenizer), so the duplex asr_server streams it unchanged. |
| ~/venv_omni/bin/python train_asr.py --run /scratch/asr/runs/A --steps 4000 [--synth /scratch/asr/synth/manifest.jsonl --synth_frac 0.2] |
| """ |
| import argparse, json, math, os, random, time, re, unicodedata |
| import numpy as np, torch, torch.nn.functional as F |
| import soundfile |
| from scipy.signal import resample_poly, fftconvolve |
| from fairseq2.models.hub import load_model |
| from fairseq2.data.tokenizers.hub import load_tokenizer |
| from fairseq2.nn.batch_layout import BatchLayout |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--card", default="omniASR_CTC_1B_v2"); ap.add_argument("--run", required=True) |
| ap.add_argument("--data", default="/scratch/asr/data"); ap.add_argument("--synth", default="") |
| ap.add_argument("--synth_frac", type=float, default=0.2); ap.add_argument("--steps", type=int, default=4000) |
| ap.add_argument("--lr", type=float, default=2e-5); ap.add_argument("--warmup", type=int, default=300) |
| ap.add_argument("--batch_sec", type=float, default=280.0); ap.add_argument("--max_sec", type=float, default=20.0) |
| ap.add_argument("--real_aug_p", type=float, default=0.15); ap.add_argument("--synth_aug_p", type=float, default=0.8) |
| ap.add_argument("--val_every", type=int, default=500); ap.add_argument("--seed", type=int, default=1) |
| ap.add_argument("--workers", type=int, default=10); ap.add_argument("--smoke", action="store_true") |
| a = ap.parse_args() |
| os.makedirs(a.run, exist_ok=True); json.dump(vars(a), open(f"{a.run}/args.json", "w"), indent=1) |
| random.seed(a.seed); np.random.seed(a.seed); torch.manual_seed(a.seed) |
| dev = torch.device("cuda") |
| if os.environ.get("MEM_FRAC"): torch.cuda.set_per_process_memory_fraction(float(os.environ["MEM_FRAC"])) |
| model = load_model(a.card, device=dev, dtype=torch.float32) |
| tok = load_tokenizer(a.card); enc = tok.create_encoder(); dec = tok.create_decoder(skip_special_tokens=True) |
| VI = tok.vocab_info; BLANK = 0; PAD = VI.pad_idx if VI.pad_idx is not None else 1 |
| print("vocab", VI.size, "pad", VI.pad_idx, "unk", VI.unk_idx, "bos", VI.bos_idx, "eos", VI.eos_idx, flush=True) |
| SPECIAL = {i for i in (VI.pad_idx, VI.bos_idx, VI.eos_idx) if i is not None} |
|
|
| def encode(t): |
| ids = [int(x) for x in enc(t).tolist()] |
| return [i for i in ids if i not in SPECIAL] |
|
|
| def rows(p): |
| return [json.loads(l) for l in open(p)] |
| train = rows(f"{a.data}/train.jsonl"); devset = rows(f"{a.data}/dev.jsonl") |
| synth = rows(a.synth) if a.synth else [] |
| for r in train: r["kind"] = "real" |
| for r in synth: r["kind"] = "synth"; r.setdefault("w", 1.0) |
| unk = VI.unk_idx |
| def ok(r): |
| ids = encode(r["text"]); r["n_ids"] = len(ids) |
| return len(ids) > 0 and (unk is None or unk not in ids) and r["dur"] <= a.max_sec |
| n0 = len(train); train = [r for r in train if ok(r)]; synth = [r for r in synth if ok(r)] |
| print(f"train {len(train)}/{n0} after tokenizer filter, synth {len(synth)}", flush=True) |
|
|
| _mm = {} |
| def load_audio(r): |
| if r.get("wav"): |
| y, sr = soundfile.read(r["wav"], dtype="float32") |
| if y.ndim > 1: y = y.mean(1) |
| if sr != 16000: y = resample_poly(y, 16000, sr).astype(np.float32) |
| return y |
| if r["pcm"] not in _mm: _mm[r["pcm"]] = np.memmap(r["pcm"], dtype=np.int16, mode="r") |
| y = _mm[r["pcm"]][r["off24"]: r["off24"] + r["n24"]].astype(np.float32) / 32768 |
| return resample_poly(y, 2, 3).astype(np.float32) |
|
|
| def augment(y, rng): |
| """noise (white/pink/babble-ish), simple synthetic reverb, phone band-limit, gain; ~ mic/channel variety""" |
| if rng.random() < 0.5: |
| rt = rng.uniform(0.15, 0.7); n = int(16000 * rt); t = np.arange(n) / 16000 |
| rir = rng.standard_normal(n) * np.exp(-6.9 * t / rt); rir[0] = 1.0; rir /= np.sqrt((rir ** 2).sum()) |
| y = fftconvolve(y, rir)[: len(y)].astype(np.float32) |
| if rng.random() < 0.35: |
| y = resample_poly(resample_poly(y, 1, 2), 2, 1)[: len(y)].astype(np.float32) if rng.random() < 0.5 else \ |
| resample_poly(resample_poly(y, 8, 16), 16, 8)[: len(y)].astype(np.float32) |
| if rng.random() < 0.5: y = np.tanh(y * rng.uniform(1.5, 4.0)).astype(np.float32) |
| if rng.random() < 0.7: |
| n = rng.standard_normal(len(y)).astype(np.float32) |
| if rng.random() < 0.5: n = np.cumsum(n); n -= np.convolve(n, np.ones(400) / 400, "same") |
| snr = rng.uniform(5, 30); ps = (y ** 2).mean() + 1e-8; pn = (n ** 2).mean() + 1e-8 |
| y = y + n * np.sqrt(ps / (pn * 10 ** (snr / 10))) |
| y = y * rng.uniform(0.3, 1.5) |
| return np.clip(y, -1, 1).astype(np.float32) |
|
|
| class DS(torch.utils.data.Dataset): |
| def __init__(self, items, train_mode): self.items, self.tm = items, train_mode |
| def __len__(self): return len(self.items) |
| def __getitem__(self, i): |
| r = self.items[i]; y = load_audio(r) |
| if self.tm: |
| rng = np.random.default_rng((i * 7919 + int(time.time() * 1000)) % 2 ** 32) |
| p = a.synth_aug_p if r["kind"] == "synth" else a.real_aug_p |
| if rng.random() < p: y = augment(y, rng) |
| return torch.from_numpy(y), torch.tensor(encode(r["text"])), i |
|
|
| def epoch_order(rng): |
| """weighted sampling (source weights) + synth mixed at synth_frac of audio hours; bucketed by duration""" |
| pick = [] |
| for r in train: |
| k = r["w"]; n = int(k) + (1 if rng.random() < k - int(k) else 0); pick += [r] * n |
| if synth: |
| real_h = sum(r["dur"] for r in pick); want = real_h * a.synth_frac / (1 - a.synth_frac); have = sum(r["dur"] for r in synth) |
| rep = want / have; s = [] |
| for r in synth: |
| n = int(rep) + (1 if rng.random() < rep - int(rep) else 0); s += [r] * n |
| pick += s |
| rng.shuffle(pick) |
| batches = [] |
| for c in range(0, len(pick), 4000): |
| ch = sorted(pick[c: c + 4000], key=lambda r: r["dur"]); cur, cs = [], 0.0 |
| for r in ch: |
| if cs + r["dur"] > a.batch_sec and cur: batches.append(cur); cur, cs = [], 0.0 |
| cur.append(r); cs += r["dur"] |
| if cur: batches.append(cur) |
| rng.shuffle(batches); return batches |
|
|
| def collate(b): |
| ys = [F.layer_norm(x, x.shape) for x, _, _ in b]; L = torch.tensor([len(x) for x in ys]) |
| X = torch.zeros(len(ys), int(L.max())); [X[i, : len(x)].copy_(x) for i, x in enumerate(ys)] |
| T = [t for _, t, _ in b]; TL = torch.tensor([len(t) for t in T]); Y = torch.cat(T) |
| return X, L, Y, TL |
|
|
| def forward(X, L): |
| bl = BatchLayout(X.shape, seq_lens=L.tolist(), device=dev) |
| logits, blo = model(X, bl) |
| return logits, torch.tensor(blo.seq_lens, device=dev) |
|
|
| for p in model.encoder_frontend.feature_extractor.parameters(): p.requires_grad = False |
| params = [p for p in model.parameters() if p.requires_grad] |
| print("trainable params %.1fM" % (sum(p.numel() for p in params) / 1e6), flush=True) |
| opt = torch.optim.AdamW(params, lr=a.lr, betas=(0.9, 0.98), eps=1e-8, weight_decay=0.01) |
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / a.warmup) * max(0.02, 1 - max(0, s - a.warmup) / max(1, a.steps - a.warmup))) |
|
|
| def cer_eval(items, bs=16): |
| model.eval(); hyps = [] |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| for i in range(0, len(items), bs): |
| b = [(torch.from_numpy(load_audio(r)), torch.tensor([1]), 0) for r in items[i: i + bs]] |
| X, L, _, _ = collate(b); lg, ol = forward(X.to(dev), L) |
| pr = lg.argmax(-1) |
| for j in range(pr.shape[0]): |
| s = pr[j, : ol[j]]; m = torch.ones_like(s, dtype=torch.bool); m[1:] = s[1:] != s[:-1]; s = s[m]; s = s[s != BLANK] |
| hyps.append(str(dec(s))) |
| model.train() |
| import jiwer |
| def n(x): x = unicodedata.normalize("NFC", x.lower()); x = re.sub(r"[^\w\s']", " ", x); return re.sub(r"\s+", " ", x).strip() or "x" |
| out = {} |
| for lg in sorted({r["lang"] for r in items}): |
| idx = [k for k, r in enumerate(items) if r["lang"] == lg] |
| out[lg] = round(jiwer.cer([n(items[k]["text"]) for k in idx], [n(hyps[k]) for k in idx]), 4) |
| out["mean"] = round(float(np.mean([v for k, v in out.items()])), 4) |
| return out, hyps |
|
|
| log = open(f"{a.run}/log.jsonl", "a") |
| v0, _ = cer_eval(devset); print("dev step0", v0, flush=True); log.write(json.dumps(dict(step=0, dev=v0)) + "\n"); log.flush() |
| best = v0["mean"]; step = 0; rng = random.Random(a.seed); model.train(); t0 = time.time(); ema = None |
| steps_total = 60 if a.smoke else a.steps |
| while step < steps_total: |
| batches = epoch_order(rng); print(f"epoch with {len(batches)} batches", flush=True) |
| lut = {} |
| flat = [] |
| for b in batches: |
| flat.append([len(lut) + k for k in range(len(b))]); [lut.setdefault(len(lut), r) for r in b] |
| items = [lut[k] for k in range(len(lut))] |
| dl = torch.utils.data.DataLoader(DS(items, True), batch_sampler=flat, num_workers=a.workers, collate_fn=collate, |
| prefetch_factor=4, persistent_workers=False) |
| for X, L, Y, TL in dl: |
| X, Y = X.to(dev, non_blocking=True), Y.to(dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| lg, ol = forward(X, L) |
| lp = F.log_softmax(lg.float(), -1).transpose(0, 1) |
| loss = F.ctc_loss(lp, Y, ol, TL.to(dev), blank=BLANK, reduction="sum", zero_infinity=True) / TL.sum() |
| opt.zero_grad(set_to_none=True); loss.backward() |
| gn = torch.nn.utils.clip_grad_norm_(params, 5.0) |
| if torch.isfinite(gn): opt.step() |
| sched.step(); step += 1 |
| ema = loss.item() if ema is None else 0.98 * ema + 0.02 * loss.item() |
| if step % 25 == 0: |
| el = time.time() - t0 |
| print(f"step {step} loss {ema:.4f} gn {gn:.2f} lr {sched.get_last_lr()[0]:.2e} {el / step:.2f}s/step mem {torch.cuda.max_memory_allocated() / 2**30:.1f}G", flush=True) |
| log.write(json.dumps(dict(step=step, loss=round(ema, 4), sps=round(el / step, 3))) + "\n"); log.flush() |
| if step % a.val_every == 0 or step == steps_total: |
| v, _ = cer_eval(devset); print("dev", step, v, flush=True); log.write(json.dumps(dict(step=step, dev=v)) + "\n"); log.flush() |
| sd = {k: t.to(torch.bfloat16) for k, t in model.state_dict().items()} |
| if v["mean"] < best and not a.smoke: |
| best = v["mean"]; torch.save(dict(model=sd, step=step, dev=v, card=a.card), f"{a.run}/best.pt"); print("saved best", flush=True) |
| if step >= steps_total: break |
| if not a.smoke: |
| torch.save(dict(model={k: t.to(torch.bfloat16) for k, t in model.state_dict().items()}, step=step, card=a.card), f"{a.run}/last.pt") |
| print("TRAIN_DONE best", best, flush=True) |
|
|
