Add NVIDIA NeMo Sortformer speaker-diarization model - #410
Merged
Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a new speaker-diarization export path to mobius by implementing NVIDIA NeMo’s Sortformer streaming diarization model as an ONNX graph (offline forward path), along with a new diarization task and tests to validate graph construction and NeMo parity.
Changes:
- Added
SortformerConfig,SortformerDiarizationModel, and a.nemoloader (build_sortformer) to reconstruct the model and apply NeMo weights. - Introduced
DiarizationTask(input_features → speaker_probs) and registered it under thediarizationtask name. - Added unit graph-build tests and an integration parity test against NeMo.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/sortformer_integration_test.py | Integration test comparing ONNX output to NeMo PyTorch offline inference. |
| tests/build_graph_test.py | New L1-style graph build + ORT execution smoke tests for Sortformer + DiarizationTask. |
| src/mobius/tasks/_diarization.py | New encoder-only diarization task wiring input_features to speaker_probs. |
| src/mobius/tasks/init.py | Exports/registers DiarizationTask and the diarization task mapping. |
| src/mobius/models/sortformer.py | New Sortformer model implementation + .nemo extraction/loading helper. |
| src/mobius/models/init.py | Exports Sortformer* symbols and build_sortformer. |
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
justinchuby
reviewed
Jul 17, 2026
themason2011
added a commit
that referenced
this pull request
Jul 17, 2026
Applies @justinchuby's PR #410 review feedback: - Use the opset-24 `op.Swish` op directly for Conformer feed-forward and convolution activations instead of the manual `x * sigmoid(x)` helper (removes the `_swish` function). - Replace the manual scaled-dot-product attention in `_TransformerAttention` with the opset-24 `op.Attention` op (bidirectional, unmasked). - Move the `.nemo` archive loader `build_sortformer` out of the modeling code into a dedicated NeMo integration package (`mobius.integrations.nemo`), mirroring the GGUF integration layout, and drop it from the public `mobius.models` namespace. Verified: sortformer graph unit tests pass and offline output stays bit-close to the NeMo reference (max abs diff 1.2e-7 vs golden). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <masoncorey@microsoft.com>
justinchuby
reviewed
Jul 17, 2026
Wire the NeMo Sortformer speaker-diarization model into mobius's central build_from_nemo loader instead of a bespoke build_sortformer entry point. - Register SortformerDiarizationModel under model_type 'sortformer' with the 'diarization' task in _registry.py. - Export SortformerConfig/SortformerDiarizationModel from models/__init__.py and DiarizationTask from tasks/__init__.py (TASK_REGISTRY 'diarization'). - Map the NeMo SortformerEncLabelModel target to model_type 'sortformer' and dispatch nemo_to_config to SortformerConfig.from_nemo_yaml before the FastConformer-RNNT encoder validation. - Add model_type field to SortformerConfig so the generic pipeline can resolve the model class and task from config.model_type. - Re-add graph-build unit tests and rewrite the integration test to use build_from_nemo. Verified: build_from_nemo output matches the NeMo golden reference (max abs diff 1.15e-7); 5 graph tests pass; ruff clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <masoncorey@microsoft.com>
themason2011
force-pushed
the
themason2011/sortformer-diarization
branch
from
July 20, 2026 21:30
c235538 to
ac3670a
Compare
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
|
The author of this PR, themason2011, is not an activated member of this organization on Codecov. |
justinchuby
approved these changes
Jul 21, 2026
justinchuby
left a comment
justinchuby
left a comment
Member
There was a problem hiding this comment.
Thanks! Could you address the copilot comments and ensure the unit test CI + lint CI are green?
Member
|
If possible, could you also add the L4/L5 tests? |
|
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (9)
src/mobius/models/sortformer.py:403
scaleis a float32 Constant, butmatrix_ac/matrix_bdfollow the model compute dtype (can be f16/bf16 whenbuild_from_nemo(..., dtype=...)is used). ONNXMulrequires matching dtypes, so this can produce an invalid graph for non-float32 builds. Cast the scalar to the scores dtype (the codebase typically usesop.CastLike).
scale = op.Constant(value=ir.tensor(np.array(self._head_dim**-0.5, dtype=np.float32)))
scores = op.Mul(op.Add(matrix_ac, matrix_bd), scale)
src/mobius/models/sortformer.py:486
_build_pos_emb()constructs positional embeddings in float32 (Cast(..., FLOAT)+ float32div_term), but those embeddings are fed intolinear_poswhich will use the model dtype (potentially f16/bf16 viabuild_from_nemo(dtype=...)). This dtype mismatch will breakMatMulinsideLinear. Cast the finalpos_embto the compute dtype before returning (keeping Sin/Cos in f32 if desired).
t_scalar = op.Squeeze(op.Shape(x, start=1, end=2)) # scalar T
one = op.Constant(value_int=1)
start = op.Sub(t_scalar, one) # T - 1
limit = op.Neg(t_scalar) # -T (exclusive) -> last value -(T-1)
positions = op.Range(start, limit, op.Constant(value_int=-1)) # [2T-1]
pos_f = op.Unsqueeze(op.Cast(positions, to=ir.DataType.FLOAT), [-1]) # [2T-1,1]
div_term = op.Constant(value=ir.tensor(self._div_term)) # [d/2]
angles = op.Mul(pos_f, div_term) # [2T-1, d/2]
sin = op.Unsqueeze(op.Sin(angles), [-1])
cos = op.Unsqueeze(op.Cos(angles), [-1])
interleaved = op.Concat(sin, cos, axis=-1) # [2T-1, d/2, 2]
pe = op.Reshape(interleaved, [0, -1]) # [2T-1, d]
return op.Unsqueeze(pe, [0]) # [1, 2T-1, d]
src/mobius/tasks/_diarization.py:42
feat_in if feat_in else "feat"treatsfeat_in=0as missing and falls back to a symbolic dimension. Iffeat_inis present, the check should beis not Noneto avoid surprising behavior and keep the shape contract consistent.
feat_in = getattr(config, "feat_in", None)
input_features = builder.input(
"input_features",
dtype=config.dtype,
shape=["batch", feat_in if feat_in else "feat", "time"],
)
src/mobius/models/sortformer.py:494
- The xscaling multiplier is emitted as a float32 constant, but
xis in the model compute dtype (potentially f16/bf16). This can make the graph invalid for non-float32 builds. UseCastLikeon the scalar constant (same pattern used in other models).
# Scale inputs to the attention layers by sqrt(d_model) (xscaling).
x = op.Mul(x, op.Constant(value=ir.tensor(np.array(self._xscale, np.float32))))
pos_emb = self._build_pos_emb(op, x)
src/mobius/integrations/nemo/_config_mapping.py:23
- The module-level docstring still claims this mapping only targets
ArchitectureConfig/ FastConformer-RNNT, but the module now also maps Sortformer diarization toSortformerConfig. Updating the top docstring prevents misleading documentation for NeMo support.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/models/sortformer.py:16
- The module docstring claims the subsampled time length is
T = T_mel / 8, but the dw_striding Conv2d stack uses stride=2 with padding=1/kernel=3, which yieldsceil(T_mel / 8)for odd lengths. This doc mismatch can confuse downstream users when shapes don’t divide evenly.
input_features [B, feat, T_mel]
-> transpose [B, T_mel, feat]
-> FastConformer encoder [B, T, fc_d_model] (T = T_mel / 8)
-> encoder_proj (Linear) [B, T, tf_d_model]
-> Transformer encoder (post-LN) [B, T, tf_d_model]
-> speaker sigmoid head [B, T, num_spks]
src/mobius/tasks/_diarization.py:25
- The output-frame formula in the docstring is inaccurate for the Sortformer subsampling stem (stride-2 convs with padding). The time dimension becomes
ceil(time / subsampling_factor)rather than a simple division whentimeisn’t an exact multiple.
Input: ``input_features`` — ``[batch, feat, time]`` mel spectrogram.
Output: ``speaker_probs`` — ``[batch, frames, num_spks]`` sigmoid
probabilities (``frames = time / subsampling_factor``).
"""
src/mobius/integrations/nemo/_config_mapping.py:23
- This file imports
BaseModelConfigfrom the private modulemobius._configs._base. Elsewhere (e.g.src/mobius/tasks/_diarization.py) the convention is to import config types from the publicmobius._configspackage. Using the public import keeps internal module boundaries consistent.
This issue also appears on line 22 of the same file.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/models/sortformer.py:87
- The PR description mentions a
build_sortformer(.nemo)loader implemented inmodels/sortformer.py, but there is no such function in this file (andbuild_sortformeris not found in the repo). Either add the helper (if intended API surface) or update the PR description to match the currentbuild_from_nemoentry point.
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
@dataclasses.dataclass
class SortformerConfig(BaseModelConfig):
"""Configuration for the Sortformer diarization model.
Fields mirror the NeMo ``.nemo`` ``model_config.yaml`` (``encoder``,
``transformer_encoder`` and ``sortformer_modules`` sections).
"""
# Mel feature dimension (preprocessor ``features``).
feat_in: int = 128
# FastConformer encoder.
fc_d_model: int = 512
fc_num_layers: int = 17
fc_num_heads: int = 8
fc_ff_expansion: int = 4
fc_conv_kernel: int = 9
fc_subsampling_conv_channels: int = 256
fc_subsampling_factor: int = 8
# Transformer encoder (operates on projected embeddings).
tf_d_model: int = 192
tf_num_layers: int = 18
tf_num_heads: int = 8
tf_inner_size: int = 768
tf_hidden_act: str = "relu"
# Diarization head.
num_spks: int = 4
# HuggingFace/registry model_type — consumed by the generic
# ``build_from_nemo`` pipeline to resolve the model class and task.
model_type: str | None = "sortformer"
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <36940706+themason2011@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <36940706+themason2011@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (4)
src/mobius/models/sortformer.py:489
_build_pos_emb()returnspeas FLOAT (Sin/Cos are computed in float32), but the attention path later feeds this intoLinear/MatMultensors that followconfig.dtype(f16/bf16 possible viabuild_from_nemo(..., dtype=...)). Without castingpeto the compute dtype, ONNX will reject the graph due to mismatched element types.nemo_rnnt.pyhandles this by casting the positional encoding to the model dtype before projection (seesrc/mobius/models/nemo_rnnt.py:595-599).
interleaved = op.Concat(sin, cos, axis=-1) # [2T-1, d/2, 2]
pe = op.Reshape(interleaved, [0, -1]) # [2T-1, d]
return op.Unsqueeze(pe, [0]) # [1, 2T-1, d]
src/mobius/models/sortformer.py:434
op.Constant(value_float=0.5)creates a float32 scalar. If the model is built withconfig.dtype= f16/bf16, this will cause a dtype mismatch inMul/Add(ONNX requires matching element types). The repo already uses a cast-to-ref pattern for scalars (e.g._scalar_like()insrc/mobius/models/nemo_rnnt.py:55-63). Consider casting the scalar once and reusing it for both Macaron half-step residuals.
# Macaron feed-forward (half-step residual)
ff = self.feed_forward1(op, self.norm_feed_forward1(op, x))
x = op.Add(x, op.Mul(ff, op.Constant(value_float=0.5)))
src/mobius/models/sortformer.py:497
- The xscaling constant is emitted as a float32 tensor (
np.float32) and multiplied intox. If the graph is built withconfig.dtype= f16/bf16, this introduces a float32×float16Mul, which ONNX will reject. Please cast the scalar tox's dtype (same pattern as_scalar_like()insrc/mobius/models/nemo_rnnt.py:55-63).
# Scale inputs to the attention layers by sqrt(d_model) (xscaling).
x = op.Mul(x, op.Constant(value=ir.tensor(np.array(self._xscale, np.float32))))
pos_emb = self._build_pos_emb(op, x)
src/mobius/integrations/nemo/_config_mapping.py:23
BaseModelConfigis part of the publicmobius._configsAPI (re-exported fromsrc/mobius/_configs/__init__.py:24-29), but this file imports it from the private modulemobius._configs._base. SinceArchitectureConfigis already imported frommobius._configs, it would be more consistent to importBaseModelConfigfrom the same public module to avoid leaking internal paths into other packages.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <36940706+themason2011@users.noreply.github.com>
Contributor
Author
Added L4/L5 tests |
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Signed-off-by: Mason Corey <36940706+themason2011@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/models/sortformer.py:486
_build_pos_emb()currently referencessinandcoswithout defining them, so the Sortformer graph cannot be built (NameError during graph construction). Implement the standard Transformer-XL sinusoidal relative positional embedding (arange(T-1, -T, -1) × div_term, then sin/cos interleave) before concatenation/reshape.
t_scalar = op.Squeeze(op.Shape(x, start=1, end=2)) # scalar T
one = op.Constant(value_int=1)
start = op.Sub(t_scalar, one) # T - 1
limit = op.Neg(t_scalar) # -T (exclusive) -> last value -(T-1)
# Positions: [T-1, T-2, ..., -(T-1)] -> shape [2T-1]
positions = op.Range(start, limit, op.Constant(value_int=-1))
src/mobius/integrations/nemo/_config_mapping.py:23
BaseModelConfigis imported from the private modulemobius._configs._base, but it is already re-exported frommobius._configs(and that’s how other tasks import it). Prefer the public import to avoid coupling to private module structure.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/tasks/_diarization.py:34
DiarizationTask.build()leaves themoduleparameter untyped. This repo uses MyPy strict mode and other tasks type this asonnxscript.nn.Module; leaving it untyped is likely to fail type-checking and is inconsistent with nearby task implementations.
def build(
self,
module,
config: BaseModelConfig,
) -> ModelPackage:
Co-authored-by: themason2011 <36940706+themason2011@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/models/sortformer.py:92
fc_subsampling_factoris parsed from NeMo YAML but the model implementation hardcodes 8x subsampling (3 stride-2 conv stages) and never uses this field. If a config ever has a different subsampling factor, the graph output length will be wrong while the config suggests otherwise. Add validation (or plumb the factor into_ConvSubsampling) so mismatches fail fast.
def validate(self) -> None:
if self.fc_d_model % self.fc_num_heads != 0:
raise ValueError("fc_d_model must be divisible by fc_num_heads")
if self.tf_d_model % self.tf_num_heads != 0:
raise ValueError("tf_d_model must be divisible by tf_num_heads")
src/mobius/tasks/_diarization.py:25
- The docstring claims
frames = time / subsampling_factor, but the Sortformer export uses strided conv subsampling with padding, which yieldsceil(time / subsampling_factor)in general (only exact division whentimeis a multiple of the factor). Update the docstring to avoid documenting an incorrect shape relationship.
Input: ``input_features`` — ``[batch, feat, time]`` mel spectrogram.
Output: ``speaker_probs`` — ``[batch, frames, num_spks]`` sigmoid
probabilities (``frames = time / subsampling_factor``).
tests/sortformer_integration_test.py:90
- Parsing the NPZ
metafield viajson.loads(str(golden["meta"]))is brittle becausestr(...)depends on NumPy's formatting. Use.item()to extract the scalar string reliably before JSON parsing.
meta = json.loads(str(golden["meta"]))
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.
Suppressed comments (4)
src/mobius/integrations/nemo/_config_mapping.py:23
BaseModelConfigis imported from the private module pathmobius._configs._base, while the rest of the codebase imports config types from the publicmobius._configspackage. Keeping the public import avoids leaking internal module structure and prevents accidental circular-import/regression if_baselayout changes.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/integrations/nemo/_config_mapping.py:104
- The new
model_type == "sortformer"branch innemo_to_config()is currently untested. Existing_config_mapping_test.pyonly exercises the RNNT path, so regressions in Sortformer field extraction (e.g.,feat_in, layer counts) won’t be caught by unit tests.
def nemo_to_config(nemo_config: dict[str, Any]) -> BaseModelConfig:
"""Build a mobius config from a NeMo ``model_config.yaml`` dict.
Dispatches on the NeMo ``target`` class path: FastConformer-RNNT models
produce an :class:`ArchitectureConfig`; Sortformer diarization models
produce a :class:`SortformerConfig`.
"""
target = str(nemo_config.get("target", ""))
model_type = nemo_model_type(target)
if model_type == "sortformer":
# Imported lazily to avoid a models→integrations import cycle.
from mobius.models.sortformer import SortformerConfig
return SortformerConfig.from_nemo_yaml(nemo_config)
src/mobius/tasks/_diarization.py:25
- The task docstring says
frames = time / subsampling_factor, but the Sortformer subsampling stem uses stride-2 convs with padding, which yieldsframes = ceil(time / subsampling_factor)in general. This is a user-facing contract comment, so it should match the actual shape behavior.
Input: ``input_features`` — ``[batch, feat, time]`` mel spectrogram.
Output: ``speaker_probs`` — ``[batch, frames, num_spks]`` sigmoid
probabilities (``frames = time / subsampling_factor``).
src/mobius/models/sortformer.py:17
- Top-level model docstring says
T = T_mel / 8, but the conv subsampling stem actually producesT = ceil(T_mel / 8)due to stride-2 + padding (and_ConvSubsamplingalready documents theceil). Aligning this avoids confusion for callers that feed lengths not divisible by 8.
input_features [B, feat, T_mel]
-> transpose [B, T_mel, feat]
-> FastConformer encoder [B, T, fc_d_model] (T = T_mel / 8)
-> encoder_proj (Linear) [B, T, tf_d_model]
Contributor
Author
|
Ran lintrunner with copilot |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 10 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/tasks/_diarization.py:34
DiarizationTask.build()leavesmoduleuntyped. Most tasks annotate this asmodule: nn.Module(e.g.src/mobius/tasks/_feature_extraction.py:32-36,src/mobius/tasks/_image_classification.py:30-35), which helps MyPy strict mode and keeps task signatures consistent.
def build(
self,
module,
config: BaseModelConfig,
) -> ModelPackage:
src/mobius/integrations/nemo/_config_mapping.py:104
- The new
model_type == "sortformer"branch innemo_to_config()isn't covered by a unit test (currentsrc/mobius/integrations/nemo/_config_mapping_test.pyonly exercises RNNT). Adding a minimal Sortformer config fixture + assertions would prevent silent mapping regressions (e.g., wrong key names / defaults).
if model_type == "sortformer":
# Imported lazily to avoid a models→integrations import cycle.
from mobius.models.sortformer import SortformerConfig
return SortformerConfig.from_nemo_yaml(nemo_config)
src/mobius/models/sortformer.py:679
SortformerDiarizationModel.preprocess_weights()is new, but Sortformer isn't currently included in thetests/weight_alignment_test.pyparametrization (it only usesALL_CAUSAL_LM_CONFIGS,ENCODER_CONFIGS,VISION_CONFIGS, etc.; no speech/diarization configs and nosortformerentry intests/_test_configs.py). This means the standard Tier-1 weight-alignment coverage won't catch accidental key drops/renames here.
def preprocess_weights(
self, state_dict: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]:
"""Map NeMo checkpoint names onto module parameter names.
NeMo prefixes align 1:1 with this module tree (``encoder.*``,
``transformer_encoder.*``, ``sortformer_modules.*``). Only the mel
preprocessor, batch-norm ``num_batches_tracked`` counters and the
unused streaming ``hidden_to_spks`` head are dropped.
"""
renamed: dict[str, torch.Tensor] = {}
for key, value in state_dict.items():
if key.startswith("preprocessor."):
continue
if key.endswith("num_batches_tracked"):
continue
# Unused in the offline forward path (streaming FIFO head).
if key.startswith("sortformer_modules.hidden_to_spks."):
continue
renamed[key] = value
The model_coverage_test enforces that every registered model_type has either a shared test config (tests/_test_configs.py + _registry.py test_model_id) or an entry in _COVERAGE_SKIP. Sortformer is a NeMo .nemo diarization model with a custom SortformerConfig and no HuggingFace config.json, so it uses the same mechanism as the sibling fastconformer_rnnt model: skip the generic coverage and rely on the dedicated tests/sortformer_integration_test.py (L4/L5 golden) plus TestBuildGraphSortformer (L1). Fixes 4 CI failures in tests/model_coverage_test.py for model_type 'sortformer'. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <masoncorey@microsoft.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.
Suppressed comments (2)
src/mobius/integrations/nemo/_config_mapping.py:23
- Importing BaseModelConfig from the private module
mobius._configs._baseis inconsistent with the rest of the codebase (tasks and other call sites import it from the publicmobius._configs). Using the public import avoids coupling to internal module layout.
from mobius._configs import ArchitectureConfig
from mobius._configs._base import BaseModelConfig
src/mobius/tasks/_diarization.py:25
- The docstring states
frames = time / subsampling_factor, but Sortformer’s dw_striding subsampling uses stride-2, pad-1 convs, which yieldceil(time / subsampling_factor)for odd-length inputs. This is important for downstream consumers that rely on the exact frame count behavior.
Input: ``input_features`` — ``[batch, feat, time]`` mel spectrogram.
Output: ``speaker_probs`` — ``[batch, frames, num_spks]`` sigmoid
probabilities (``frames = time / subsampling_factor``).
justinchuby
approved these changes
Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds support for the NVIDIA NeMo Sortformer streaming speaker-diarization model (
nvidia/diar_streaming_sortformer_4spk-v2.1) to mobius.This is not a HuggingFace
transformersarchitecture — it ships as a NeMo.nemocheckpoint with noconfig.json. The architecture was reconstructed from the checkpoint and NeMo source. The exported ONNX graph implements the offline forward path (frontend_encoder+forward_infer): mel-spectrogram features in, per-frame speaker-activity probabilities out.Architecture
ConformerEncoder): rel_pos Transformer-XL attention, 8xdw_stridingconv subsampling, Macaron feed-forward blocks, batch-norm convolution module, xscaling.relu -> Linear -> relu -> Linear -> sigmoid).Module attribute names mirror the NeMo weight prefixes, so
preprocess_weightsonly drops unused keys (mel preprocessor,num_batches_tracked, the streaming FIFO head). The mel preprocessor (STFT) is left out of the graph, matching the mobius audio convention.Changes
models/sortformer.py—SortformerConfig(+from_nemo_yaml),SortformerDiarizationModel, and abuild_sortformer(.nemo)loader that extracts the archive, builds the graph, and applies checkpoint weights.tasks/_diarization.py—DiarizationTask(input_features -> speaker_probs), registered as thediarizationtask.models/__init__.pyand the task registry.Testing
tests/build_graph_test.py::TestBuildGraphSortformer): tiny-config graph build, I/O contract, initializers, task lookup, and an ORT forward run. All pass.tests/sortformer_integration_test.py,-m integration): compares the ONNX output against the real NeMo PyTorch reference. Offline output matches to 1.2e-7 (float32).Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com