transcribe-cpp-sys 0.4.1

// arch/parakeet/parakeet.h - Parakeet-family internal model and context
// types: the concrete classes deriving from transcribe_model /
// transcribe_session. Internal to the Parakeet sources (also visible to
// tests/transcribe_tokenizer_smoke via the test target's PRIVATE include
// path). The central dispatcher talks to the family only through the Arch
// trait and the base classes.

#pragma once

#include "../sortformer/sortformer.h"  // embedded diarizer (multitalker bundle)
#include "decoder.h"
#include "transcribe-backend.h"
#include "transcribe-load-common.h"
#include "transcribe-mel.h"
#include "transcribe-model.h"
#include "transcribe-session.h"
#include "transcribe-tokenizer.h"
#include "weights.h"

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>

struct ggml_context;
struct ggml_cgraph;
struct ggml_tensor;
struct ggml_backend;
struct ggml_backend_buffer;
struct ggml_backend_sched;
typedef struct ggml_backend *        ggml_backend_t;
typedef struct ggml_backend_buffer * ggml_backend_buffer_t;
typedef struct ggml_backend_sched *  ggml_backend_sched_t;

namespace transcribe::parakeet {

// Output length of one stride-2, kernel-3 pre-encode convolution.
inline int pre_encode_time_out(int input_frames, bool causal) {
    return causal ? (input_frames / 2 + 1) : ((input_frames - 1) / 2 + 1);
}

// Host-side attention-mask helpers for the chunked-attention paths. Declared
// here so the per-mask unit tests can call them without a ggml backend
// or a GGUF on disk.
//
// chunked_limited_with_rc mask (buffered streaming). Mirrors NeMo
// ConformerEncoder._create_masks:
//   c_q = q / C
//   window_start = max(0, c_q * C - L)
//   window_end   = min(T - 1, c_q * C + C - 1 + R)
//   mask[q, k]   = 0 if window_start <= k <= window_end else -INF
// Output is row-major [T_q, T_k] with T_q == T_k == T (>= T*T floats);
// broadcasts over heads in rel_pos_mhsa.
//
// `pad_length` (optional) folds in NeMo's conv-overhang pad_mask: cells
// where q >= pad_length OR k >= pad_length are masked regardless of the
// window. pad_length = T disables it (offline). For buffered streaming
// pass effective_T = ctx.total / encoder_frame2audio_samples — the conv
// frontend may produce T_enc > effective_T, and those trailing frames
// must be masked or they contaminate valid frames' attention scores.
void compute_chunked_limited_with_rc_mask(float * out_buf,
                                          int     T,
                                          int     left_context_frames,
                                          int     chunk_size_frames,
                                          int     right_context_frames,
                                          int     pad_length);

// Bounded offline ChunkedLimited mask. The dense mask has T*T cells, but
// every query chunk only attends to `left_chunks` complete chunks on its
// left plus itself. This helper emits the equivalent compact layout used by
// rel_pos_mhsa: [window, chunk_size, 1, n_chunks], contiguous in `window`,
// where window = (left_chunks + 1)*chunk_size and
// n_chunks = ceil(T_alloc/chunk_size). T is the utterance's valid length and
// T_alloc >= T the graph's padded length (the batch T_max; T itself when
// unpadded): keys and queries at or past T are the utterance's padded tail
// and are -INF. Padded query rows retain one finite cell to avoid an
// all-masked softmax; those rows never reach the residual path.
void compute_chunked_limited_window_mask(float * out_buf, int T, int T_alloc, int chunk_size, int left_chunks);

// Bounded Regular-local mask (NeMo rel_pos_local_attn / att_context
// [left, right]): query q attends to keys k with q - left <= k <= q + right.
// Same compact layout and T / T_alloc contract as above with
// window = left + chunk_size + right and block n's keys starting at
// n*chunk_size - left.
void compute_local_window_mask(float * out_buf, int T, int T_alloc, int chunk_size, int left, int right);

// Family defaults — applied before transcribe::read_capability_kv runs
// (KV present overrides, KV absent leaves the default). Defined in
// capabilities.cpp. There is no per-variant resolver: variant identity
// is descriptive metadata; v2/v3 differences are expressed as
// stt.capability.* / general.languages KV.
void apply_family_invariants(transcribe_model & model);

// Concrete model. Owns ctx_meta, which holds every weight tensor's
// data buffer; the destructor frees it, invalidating every borrowed
// ggml_tensor* in `weights`.
struct ParakeetModel final : public transcribe_model {
    Tokenizer       tok;
    ParakeetHParams hparams;
    ParakeetWeights weights;
    ggml_context *  ctx_meta = nullptr;

    // Runtime backend plan, resolved at load() via
    // load_common::init_backends (see transcribe-backend.h). plan's
    // scheduler_list owns the backends; the destructor frees them in
    // reverse order.
    transcribe::BackendPlan                plan;
    transcribe::load_common::WeightStorage weight_storage;

    // Fused BN parameters live in a separate ggml context + buffer,
    // computed at load time from the raw BN tensors. Freed in dtor.
    ggml_context *        bn_fused_ctx    = nullptr;
    ggml_backend_buffer_t bn_fused_buffer = nullptr;

    // On a CPU primary backend, the conformer 1×1 pointwise conv weights
    // are dequantized from F16 to F32 at load: CPUs without native F16
    // arithmetic (Zen 2 and earlier) pay a per-matmul F16→F32 upconvert
    // that erases the bandwidth win, so promoting once at load is cheaper.
    // No-op on GPU backends and on non-CPU primary backends. The F16
    // source tensors stay resident (ggml can't free individual slots);
    // the weight slot is repointed at the F32 copy.
    ggml_context *        conv_pw_f32_ctx    = nullptr;
    ggml_backend_buffer_t conv_pw_f32_buffer = nullptr;

    // Mel front-end, constructed once at load() from the hparams
    // (precomputes Hann window + Slaney mel filterbank).
    // const-after-ctor and thread-safe, so shared by every context.
    // std::optional because MelFrontend has no default constructor.
    std::optional<transcribe::MelFrontend> mel;

    // Host mirror of the predictor + joint weights, populated at load()
    // via build_host_decoder_weights. The decoder runs on host (see
    // decoder.h); the const-after-load mirror lets every context share it
    // without a per-run backend readback.
    HostDecoderWeights host_decoder;

    // Embedded Sortformer diarizer (multitalker bundle GGUFs). Present iff
    // stt.parakeet.diarizer.embedded = true in the GGUF. The weight slots
    // are borrowed views into this model's ctx_meta (resolved under the
    // bundle's tensor prefix); the diarizer conformer's fused-BN params
    // live in their own ctx + buffer, freed in the dtor.
    std::unique_ptr<transcribe::sortformer::SortformerEmbedded> diar;
    ggml_context *                                              diar_bn_ctx    = nullptr;
    ggml_backend_buffer_t                                       diar_bn_buffer = nullptr;

    // Diarizer mel front-end. Same acoustic settings as the ASR frontend
    // (the compose step enforces agreement) but with the sortformer port's
    // ceil framing, so the embedded forward sees exactly what the
    // standalone sortformer port validated against.
    std::optional<transcribe::MelFrontend> diar_mel;

    ParakeetModel() = default;
    ~ParakeetModel() override;

    const transcribe::Tokenizer * tokenizer() const override { return &tok; }
};

// Per-context streaming encoder state, allocated lazily on the first
// stream_begin for ChunkedLimited variants. Mirrors NeMo's
// get_initial_cache_state tuple. Layout cheat sheet:
//   cache_last_channel  [d_model, T_cache, n_layer] f32 — per-layer
//     post-attn-LN cache, concat-prepended to x_norm before Q/K/V.
//   cache_last_time     [k_minus_1, d_model, n_layer] f32 — per-layer
//     post-pw1+GLU cache, replaces the depthwise conv's left zero-pad.
//   cache_last_channel_len  scalar — valid frames (0..T_cache); drives
//     the streaming-mask offset until the cache fills.
// Sizes for nemotron-speech-streaming-en-0.6b at [70,13]: T_cache=70,
// k_minus_1=8, n_layer=24, d_model=1024 (~7.5 MB total).
struct ParakeetStreamingCaches {
    ggml_context *        ctx    = nullptr;
    ggml_backend_buffer_t buffer = nullptr;

    // Per-layer cache tensors. Shapes:
    //   last_channel[i] = ne[d_model, T_cache, 1, 1]
    //   last_time[i]    = ne[k_minus_1, d_model, 1, 1]
    std::vector<ggml_tensor *> last_channel;
    std::vector<ggml_tensor *> last_time;

    // KV-cache variant of last_channel: per-layer pre-projected K/V
    // (bias included) over the T_cache window, ne[d_model, T_cache, 1, 1]
    // each. Equivalent to recomputing K/V from last_channel every chunk
    // but the per-chunk projections run on T_q_new columns instead of
    // T_cache + T_q_new. The default graph consumes these; an active dump
    // harness falls back to the channel recompute path. Cleared at
    // stream_begin.
    std::vector<ggml_tensor *> last_k;
    std::vector<ggml_tensor *> last_v;

    // Rel-pos projection memoization: per-layer
    // [head_dim, pos_proj_len, n_head, 1] f32 tensors holding
    // attn_pos_w @ pos_emb. Geometry-keyed (pos_proj_len), so it survives
    // across chunks and streams; rebuilt by ensure_pos_proj_cache on a
    // pos_len change. Own ctx + buffer (size varies with geometry).
    ggml_context *             pos_proj_ctx = nullptr;
    ggml_backend_buffer_t      pos_proj_buf = nullptr;
    std::vector<ggml_tensor *> pos_proj;
    int                        pos_proj_len = -1;

    int32_t channel_len = 0;

    // Absolute mel-frame cursor across the whole stream, anchoring
    // token/segment timestamps to original-audio time even when chunks
    // vary in size. Counts mel frames consumed (minus history prepend).
    int64_t mel_frames_consumed = 0;

    // Absolute sample index of stream_pcm_buffer[0]. The buffer is a
    // sliding window: prefix samples no longer needed by any unemitted
    // frame's STFT window are trimmed (hop-aligned), bounding per-feed mel
    // cost to O(new audio). Always a hop_length multiple so mel-frame
    // indexing translates cleanly between absolute and buffer-relative.
    int64_t pcm_start_sample = 0;

    // Streaming geometry, resolved at stream_begin from ParakeetHParams +
    // the caller-selected att_context_right. Constant across chunks within
    // a stream. First- vs subsequent-chunk geometry differ (NeMo's
    // setup_streaming_params); the driver picks the pair via
    // is_first_chunk. For nemotron-speech-streaming-en at R=13:
    //   chunk_size_first=105, chunk_size_subsequent=112,
    //   mel_fed_first=105, mel_fed_subsequent=121 (9 cache + 112 new),
    //   drop_extra_first=0, drop_extra_subsequent=2.
    int  att_context_right     = 0;
    int  att_context_left      = 0;
    int  chunk_size_first      = 0;
    int  chunk_size_subsequent = 0;
    int  mel_fed_first         = 0;
    int  mel_fed_subsequent    = 0;
    int  drop_extra_first      = 0;
    int  drop_extra_subsequent = 0;
    // Whether the next emit_streaming_chunk is the FIRST in this stream.
    bool is_first_chunk        = true;

    // Per-stream chunk counter (NeMo's step_num). Resets at stream_begin,
    // increments per emit; indexes dump filenames.
    int chunk_step = 0;

    bool initialized = false;
};

// Per-context streaming RNN-T decoder state, owned on the context so it
// carries across stream_feed calls. Reset at each stream_begin to a fresh
// start-of-sequence state.
struct ParakeetStreamingDecoderState {
    // LSTM (h, c) per layer; sized at stream_begin.
    LstmState lstm_state;

    // Last emitted non-blank token id (predictor input); -1 = SOS.
    int32_t prev_token_id = -1;

    // Absolute encoder-frame offset for converting per-chunk step_at_emit
    // into stream-wide frame indices.
    int64_t frame_offset = 0;

    bool initialized = false;
};

// Scheduler-allocated steady-state streaming graph. ggml graphs are
// multi-use for computation once allocated; retaining the fixed-geometry
// graph avoids rebuilding and re-allocating the same encoder topology for
// every chunk. All tensor handles are borrowed from compute_ctx and become
// invalid together when reset() is called before that context is freed.
struct ParakeetStreamingGraph {
    ggml_cgraph * graph = nullptr;

    ggml_tensor * mel_in            = nullptr;
    ggml_tensor * pos_emb_in        = nullptr;
    ggml_tensor * chunked_mask_in   = nullptr;
    ggml_tensor * prompt_one_hot_in = nullptr;
    ggml_tensor * spk_mask_in       = nullptr;
    ggml_tensor * bg_mask_in        = nullptr;
    ggml_tensor * out               = nullptr;

    std::vector<ggml_tensor *> channel_out;
    std::vector<ggml_tensor *> time_out;
    std::vector<ggml_tensor *> k_out;
    std::vector<ggml_tensor *> v_out;

    int  n_mel_chunk_frames     = -1;
    int  drop_extra_pre_encoded = -1;
    int  pos_proj_len           = -1;
    int  kv_type                = -1;
    bool spk_supervision        = false;
    bool ready                  = false;

    void reset() {
        graph                  = nullptr;
        mel_in                 = nullptr;
        pos_emb_in             = nullptr;
        chunked_mask_in        = nullptr;
        prompt_one_hot_in      = nullptr;
        spk_mask_in            = nullptr;
        bg_mask_in             = nullptr;
        out                    = nullptr;
        n_mel_chunk_frames     = -1;
        drop_extra_pre_encoded = -1;
        pos_proj_len           = -1;
        kv_type                = -1;
        spk_supervision        = false;
        ready                  = false;
        channel_out.clear();
        time_out.clear();
        k_out.clear();
        v_out.clear();
    }
};

// Concrete context. The per-call compute context and the multi-backend
// scheduler that dispatches encoder graph ops to the best available backend
// are owned by the transcribe_session base (sched / compute_ctx).
struct ParakeetSession final : public transcribe_session {
    // Encoder forward output, borrowed into compute_ctx; invalidated when
    // compute_ctx is reset next run().
    ggml_tensor * encoder_out = nullptr;

    // Per-context scratch reused across runs.
    std::vector<float>    mel_buf;
    std::vector<float>    pos_buf;
    std::vector<float>    pos_div_term;
    std::vector<float>    enc_host;
    std::vector<TdtToken> raw_tokens;

    // Multitalker bundle scratch: streaming-diarizer AOSC/FIFO state for
    // the embedded sortformer forward (diarize=ON runs only).
    transcribe::sortformer::DiarStreamScratch diar_scratch;

    // Per-call timings and the TDT decode result (tokens / words /
    // segments / full_text / result_kind / has_result) live on the base
    // transcribe_session; no per-family copies here.

    // ---- streaming-of-whole state ----
    //
    // stream_begin captures run_params; stream_feed appends PCM;
    // stream_finalize drains the buffer through the inference helper;
    // stream_reset clears (keeping capacity). NOT touched by
    // clear_result (the family owns its per-utterance audio scratch).
    std::vector<float>    stream_pcm_buffer;
    transcribe_run_params stream_run_params{};
    // Convert from cumulative samples so odd feed sizes do not lose fractions.
    int64_t               stream_audio_input_samples = 0;

    // ---- incremental streaming state (cache-aware) ----
    //
    // Allocated on the first stream_begin for ChunkedLimited variants,
    // zeroed at each stream_begin, persists across feed/finalize within
    // an utterance, freed in the dtor.
    ParakeetStreamingCaches       stream_caches;
    ParakeetStreamingDecoderState stream_dec_state;
    ParakeetStreamingGraph        stream_graph;

    // ---- Buffered streaming state (parakeet-unified-en-0.6b) ----
    //
    // Mirrors NeMo's StreamingBatchedAudioBuffer + reference inference
    // loop. Variable-stride: step 0 consumes samples_chunk + samples_right
    // of new audio; subsequent feeds consume samples_chunk. Finalize folds
    // retained right context into the decoded chunk and appends right-context
    // silence without advancing the real-audio cursor. The RNN-T state rides
    // on stream_dec_state (same predictor/joint path, no per-layer encoder cache).
    int32_t buf_left_frames     = 0;  // L (expected)
    int32_t buf_chunk_frames    = 0;  // C
    int32_t buf_right_frames    = 0;  // R
    int32_t buf_samples_left    = 0;  // L * encoder_frame2audio_samples
    int32_t buf_samples_chunk   = 0;
    int32_t buf_samples_right   = 0;
    // Next sample index in stream_pcm_buffer to feed (= NeMo's
    // `left_sample` at the start of the next step). 0 at stream_begin;
    // advances by num_new_samples after each emit.
    int64_t buf_next_audio_read = 0;
    // Buffer's internal ContextSize. (0,0,0) at stream_begin; ramps up
    // as audio arrives until it saturates at (samples_left, chunk, right).
    int64_t buf_ctx_left        = 0;
    int64_t buf_ctx_chunk       = 0;
    int64_t buf_ctx_right       = 0;
    // Whether step 0 (the initial fill of samples_chunk+samples_right)
    // has run. False at stream_begin; first emit flips it.
    bool    buf_initialized     = false;
    // Per-stream chunk step counter (== NeMo's `step_num`). Used for
    // per-chunk dump file naming.
    int32_t buf_chunk_step      = 0;
    bool    buf_active          = false;

    ParakeetSession() = default;
    ~ParakeetSession() override;
    void on_scratch_released() noexcept override;
};

// ---- Multitalker (bundle) internals ----------------------------------- //

// Per-speaker supervision for one multitalker decode pass, at the diar
// frame rate (80 ms). Mirrors NeMo's SpeakerTaggedASR hand-off:
//   masked mode (use_kernel=false, the NeMo example default): the mel is
//     masked per feature frame with the 8x-expanded binarized target
//     activity (mask_features: masked = mel * mask; original-zero
//     features -> mask_value).
//   kernel mode (use_kernel=true): spk (raw sigmoids) and bg (binarized
//     union of other actives) drive the layer-0 speaker-kernel injection
//     as graph inputs (set_speaker_targets path).
struct MultitalkerPass {
    bool                 use_kernel = false;
    float                threshold  = 0.5f;
    float                mask_value = -16.6355f;
    std::vector<float>   spk;  // [T_diar] target-speaker activity
    std::vector<float>   bg;   // [T_diar] binarized union of other actives
    // Kernel-mode chunk gating (reference cache_gating): encoder-frame
    // indices of this speaker's ACTIVE 14-frame chunks, in order. The
    // conformer + decoder see only these frames (gathered after the
    // continuous pre_encode), mirroring the reference per-speaker cache
    // that only advances on active chunks. Empty = no gating (masked
    // mode / single-speaker).
    std::vector<int32_t> keep_enc_frames;
};

// Shared one-utterance inference helper (mel -> encoder -> host decode ->
// result build on pc). Defined in model.cpp; also driven per speaker by
// the multitalker orchestrator with a supervision pass.
transcribe_status run_one_shot_inner(ParakeetSession *             pc,
                                     ParakeetModel *               pm,
                                     const float *                 pcm,
                                     int                           n_samples,
                                     const transcribe_run_params * params,
                                     const MultitalkerPass *       mt = nullptr);

// Build tokens/words/segments/full_text on pc from pc->raw_tokens
// (defined in model.cpp; shared by the offline decode and the
// multitalker streaming pass).
transcribe_status build_result_from_raw_tokens(ParakeetSession *             pc,
                                               ParakeetModel *               pm,
                                               const transcribe_run_params * params);

// Streaming primitives (defined in model.cpp) shared with the
// multitalker streaming pass, which drives one cache/decoder-state
// instance per speaker by swapping it into pc around each call.
transcribe_status init_streaming_caches(ParakeetSession * pc, ParakeetModel * pm);
void              reset_streaming_decoder_state(ParakeetSession * pc, const ParakeetModel * pm);
transcribe_status emit_streaming_chunk(ParakeetSession * pc,
                                       ParakeetModel *   pm,
                                       const float *     mel_chunk_data,
                                       int               n_mel_chunk_frames,
                                       int               drop_extra_pre_encoded,
                                       int               mel_frames_advance,
                                       const float *     mt_spk      = nullptr,
                                       const float *     mt_bg       = nullptr,
                                       int               mt_mask_len = 0);

// Multitalker offline run (diarize=ON on a bundle model): embedded
// sortformer streaming forward -> active-speaker selection -> one decode
// pass per active speaker -> merged speaker-tagged result. Defined in
// multitalker.cpp.
transcribe_status run_multitalker(ParakeetSession *             pc,
                                  ParakeetModel *               pm,
                                  const float *                 pcm,
                                  int                           n_samples,
                                  const transcribe_run_params * params);

}  // namespace transcribe::parakeet