Skip to content

Commit 43bcec7

Browse files
themason2011Copilot
andcommitted
Add streaming multi-talker Parakeet RNN-T ASR model
Add NVIDIA NeMo `EncDecMultiTalkerRNNTBPEModel` (multitalker-parakeet-streaming-0.6b-v1) as its own model type, exported as three ONNX graphs matching the nemotron-speech-streaming cache-aware streaming deployment contract: - encoder: FastConformer with speaker-kernel injection and streaming channel/time caches (causal dw_striding subsampling, bias-free linears, LayerNorm conv, chunked-limited attention, rel-pos KV cache). - decoder: RNN-T LSTM prediction network (2-layer ONNX LSTM with PyTorch->ONNX gate reordering). - joint: RNN-T joint network emitting per-(time, target) logits. Adds MultiTalkerRNNTTask, the ParakeetMultiTalkerModel/Config classes with a .nemo loader (build_parakeet_multitalker), unit tests (graph build + ORT smoke run of all three graphs), and an integration test verifying streaming parity against the NeMo PyTorch reference (encoder 3.9e-7, decoder 1.1e-6, joint argmax match). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Mason Corey <masoncorey@microsoft.com>
1 parent bcf8e15 commit 43bcec7

6 files changed

Lines changed: 1605 additions & 0 deletions

File tree

src/mobius/models/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
"Blip2Model",
1515
"BloomCausalLMModel",
1616
"build_sortformer",
17+
"build_parakeet_multitalker",
1718
"CLIPVisionModel",
1819
"SigLIPVisionModel",
1920
"CTRLCausalLMModel",
@@ -124,6 +125,8 @@
124125
"SmolLM3CausalLMModel",
125126
"SortformerConfig",
126127
"SortformerDiarizationModel",
128+
"ParakeetMultiTalkerConfig",
129+
"ParakeetMultiTalkerModel",
127130
"StarCoder2CausalLMModel",
128131
"T2IAdapterModel",
129132
"T5ForConditionalGeneration",
@@ -258,6 +261,11 @@
258261
)
259262
from mobius.models.sensevoice_small import SenseVoiceSmallModel
260263
from mobius.models.smollm import SmolLM3CausalLMModel
264+
from mobius.models.parakeet_multitalker import (
265+
ParakeetMultiTalkerConfig,
266+
ParakeetMultiTalkerModel,
267+
build_parakeet_multitalker,
268+
)
261269
from mobius.models.sortformer import (
262270
SortformerConfig,
263271
SortformerDiarizationModel,

0 commit comments

Comments
 (0)