| |
| """CAFE-25lang — reference inference (language supplied by the caller). |
| |
| This model routes by language. English is transcribed by the bare trunk; every other language |
| additionally needs the language-conditional adapters, which are shipped alongside the checkpoint |
| in `adapters.safetensors` and injected at inference time. |
| |
| python inference.py --lang en audio_en.wav |
| python inference.py --lang other audio_de.wav audio_pt.wav |
| |
| `--lang en` runs the trunk on its own; `--lang other` injects the adapters. Passing the wrong one |
| does not fail, it just costs accuracy, so route deliberately. If you do not know the language of |
| the audio, use `inference_gated.py`, which decides per utterance. |
| |
| Requirements: nemo_toolkit[asr], torch, safetensors. Audio must be 16 kHz mono. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import os |
|
|
| import torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| ADAPTER_LAYERS = (3, 7, 11, 15, 19, 23) |
|
|
|
|
| class Bottleneck(nn.Module): |
| """Residual adapter: x -> x + up(gelu(down(x))).""" |
|
|
| 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))) |
|
|
|
|
| def attach_adapters(model, path: str, device, dtype=torch.bfloat16): |
| """Inject the language-conditional adapters as residual forward hooks. |
| |
| The .nemo export carries the trunk only, so the adapters have to be re-attached here; without |
| them the model is on its English path. |
| """ |
| sd = load_file(path) |
| idxs = sorted({int(k.split(".")[1]) for k in sd}) |
| rank, dim = sd["adapter.0.down.weight"].shape |
| if len(idxs) != len(ADAPTER_LAYERS): |
| raise SystemExit(f"{len(idxs)} adapters for {len(ADAPTER_LAYERS)} layers") |
| handles, mods = [], [] |
| for i, layer in zip(idxs, ADAPTER_LAYERS): |
| 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"]}) |
| a = a.to(device=device, dtype=dtype).eval() |
| mods.append(a) |
|
|
| 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))) |
| |
| model._adapters, model._adapter_handles = mods, handles |
| print(f"[CAFE] adapters injected at encoder layers {list(ADAPTER_LAYERS)} " |
| f"(dim={dim}, rank={rank})") |
|
|
|
|
| def safe_decoding(model, cuda_graph: bool = False): |
| """Greedy batched decoding with the CUDA-graph path off by default. |
| |
| The graph decoder is only a throughput optimisation and is sensitive to driver / NeMo |
| versions; transcripts are identical either way. |
| """ |
| 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) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("audio", nargs="+", help="16 kHz mono wav files") |
| ap.add_argument("--lang", choices=["en", "other"], required=True, |
| help="'en' = bare trunk; 'other' = inject the adapters") |
| 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("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| ap.add_argument("--batch_size", type=int, default=8) |
| 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) |
| if a.lang == "other": |
| attach_adapters(model, a.adapters, device, dtype) |
| else: |
| print("[CAFE] English path: bare trunk, adapters not attached") |
|
|
| with torch.inference_mode(): |
| out = model.transcribe(list(a.audio), batch_size=a.batch_size, verbose=False) |
| if isinstance(out, tuple): |
| out = out[0] |
| for f, h in zip(a.audio, out): |
| print(f"{f}\t{getattr(h, 'text', h)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
