| """XLSR-Mamba network, ported from https://github.com/swagshaw/XLSR-Mamba (MIT) |
| ``model.py`` (SSLModel, Model) + ``mamba_blocks.py`` (MixerModel, create_block). |
| DPDA-GA's ``nets/mamba.py`` + ``nets/mamba_blocks.py`` are byte-identical copies. |
| |
| Module / parameter names are identical to the source so the released state_dicts load |
| strictly. Edits (inference only; see implementation-notes.md): |
| 1. SSLModel builds the fairseq wav2vec 2.0 XLS-R 300M architecture from the base |
| ``xlsr2_300m.pt`` config via merge_with_parent (cross-version fairseq fix, as in |
| benchmarks/W2V2-AASIST/_net.py) instead of load_model_ensemble_and_task; the |
| fine-tuned checkpoint then overwrites every weight. |
| 2. extract_feat: removed the ``self.model.train()`` call (it fires on a device/dtype |
| mismatch and would re-enable dropout) and the second, unused XLS-R forward |
| (``layerresult``): the source ran XLS-R twice and discarded the second result. |
| 3. Model takes the two hyper-parameters directly (emb_size, num_encoders) instead of |
| argparse ``args``; no ``device`` / print. |
| 4. Mamba mixer, Block, RMSNorm and the fused add+RMSNorm come from the dir-local |
| ``_mamba_v1.py`` (pure PyTorch; mamba_ssm v1 subset) instead of the |
| mamba_ssm CUDA / Triton kernels. Same math. |
| 5. Only the code path the released checkpoints use is kept (MixerModel with |
| if_bidirectional=True, fused_add_norm=True, rms_norm=True); BiBlock / BiMixerModel |
| (unused by Model) are dropped. The quirk of that path is preserved verbatim: the |
| final fused add+norm adds the *block input* ``hidden_states`` (not the last mixer |
| output) to each direction's residual. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| import sys |
| from dataclasses import dataclass, field |
| from functools import partial |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| if _HERE not in sys.path: |
| sys.path.insert(0, _HERE) |
|
|
| from _mamba_v1 import Block, Mamba, RMSNorm, rms_norm_fn |
|
|
| WINDOW = 66800 |
| _XLSR_PATH = os.path.join(_HERE, "xlsr2_300m.pt") |
|
|
|
|
| def pad_repeat(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray: |
| """Source utils.pad: first max_len samples; shorter clips tiled to max_len.""" |
| x = np.asarray(x, dtype=np.float32).reshape(-1) |
| n = x.shape[0] |
| if n >= max_len: |
| return x[:max_len] |
| reps = int(max_len / n) + 1 |
| return np.tile(x, reps)[:max_len].astype(np.float32) |
|
|
|
|
| def load_state(path: str) -> dict: |
| if path.endswith(".safetensors"): |
| from safetensors.torch import load_file |
| return load_file(path, device="cpu") |
| sd = torch.load(path, map_location="cpu", weights_only=True) |
| return sd.get("state_dict", sd) |
|
|
|
|
| @dataclass |
| class MambaConfig: |
| d_model: int = 64 |
| n_layer: int = 6 |
| vocab_size: int = 50277 |
| ssm_cfg: dict = field(default_factory=dict) |
| rms_norm: bool = True |
| residual_in_fp32: bool = True |
| fused_add_norm: bool = True |
| pad_vocab_size_multiple: int = 8 |
|
|
|
|
| class SSLModel(nn.Module): |
| def __init__(self, cp_path: str = _XLSR_PATH): |
| super().__init__() |
| import fairseq |
| from fairseq.dataclass.utils import merge_with_parent |
| from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model |
| state = fairseq.checkpoint_utils.load_checkpoint_to_cpu(cp_path) |
| cfg = merge_with_parent(Wav2Vec2Config(), state["cfg"]["model"], remove_missing=True) |
| ssl = Wav2Vec2Model.build_model(cfg, task=None) |
| ssl.load_state_dict(state["model"], strict=True) |
| self.model = ssl |
| self.out_dim = 1024 |
|
|
| def extract_feat(self, input_data): |
| input_tmp = input_data[:, :, 0] if input_data.ndim == 3 else input_data |
| return self.model(input_tmp, mask=False, features_only=True)["x"] |
|
|
|
|
| def create_block(d_model, ssm_cfg=None, norm_epsilon=1e-5, rms_norm=False, residual_in_fp32=False, |
| fused_add_norm=False, layer_idx=None, device=None, dtype=None): |
| ssm_cfg = ssm_cfg or {} |
| factory_kwargs = {"device": device, "dtype": dtype} |
| mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs) |
| norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs) |
| block = Block(d_model, mixer_cls, norm_cls=norm_cls, fused_add_norm=fused_add_norm, |
| residual_in_fp32=residual_in_fp32) |
| block.layer_idx = layer_idx |
| return block |
|
|
|
|
| class MixerModel(nn.Module): |
| """mamba_blocks.MixerModel, bidirectional + fused_add_norm path only.""" |
|
|
| def __init__(self, d_model, n_layer, ssm_cfg=None, norm_epsilon=1e-5, rms_norm=False, |
| fused_add_norm=False, residual_in_fp32=False): |
| super().__init__() |
| assert rms_norm and fused_add_norm, "only the released configuration is ported" |
| self.residual_in_fp32 = residual_in_fp32 |
| self.fused_add_norm = fused_add_norm |
| kw = dict(ssm_cfg=ssm_cfg, norm_epsilon=norm_epsilon, rms_norm=rms_norm, |
| residual_in_fp32=residual_in_fp32, fused_add_norm=fused_add_norm) |
| self.forward_layers = nn.ModuleList([create_block(d_model, layer_idx=i, **kw) for i in range(n_layer)]) |
| self.backward_layers = nn.ModuleList([create_block(d_model, layer_idx=i, **kw) for i in range(n_layer)]) |
| self.norm_f = RMSNorm(d_model, eps=norm_epsilon) |
| self.f_attention_pool = nn.Linear(d_model, 1) |
| self.b_attention_pool = nn.Linear(d_model, 1) |
| self.LL = nn.Linear(d_model * 2, d_model) |
| self.dropout = nn.Dropout(p=0.1) |
| self.classifier = nn.Linear(d_model, 2) |
|
|
| def forward(self, x): |
| hidden_states = self.dropout(x) |
| f_hidden_states = hidden_states |
| b_hidden_states = hidden_states.flip([1]) |
| f_residual, b_residual = None, None |
| for layer in self.forward_layers: |
| f_hidden_states, f_residual = layer(f_hidden_states, f_residual) |
| for layer in self.backward_layers: |
| b_hidden_states, b_residual = layer(b_hidden_states, b_residual) |
| |
| |
| f_hidden_states = rms_norm_fn(hidden_states, self.norm_f.weight, self.norm_f.bias, |
| eps=self.norm_f.eps, residual=f_residual, prenorm=False, |
| residual_in_fp32=self.residual_in_fp32) |
| b_hidden_states = rms_norm_fn(hidden_states, self.norm_f.weight, self.norm_f.bias, |
| eps=self.norm_f.eps, residual=b_residual, prenorm=False, |
| residual_in_fp32=self.residual_in_fp32) |
| f_hidden_states = torch.matmul(F.softmax(self.f_attention_pool(f_hidden_states), dim=1) |
| .transpose(-1, -2), f_hidden_states).squeeze(-2) |
| b_hidden_states = torch.matmul(F.softmax(self.b_attention_pool(b_hidden_states), dim=1) |
| .transpose(-1, -2), b_hidden_states).squeeze(-2) |
| hidden_states = torch.cat((f_hidden_states, b_hidden_states), dim=1) |
| hidden_states = self.LL(hidden_states) |
| hidden_states = self.dropout(hidden_states) |
| return self.classifier(hidden_states) |
|
|
|
|
| class Model(nn.Module): |
| """XLSR-Mamba: XLS-R 300M (fine-tuned) -> Linear 1024->emb -> BN+SELU -> bidirectional Mamba.""" |
|
|
| def __init__(self, emb_size: int = 144, num_encoders: int = 12, cp_path: str = _XLSR_PATH): |
| super().__init__() |
| self.ssl_model = SSLModel(cp_path) |
| self.LL = nn.Linear(1024, emb_size) |
| self.first_bn = nn.BatchNorm2d(num_features=1) |
| self.selu = nn.SELU(inplace=True) |
| self.config = MambaConfig(d_model=emb_size, n_layer=num_encoders // 2) |
| self.conformer = MixerModel(d_model=self.config.d_model, n_layer=self.config.n_layer, |
| ssm_cfg=self.config.ssm_cfg, rms_norm=self.config.rms_norm, |
| residual_in_fp32=self.config.residual_in_fp32, |
| fused_add_norm=self.config.fused_add_norm) |
|
|
| def set_ref_scan(self, on: bool) -> None: |
| """Sequential selective_scan_ref (True) vs the chunked scan (False, default).""" |
| for m in self.modules(): |
| if isinstance(m, Mamba): |
| m.use_ref_scan = on |
|
|
| def forward(self, x): |
| x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1)) |
| x = self.LL(x_ssl_feat) |
| x = x.unsqueeze(dim=1) |
| x = self.first_bn(x) |
| x = self.selu(x) |
| x = x.squeeze(dim=1) |
| return self.conformer(x) |
|
|
