File size: 2,530 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 | """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 # shape guard: leave the graph untouched
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
|