Add streaming multi-talker Parakeet RNN-T ASR model - #411
Open
themason2011 wants to merge 1 commit into
Open
Conversation
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. |
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>
themason2011
force-pushed
the
themason2011/parakeet-multitalker
branch
from
July 17, 2026 20:36
43bcec7 to
0816c51
Compare
themason2011
changed the base branch from
themason2011/sortformer-diarization
to
main
July 17, 2026 20:36
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 NVIDIA NeMo's
EncDecMultiTalkerRNNTBPEModel(multitalker-parakeet-streaming-0.6b-v1) as its own model type. It is a cache-aware streaming multi-talker RNN-T ASR model derived fromnemotron-speech-streaming, exported as three ONNX graphs matching that deployment contract:dw_stridingsubsampling (asymmetric pad, 8x reduction), bias-free linears, LayerNorm convolution, chunked-limited attention, and a relative-position KV cache.LSTMwith PyTorch->ONNX gate reordering(i,f,g,o)->(i,o,f,g)).[B, T, U, V]logits.New public API:
ParakeetMultiTalkerModel/ParakeetMultiTalkerConfigwith a.nemoloaderbuild_parakeet_multitalker, and aMultiTalkerRNNTTask(registered asmultitalker-rnnt).Testing
tests/build_graph_test.py::TestBuildGraphParakeetMultiTalker): builds all three graphs from a tiny config and runs each through ORT with random weights; verifies I/O names/shapes and the streaming length math.tests/parakeet_multitalker_integration_test.py,-m integration): streaming parity vs the NeMo PyTorch reference — encoder3.9e-7, caches2e-6/1.8e-4, decoder1.1e-6, joint log-softmax1.5e-5with matching argmax. Verified end-to-end against real weights.ruff check/ruff formatclean.Note
Independent of #410 — the small FastConformer
_ReLU/_swishprimitives are defined locally inmodels/parakeet_multitalker.py, so this PR has no code or git dependency on the Sortformer branch and targetsmaindirectly.