Nemotron ASR Support for Streaming

[nenad1002]

[Copilot]

Pull request overview

Copilot reviewed 28 out of 28 changed files in this pull request and generated 9 comments.

---

You can also share your feedback on Copilot code review. Take the survey.

[nenad1002]

@kunal-vaishnavi can we sign this off?

[kunal-vaishnavi]

@kunal-vaishnavi can we sign this off?

There are a few issues with the ONNX models that may change the input/output structure.

Decoder

From inspecting the graph, target_length_orig and target_length are not used at all. This is an identity operation to go from target_length_orig to target_length. The input and the output can be removed from the graph.

Joiner

This is the end of the encoder and the decoder. The decoder contains a Squeeze op with axes = [1] before the final Transpose op. The Transpose op has perm = [1, 2, 0] so this means that its input (i.e. the output from the Squeeze op) must be 3D then. Because the axes are specified on the Squeeze op, the operation is only applied on those axes. This implies that the input to the Squeeze op (i.e. the output from the LSTM op) is 4D. Therefore, the Squeeze op is removing a dimension of size 1 to go from [A, 1, B, C] to [A, B, C].

Here is the joiner model. The joiner is receiving 3D inputs from both the encoder and decoder and returns 4D outputs. The decoder input is 3D and the Unsqueeze op in the decoder path has axes = [1]. This means that the Unsqueeze op is going from [A, B, C] to [A, 1, B, C].

That shape of [A, 1, B, C] is already obtained in the decoder. If the Squeeze op is removed and the Transpose op is changed to handle a 4D output, then the output of the decoder can be 4D. The decoder input to the joiner model can be 4D and then the Unsqueeze op in the decoder path can be removed. These changes can eliminate two ops across two ONNX models.

If we want to keep the input dimensions for the joiner model as the same, then we can move the Unsqueeze op in the encoder path to the output of the encoder. Then, the joiner model will take in two 4D inputs and produce one 4D output.

[nenad1002]

@kunal-vaishnavi can we sign this off?

There are a few issues with the ONNX models that may change the input/output structure.

Decoder

From inspecting the graph, `target_length_orig` and `target_length` are not used at all. This is an identity operation to go from `target_length_orig` to `target_length`. The input and the output can be removed from the graph.

Joiner

This is the end of the encoder and the decoder. The decoder contains a Squeeze op with axes = [1] before the final Transpose op. The Transpose op has perm = [1, 2, 0] so this means that its input (i.e. the output from the Squeeze op) must be 3D then. Because the axes are specified on the Squeeze op, the operation is only applied on those axes. This implies that the input to the Squeeze op (i.e. the output from the LSTM op) is 4D. Therefore, the Squeeze op is removing a dimension of size 1 to go from [A, 1, B, C] to [A, B, C].

Here is the joiner model. The joiner is receiving 3D inputs from both the encoder and decoder and returns 4D outputs. The decoder input is 3D and the `Unsqueeze` op in the decoder path has `axes = [1]`. This means that the `Unsqueeze` op is going from `[A, B, C]` to `[A, 1, B, C]`.

That shape of [A, 1, B, C] is already obtained in the decoder. If the Squeeze op is removed and the Transpose op is changed to handle a 4D output, then the output of the decoder can be 4D. The decoder input to the joiner model can be 4D and then the Unsqueeze op in the decoder path can be removed. These changes can eliminate two ops across two ONNX models.

If we want to keep the input dimensions for the joiner model as the same, then we can move the Unsqueeze op in the encoder path to the output of the encoder. Then, the joiner model will take in two 4D inputs and produce one 4D output.

Right, we are aware for the length, still waiting for the newest model, this is not related to this PR. I assume this PR will deal with ONNX model finalization and if there is any model signature change:
#2010

I can wait for the team's changes though if you prefer.