Source code
Revision control
Copy as Markdown
Other Tools
#pragma once
#include "model_loader.hpp"
#include <vector>
#include <cstdint>
namespace pk {
// Carries the LSTM hidden and cell state for stateful single-step decoding.
// One (h, c) pair PER stacked LSTM layer (PyTorch nn.LSTM with num_layers>1).
// h[l] / c[l] are each [hidden]. For a single-layer LSTM the outer vectors hold
// exactly one element.
struct PredState {
std::vector<std::vector<float>> h; // h[layer] = [hidden]
std::vector<std::vector<float>> c; // c[layer] = [hidden]
};
// Batched LSTM state: one (h,c) per layer, each holding N items' columns laid
// out [H*N] (item n at offset n*H). Generalizes PredState to a batch.
struct BatchedPredState {
std::vector<std::vector<float>> h; // h[layer] size H*N
std::vector<std::vector<float>> c; // c[layer] size H*N
};
// RNN-Transducer prediction network — NeMo RNNTDecoder prediction net.
//
// Architecture:
// Embedding(vocab+1, pred_hidden, padding_idx=blank) lookup
// Stacked LSTM with `pred_rnn_layers` layers (PyTorch convention). The 110m
// anchor has 1 layer; parakeet-tdt-0.6b-v2/-v3 have 2. The number of layers
// is read from the GGUF (parakeet.decoder.pred_rnn_layers); each layer l uses
// the weight_ih_l{l}/weight_hh_l{l}/bias_*_l{l} tensors, and layer l>0 takes
// the previous layer's hidden output h as its input.
//
// LSTM math (per step, input x_t [H], prev h,c [H]; h0=c0=0):
// z = W_ih · x_t + b_ih + W_hh · h + b_hh # [4H]
// i = sigmoid(z[0:H]); f = sigmoid(z[H:2H]); g = tanh(z[2H:3H]); o = sigmoid(z[3H:4H])
// c' = f * c + i * g
// h' = o * tanh(c')
// Gate order is PyTorch [input, forget, cell, output] stacked in the 4H dim.
//
// add_sos=true prepends a literal zero [H] "start of sequence" vector to the
// embedded sequence (matching NeMo predict(add_sos=True)), so the output has
// U+1 hidden states. The output is the sequence of h' states.
//
// Tensors (verbatim NeMo names; l = LSTM layer index 0..pred_rnn_layers-1):
// decoder.prediction.embed.weight ggml ne=[H, vocab+1]
// decoder.prediction.dec_rnn.lstm.weight_ih_l{l} ggml ne=[H, 4H]
// decoder.prediction.dec_rnn.lstm.weight_hh_l{l} ggml ne=[H, 4H]
// decoder.prediction.dec_rnn.lstm.bias_ih_l{l} ggml ne=[4H]
// decoder.prediction.dec_rnn.lstm.bias_hh_l{l} ggml ne=[4H]
class PredictionNet {
public:
explicit PredictionNet(const ModelLoader& ml);
// ids: input label ids (length U). Each id indexes embed.weight.
// add_sos: prepend a zero SOS step (output length becomes U+1).
// out: row-major [U_out, hidden] — out[u*hidden + h]
// U_out: number of output steps (U, or U+1 if add_sos)
// hidden: pred_hidden (H)
void forward(const std::vector<int32_t>& ids, bool add_sos,
std::vector<float>& out, int& U_out, int& hidden) const;
// Returns a zero-initialised LSTM state (h and c each [hidden] zeros).
PredState zero_state() const;
// Advance the LSTM by one token.
// token_id: embedding index (ignored when is_sos=true).
// is_sos: use the zero SOS input vector instead of the embedding.
// in: previous (h, c) state (use zero_state() for the first call).
// g: output — the new h' vector [hidden].
// out_state: the new (h', c') state to carry to the next step.
void step(int32_t token_id, bool is_sos,
const PredState& in,
std::vector<float>& g,
PredState& out_state) const;
// Advance the LSTM one token for N items at once.
// token_ids[n]: embedding index for item n (ignored where is_sos[n]).
// is_sos[n]: use the zero SOS input for item n (1=true).
// in: batched prior state (h[L],c[L] each [H*N]).
// g: OUT, top-layer h' for all items [H*N] (item n at n*H).
// out_state: OUT, new batched (h',c').
void step_batch(const std::vector<int32_t>& token_ids,
const std::vector<uint8_t>& is_sos,
const BatchedPredState& in,
std::vector<float>& g,
BatchedPredState& out_state) const;
int hidden_size() const { return H_; }
int num_layers() const { return n_layers_; }
private:
const ModelLoader& ml_;
int H_; // pred_hidden
int vocab_p1_; // vocab + 1 (embedding rows)
int n_layers_; // pred_rnn_layers (stacked LSTM layers)
// Host-side copy of the embedding table, lazily fetched on the first step()
// via ggml_backend_tensor_get (works for both CPU and device-resident
// weights). [vocab_p1_ * H_], row-major: embed_host_[id*H_ + h].
mutable std::vector<float> embed_host_;
};
} // namespace pk