rnnt.hpp - mozsearch

Enable keyboard shortcuts

#pragma once

#include "prediction.hpp"

#include "joint.hpp"

#include "decode_types.hpp"

#include <vector>

#include <cstdint>

namespace pk {

// Standard RNN-Transducer greedy decoding (no duration head).

//

// Ports NeMo GreedyRNNTInfer._greedy_decode (rnnt_greedy_decoding.py). Drives the

// prediction net + joint frame-by-frame. Unlike TDT there are NO duration logits:

// the joint output is exactly vocab+1 (num_durations()==0 -> V_plus=vocab+1). On a

// blank the time index advances by exactly one frame; on a non-blank a symbol is

// emitted and the loop STAYS at the same frame (capped by max_symbols).

//

// Algorithm (NeMo):

//   t = 0; hyp = []; committed_state = zeros; last_token = SOS; have_token = false

//   while t < T:

//     emitted = 0

//     while emitted < max_symbols:

//       is_sos = !have_token

//       g, out_state = pred.step(last_label, is_sos, committed_state)

//       logits = joint(enc[t], g)             # raw [vocab+1]

//       k = argmax(logits)                    # token (incl. blank)

//       if k == blank: break                  # blank -> stop emitting, advance time

//       hyp.append(k); last_token = k; committed_state = out_state; have_token = true

//       emitted += 1                          # commit on non-blank, STAY at t

//     t += 1                                  # advance exactly one frame

//

// Argmax is taken over the RAW joint logits — argmax is invariant under the

// monotonic log_softmax NeMo applies for confidence, so greedy needs no softmax.

//

// pred:       stateful prediction net (carries LSTM h,c across steps).

// joint:      RNN-T joint network (called per (t, u) with T=U=1), V_plus=vocab+1.

// enc:        encoder output, row-major [T, enc_hidden] — enc[t*enc_hidden + c].

// T:          number of encoder time frames.

// enc_hidden: encoder feature dimension (= d_model).

// blank_id:   blank token id (= vocab_size); argmax range is [0, vocab+1).

// max_symbols: max symbols emitted per time frame (NeMo default 10).

//

// Returns the emitted token-id sequence (hyp). All emitted ids are < blank_id.

//

// If `tokens` is non-null it is filled (one entry per emitted id, in order) with

// per-token metadata (see TokenInfo): frame = the encoder frame t at emission,

// conf = max_prob over the full joint output vector, span = 1. The id-only path

// (tokens == nullptr) is unchanged.

std::vector<int32_t> rnnt_greedy(const PredictionNet& pred, const Joint& joint,

                                 const std::vector<float>& enc, int T, int enc_hidden,

                                 int blank_id, int max_symbols,

                                 std::vector<TokenInfo>* tokens = nullptr);

// Carried RNN-T greedy decoding state, so the per-frame loop can be driven

// INCREMENTALLY (the encoder frames arrive in chunks but the decoder must NOT

// reset across chunk boundaries). This is exactly the loop-carried state of

// rnnt_greedy lifted into a struct (NeMo partial_hypotheses): the committed

// (non-blank) prediction-net LSTM state, the last emitted label + "have we

// emitted yet" flag (== SOS handling), and the accumulated hypothesis.

//

// Feeding the encoder frames [0..T) split into chunks of any size, via repeated

// rnnt_decode_frames calls on the SAME RnntDecodeState, produces the identical

// token sequence as rnnt_greedy on the whole [0..T) at once — because the only

// thing the frame loop carries across frames is exactly this struct, and the

// per-frame inner emit-until-blank loop is independent of how the frames were

// chunked.

struct RnntDecodeState {

    PredState state;          // committed prediction-net LSTM state (h,c per layer)

    int32_t   last_token = -1; // last EMITTED token (-1 sentinel: nothing yet → SOS)

    bool      have_token = false; // have we emitted any token (false → use SOS)

    std::vector<int32_t> hyp;  // accumulated emitted token ids (across all chunks)

};

// Initialise a fresh decode state (zeroed prediction-net LSTM state, SOS).

RnntDecodeState rnnt_decode_init(const PredictionNet& pred);

// Run the RNN-T greedy frame loop over `Tnew` NEW encoder frames, carrying and

// updating `st` (so subsequent calls continue from where this one left off).

// enc_frames: row-major [Tnew, enc_hidden], enc_frames[t*enc_hidden + c].

// Appends the newly emitted token ids to st.hyp AND returns the ids emitted in

// THIS call (for the streaming "newly finalized tokens" API). The decoder state

// (LSTM, last token, SOS flag) persists in `st` across calls.

//

// If `emit_frames` is non-null it is filled (one entry per returned token) with

// the LOCAL frame index t in [0, Tnew) at which that token was emitted — the

// caller adds the running global encoder-frame offset to recover an absolute

// frame index (used for EOU/EOB event timing, matching NeMo's per-token

// timestamp = the encoder time index that produced the symbol).

//

// If `tokens` is non-null it is APPENDED to (one entry per token emitted in this

// call, same convention as `emit_frames`) with per-token metadata matching

// NeMo's timestamps=True + 'max_prob' confidence (see TokenInfo): frame = the

// LOCAL emission frame t in [0, Tnew), conf = max_prob over the full joint output

// vector (N = V_plus = vocab+1), span = 1. The id-only paths (both out-params

// null) are unchanged.

std::vector<int32_t> rnnt_decode_frames(const PredictionNet& pred, const Joint& joint,

                                        const std::vector<float>& enc_frames,

                                        int Tnew, int enc_hidden,

                                        RnntDecodeState& st,

                                        int blank_id, int max_symbols,

                                        std::vector<int32_t>* emit_frames = nullptr,

                                        std::vector<TokenInfo>* tokens = nullptr);

} // namespace pk