mstar.model.nemotron_duplex.components.rnnt#

RNN-T transcription head (stt_model.rnnt_decoder.* / rnnt_joint.*).

The streaming user-transcription branch: a 2-layer LSTM prediction network and a joint network combining encoder + prediction features. Ported to exact checkpoint parity; the LSTM params are declared by hand (matching nn.LSTM’s weight_ih_l{n} / weight_hh_l{n} / bias_* names) to avoid meta-device init issues.

Forward math mirrors NeMo RNNTDecoder.predict / RNNTJoint.joint (branch nemotron-labs-voicechat, nemo/collections/asr/modules/rnnt.py):

predict: embed(y) -> optional zero SOS prepend -> nn.LSTM (time-first, batch_first=False, plain LSTM since normalization_mode is None) -> prediction features g [B, U(+1), H] plus LSTM (h, c) state.

joint: enc(f) [B, T, H] and pred(g) [B, U, H] broadcast-added as enc(f)[:, :, None, :] + pred(g)[:, None, :, :] -> joint_net (ReLU -> Dropout(id in eval) -> Linear 640->1025) -> raw logits [B, T, U, 1025]. log_softmax is None in the checkpoint config, so on CUDA no log-softmax is applied (raw logits are returned); argmax-equivalent either way.

Classes

`RnntDecoder`([config])
`RnntJoint`([config])

class mstar.model.nemotron_duplex.components.rnnt.RnntDecoder(config=None)[source]#

Bases: Module

Parameters:: config (RnntConfig | None)

forward(y, state=None, add_sos=True, batch_size=None)#

Prediction-network forward (NeMo RNNTDecoder.predict).

Parameters:

y – token ids [B, U] (long) or None (zero pad-token embedding).
state – optional (h, c) each [L, B, H]; None -> zeros.
add_sos (bool) – prepend a zero “start of sequence” step -> output [B, U+1, H].
batch_size (int | None) – batch size when y and state are both None.

Returns:

(g, (h, c)) with g shape [B, U(+1), H] and state [L, B, H].

load_weights(weights)[source]#

predict(y, state=None, add_sos=True, batch_size=None)[source]#

Prediction-network forward (NeMo RNNTDecoder.predict).

Parameters:

y – token ids [B, U] (long) or None (zero pad-token embedding).
state – optional (h, c) each [L, B, H]; None -> zeros.
add_sos (bool) – prepend a zero “start of sequence” step -> output [B, U+1, H].
batch_size (int | None) – batch size when y and state are both None.

Returns:

(g, (h, c)) with g shape [B, U(+1), H] and state [L, B, H].

static remap(name)[source]#

Parameters:: name (str)
Return type:: str | None

class mstar.model.nemotron_duplex.components.rnnt.RnntJoint(config=None)[source]#

Bases: Module

Parameters:: config (RnntConfig | None)

forward(encoder_outputs, decoder_outputs)[source]#

NeMo RNNTJoint.forward (non-fused logits path).

encoder_outputs [B, D=1024, T] and decoder_outputs [B, D=640, U] (time-last, as produced by the conformer / decoder); transposed to [B, T, D] / [B, U, D] then jointed -> logits [B, T, U, 1025].