| """Attach CAFE-25lang's language-conditional adapters to a restored NeMo model. |
| |
| The .nemo export carries the trunk only, which is the English path. Every other language needs the |
| adapters from `adapters.safetensors` re-attached, which this helper does with residual forward |
| hooks so the checkpoint itself is never modified. |
| |
| Used by eval_mls8.py (`--adapters ...`) and by the official Open ASR Leaderboard runners through |
| run_leaderboard.py. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
|
|
| ENCODER_LAYERS = (3, 7, 11, 15, 19, 23) |
|
|
|
|
| class Bottleneck(nn.Module): |
| """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, verbose: bool = True): |
| """Inject the adapters; returns the hook handles so they can be removed again.""" |
| 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(ENCODER_LAYERS): |
| raise SystemExit(f"{len(idxs)} adapters for {len(ENCODER_LAYERS)} encoder layers") |
| handles, mods = [], [] |
| for i, layer in zip(idxs, ENCODER_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 |
| if verbose: |
| print(f"[adapters] attached at encoder layers {list(ENCODER_LAYERS)} " |
| f"(dim={dim}, rank={rank})") |
| return handles |
|
|
