File size: 5,199 Bytes
f82a6c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | #!/usr/bin/env python3
"""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 # shape guard, never corrupt the graph
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)))
# keep references alive and verifiable
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()
|