| |
| """CAFE-25lang — inference with the automatic language gate (no language flag). |
| |
| `inference.py` needs to be told whether the audio is English. This variant decides per utterance |
| with the small gate head shipped in `language_gate.safetensors`, so the caller supplies nothing: |
| |
| python inference_gated.py audio_en.wav audio_de.wav audio_pt.wav |
| |
| How it works: a first pass runs the bare trunk up to the gate's tap layer and pools those features |
| into p(non-English); utterances above the threshold are then transcribed with the adapters |
| attached, the rest on the bare trunk. Two passes, one set of weights. `--threshold` overrides the |
| packaged threshold. |
| |
| Requirements: nemo_toolkit[asr], torch, safetensors, soundfile. Audio must be 16 kHz mono. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
|
|
| import soundfile as sf |
| import torch |
| import torch.nn as nn |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| ADAPTER_LAYERS = (3, 7, 11, 15, 19, 23) |
| SR = 16000 |
|
|
|
|
| class Bottleneck(nn.Module): |
| def __init__(self, dim: int, rank: int): |
| super().__init__() |
| self.down = nn.Linear(dim, rank) |
| self.act = nn.GELU() |
| self.up = nn.Linear(rank, dim) |
|
|
| def forward(self, x): |
| return self.up(self.act(self.down(x))) |
|
|
|
|
| class AttentiveStatsPool(nn.Module): |
| """Attention-weighted mean and standard deviation over time, with a length mask.""" |
|
|
| def __init__(self, dim: int, attn_dim: int): |
| super().__init__() |
| self.proj = nn.Sequential(nn.Linear(dim, attn_dim), nn.Tanh(), nn.Linear(attn_dim, dim)) |
|
|
| def forward(self, x, lens): |
| t = x.shape[1] |
| mask = (torch.arange(t, device=x.device)[None, :] < lens[:, None]).unsqueeze(-1) |
| w = self.proj(x).masked_fill(~mask, float("-inf")).softmax(dim=1) |
| mean = (w * x).sum(1) |
| var = (w * (x - mean.unsqueeze(1)).pow(2)).sum(1) |
| return torch.cat([mean, var.clamp_min(1e-8).sqrt()], dim=-1) |
|
|
|
|
| class LanguageGate(nn.Module): |
| """Pooling -> LayerNorm -> MLP -> one logit; positive means non-English.""" |
|
|
| def __init__(self, dim=640, attn_dim=128, hidden=256): |
| super().__init__() |
| self.pool = AttentiveStatsPool(dim, attn_dim) |
| self.norm = nn.LayerNorm(2 * dim) |
| self.mlp = nn.Sequential(nn.Linear(2 * dim, hidden), nn.GELU(), |
| nn.Dropout(0.1), nn.Linear(hidden, 1)) |
|
|
| def forward(self, x, lens): |
| return self.mlp(self.norm(self.pool(x, lens))).squeeze(-1) |
|
|
|
|
| def load_gate(path: str, device): |
| with safe_open(path, framework="pt") as f: |
| meta = f.metadata() or {} |
| tap = int(meta.get("tap_encoder_layer", 11)) |
| attn = int(meta.get("attn_dim", 128)) |
| hidden = int(meta.get("hidden", 256)) |
| thr = float(meta.get("threshold", 0.8582)) |
| g = LanguageGate(640, attn, hidden) |
| g.load_state_dict(load_file(path)) |
| return g.to(device).eval(), tap, thr |
|
|
|
|
| def make_adapters(path: str, device, dtype): |
| sd = load_file(path) |
| rank, dim = sd["adapter.0.down.weight"].shape |
| mods = [] |
| for i in sorted({int(k.split(".")[1]) for k in sd}): |
| a = Bottleneck(dim, rank) |
| a.load_state_dict({"down.weight": sd[f"adapter.{i}.down.weight"], |
| "down.bias": sd[f"adapter.{i}.down.bias"], |
| "up.weight": sd[f"adapter.{i}.up.weight"], |
| "up.bias": sd[f"adapter.{i}.up.bias"]}) |
| mods.append(a.to(device=device, dtype=dtype).eval()) |
| return mods |
|
|
|
|
| def attach(model, mods): |
| handles = [] |
| for a, layer in zip(mods, ADAPTER_LAYERS): |
| def mk(a): |
| def hook(module, inp, out): |
| x = out[0] if isinstance(out, tuple) else out |
| if x.dim() != 3 or x.shape[-1] != a.down.in_features: |
| return out |
| y = x + a(x) |
| return (y,) + tuple(out[1:]) if isinstance(out, tuple) else y |
| return hook |
| handles.append(model.encoder.layers[layer].register_forward_hook(mk(a))) |
| return handles |
|
|
|
|
| def safe_decoding(model, cuda_graph: bool = False): |
| from omegaconf import open_dict |
|
|
| cfg = model.cfg.decoding |
| if cfg.strategy != "beam": |
| cfg.strategy = "greedy_batch" |
| with open_dict(cfg): |
| if hasattr(cfg, "greedy"): |
| cfg.greedy.use_cuda_graph_decoder = bool(cuda_graph) |
| cfg.greedy.loop_labels = True |
| cfg.fused_batch_size = -1 |
| model.change_decoding_strategy(cfg) |
|
|
|
|
| @torch.no_grad() |
| def gate_probs(model, gate, tap: int, files, device, max_len_s: float = 20.0): |
| """First pass: bare-trunk features at the tap layer -> p(non-English) per file.""" |
| buf = {} |
|
|
| def hook(module, inp, out): |
| x = out[0] if isinstance(out, tuple) else out |
| if x.dim() == 3: |
| buf["x"] = x |
| return out |
|
|
| h = model.encoder.layers[tap].register_forward_hook(hook) |
| probs = [] |
| try: |
| for f in files: |
| x, sr = sf.read(f, dtype="float32") |
| if x.ndim > 1: |
| x = x.mean(axis=1) |
| if sr != SR: |
| raise SystemExit(f"{f}: expected {SR} Hz mono, got {sr}") |
| if max_len_s and len(x) > int(max_len_s * SR): |
| x = x[: int(max_len_s * SR)] |
| audio = torch.from_numpy(x)[None].to(device) |
| alen = torch.tensor([audio.shape[-1]], device=device) |
| _, enc_len = model.forward(input_signal=audio.float(), input_signal_length=alen) |
| probs.append(float(torch.sigmoid(gate(buf["x"].float(), enc_len).float())[0])) |
| buf.clear() |
| finally: |
| h.remove() |
| return probs |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("audio", nargs="+", help="16 kHz mono wav files") |
| ap.add_argument("--model", default=os.path.join(HERE, "CAFE-25lang.nemo")) |
| ap.add_argument("--adapters", default=os.path.join(HERE, "adapters.safetensors")) |
| ap.add_argument("--gate", default=os.path.join(HERE, "language_gate.safetensors")) |
| ap.add_argument("--threshold", type=float, default=None, help="overrides the packaged value") |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| ap.add_argument("--batch_size", type=int, default=8) |
| ap.add_argument("--json", action="store_true", help="emit one JSON object per line") |
| ap.add_argument("--cuda-graph", dest="cuda_graph", action="store_true") |
| a = ap.parse_args() |
|
|
| import nemo.collections.asr as nemo_asr |
|
|
| device = torch.device(a.device) |
| dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 |
| model = nemo_asr.models.ASRModel.restore_from(a.model, map_location=device) |
| model.to(dtype).eval() |
| safe_decoding(model, a.cuda_graph) |
| gate, tap, thr = load_gate(a.gate, device) |
| if a.threshold is not None: |
| thr = a.threshold |
| print(f"[CAFE] gate on encoder layer {tap}, threshold {thr:.4f}") |
|
|
| probs = gate_probs(model, gate, tap, a.audio, device) |
| non_en = [f for f, p in zip(a.audio, probs) if p >= thr] |
| en = [f for f, p in zip(a.audio, probs) if p < thr] |
| print(f"[CAFE] {len(en)} English (bare trunk), {len(non_en)} non-English (adapters)") |
|
|
| res = {} |
| with torch.inference_mode(): |
| if en: |
| out = model.transcribe(en, batch_size=a.batch_size, verbose=False) |
| res.update(zip(en, out[0] if isinstance(out, tuple) else out)) |
| if non_en: |
| handles = attach(model, make_adapters(a.adapters, device, dtype)) |
| try: |
| out = model.transcribe(non_en, batch_size=a.batch_size, verbose=False) |
| res.update(zip(non_en, out[0] if isinstance(out, tuple) else out)) |
| finally: |
| for h in handles: |
| h.remove() |
| for f, p in zip(a.audio, probs): |
| text = getattr(res[f], "text", res[f]) |
| if a.json: |
| print(json.dumps({"audio": f, "p_non_english": round(p, 4), |
| "path": "adapters" if p >= thr else "trunk", |
| "text": text}, ensure_ascii=False)) |
| else: |
| print(f"{f}\tp_non_en={p:.4f}\t{'ML' if p >= thr else 'EN'}\t{text}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
