transcribe-cpp-sys 0.0.7

// arch/moonshine/weights.h - canonical Moonshine ASR tensor catalog and
// per-instance weight slots.
//
// Architectural highlights (see `reports/porting/moonshine/forward-map.md`
// for the full forward pass):
//
//   - Encoder: 3-conv stem on raw 16 kHz PCM (Conv1: k=127 s=64 no-bias
//     → tanh → GroupNorm(num_groups=1); Conv2: k=7 s=3 → GELU; Conv3:
//     k=3 s=2 → GELU), then n_layers pre-LN transformer blocks (MHSA
//     with partial RoPE 0.9 on Q/K + GELU MLP), then a final LN.
//
//   - Decoder: token embedding (tied to lm_head), n_layers pre-LN blocks
//     (self-attn with partial RoPE 0.9 on Q/K, cross-attn with NO RoPE,
//     SwiGLU MLP fc1 hidden→2·intermediate split-gate-silu), final LN,
//     tied logits head.
//
//   - Attention biases: ALL projections (q/k/v/o on both self- and
//     cross-attn) are bias-less (`config.attention_bias=false`). Only
//     MLP fc1/fc2 carry bias.
//
//   - LayerNorm biases: every `nn.LayerNorm` in modeling_moonshine.py
//     is constructed with `bias=False`. The catalog has no `_b` slot
//     for any norm.
//
//   - Logits head: tied to `dec.token_embd.weight` (no separate head
//     tensor in the safetensors).
//
//   - Frontend buffers: NONE — moonshine has no mel/STFT. The conv
//     stem operates directly on raw PCM.
//
// Tensor naming follows `scripts/convert-moonshine.py`:
//
//   Encoder top:   enc.conv.{0,1,2}.{weight,bias}, enc.conv.norm.{w,b},
//                  enc.final_norm.weight  (no bias)
//   Encoder block: enc.blocks.{i}.norm_attn.weight  (no bias)
//                  enc.blocks.{i}.attn.{q,k,v,out}.weight  (no biases)
//                  enc.blocks.{i}.norm_ffn.weight   (no bias)
//                  enc.blocks.{i}.ffn.fc{1,2}.{weight,bias}
//   Decoder top:   dec.token_embd.weight  (tied to lm_head),
//                  dec.final_norm.weight  (no bias)
//   Decoder block: dec.blocks.{i}.norm_self.weight   (no bias)
//                  dec.blocks.{i}.self_attn.{q,k,v,out}.weight  (no biases)
//                  dec.blocks.{i}.norm_cross.weight  (no bias)
//                  dec.blocks.{i}.cross_attn.{q,k,v,out}.weight  (no biases)
//                  dec.blocks.{i}.norm_ffn.weight    (no bias)
//                  dec.blocks.{i}.ffn.fc1.{weight,bias}  (fc1 emits 2·inter)
//                  dec.blocks.{i}.ffn.fc2.{weight,bias}

#pragma once

#include "transcribe.h"

#include <cstdint>
#include <string>
#include <vector>

struct gguf_context;
struct ggml_context;
struct ggml_tensor;

namespace transcribe::moonshine {

// ---------------------------------------------------------------------------
// Hyperparameters
// ---------------------------------------------------------------------------

struct MoonshineHParams {
    // Encoder.
    int32_t enc_n_layers     = 0;
    int32_t enc_d_model      = 0;
    int32_t enc_n_heads      = 0;
    int32_t enc_n_kv_heads   = 0;   // moonshine has no GQA on tiny: == n_heads
    int32_t enc_ffn_dim      = 0;   // intermediate_size
    std::string enc_activation;     // "gelu"

    // Decoder.
    int32_t dec_n_layers     = 0;
    int32_t dec_d_model      = 0;   // == enc_d_model on every shipped variant
    int32_t dec_n_heads      = 0;
    int32_t dec_n_kv_heads   = 0;
    int32_t dec_ffn_dim      = 0;   // intermediate_size (fc1 emits 2·this)
    int32_t dec_max_position_embeddings = 0;  // 194 for tiny; both target
                                              // length cap and rope max-seq
    int32_t dec_vocab_size   = 0;
    std::string dec_activation;     // "silu"
    bool    dec_tie_word_embeddings = true;

    // Special tokens.
    int32_t bos_token_id           = -1;  // 1, decoder_start
    int32_t eos_token_id           = -1;  // 2
    int32_t pad_token_id           = -1;  // 2 (== eos)
    int32_t decoder_start_token_id = -1;  // 1

    // Attention / RoPE.
    float   partial_rotary_factor = 0.9f;
    float   rope_theta            = 10000.0f;
    bool    attention_bias        = false;
    int32_t pad_head_dim_multiple = 0;   // 8 for tiny; 0 means "do not pad"

    // Conv stem (3-layer raw-PCM frontend).
    std::vector<int32_t> conv_channels;     // [enc_d_model, 2·enc_d_model, enc_d_model]
    std::vector<int32_t> conv_kernel_sizes; // [127, 7, 3]
    std::vector<int32_t> conv_strides;      // [64, 3, 2]
    int32_t conv_groupnorm_num_groups = 1;
    float   conv_groupnorm_eps        = 1e-5f;

    // Frontend (raw waveform — no mel).
    std::string fe_type;            // "raw"
    int32_t     fe_sample_rate = 0; // 16000

    // Capability flags read from stt.capability.*.
    bool cap_lang_detect = false;
    bool cap_translate   = false;
    bool cap_timestamps  = false;

    // Derived helpers.
    int32_t enc_head_dim() const { return enc_n_heads > 0 ? enc_d_model / enc_n_heads : 0; }
    int32_t dec_head_dim() const { return dec_n_heads > 0 ? dec_d_model / dec_n_heads : 0; }

    // Padded head dim used inside attention. HF Moonshine pads each
    // head's q/k/v to a multiple of `pad_head_dim_multiple`, runs the
    // matmul / FA, then slices the padding off after `o_proj`. For
    // tiny: head_dim=36 → padded=40 (pad=4). For variants where
    // head_dim is already a multiple of `pad_head_dim_multiple` the
    // padding is zero and the path is a no-op.
    int32_t enc_head_dim_padded() const {
        const int32_t hd = enc_head_dim();
        const int32_t m  = pad_head_dim_multiple;
        if (m <= 0 || hd <= 0) return hd;
        return ((hd + m - 1) / m) * m;
    }
    int32_t dec_head_dim_padded() const {
        const int32_t hd = dec_head_dim();
        const int32_t m  = pad_head_dim_multiple;
        if (m <= 0 || hd <= 0) return hd;
        return ((hd + m - 1) / m) * m;
    }

    // Number of leading head-dim positions rotated by partial RoPE.
    // HF rounds `int(head_dim · factor)` down to an even number so
    // the NEOX rotate_half halves match. For tiny: 36 · 0.9 = 32.4 →
    // floor 32, already even. The C++ port uses the floor-and-mask-odd
    // formula (`(int(hd*f)) & ~1`) for forward compatibility with
    // base / future variants.
    int32_t enc_head_dim_rot() const {
        const int32_t hd = enc_head_dim();
        const int32_t r  = static_cast<int32_t>(static_cast<float>(hd) * partial_rotary_factor);
        return r & ~int32_t{1};
    }
    int32_t dec_head_dim_rot() const {
        const int32_t hd = dec_head_dim();
        const int32_t r  = static_cast<int32_t>(static_cast<float>(hd) * partial_rotary_factor);
        return r & ~int32_t{1};
    }
};

transcribe_status read_moonshine_hparams(const gguf_context * gguf,
                                         MoonshineHParams &   hp);

// ---------------------------------------------------------------------------
// Weight slots
// ---------------------------------------------------------------------------

// Encoder conv stem: 3 Conv1d layers + GroupNorm right after conv0+tanh.
struct MoonshineEncStem {
    ggml_tensor * conv0_w     = nullptr;  // [127, 1,         d_model],   no bias
    ggml_tensor * conv1_w     = nullptr;  // [7,   d_model,   2·d_model]
    ggml_tensor * conv1_b     = nullptr;  // [2·d_model]
    ggml_tensor * conv2_w     = nullptr;  // [3,   2·d_model, d_model]
    ggml_tensor * conv2_b     = nullptr;  // [d_model]
    ggml_tensor * gn_w        = nullptr;  // [d_model]   (num_groups=1, ε=1e-5)
    ggml_tensor * gn_b        = nullptr;  // [d_model]
};

struct MoonshineEncTop {
    ggml_tensor * final_norm_w = nullptr;  // [d_model]   (no bias)
};

// One encoder transformer block. attention_bias=false: every q/k/v/o is
// weight-only. norm_* have no bias.
struct MoonshineEncBlock {
    ggml_tensor * norm_attn_w = nullptr;
    ggml_tensor * attn_q_w    = nullptr;
    ggml_tensor * attn_k_w    = nullptr;
    ggml_tensor * attn_v_w    = nullptr;
    ggml_tensor * attn_out_w  = nullptr;

    ggml_tensor * norm_ffn_w  = nullptr;
    ggml_tensor * ffn_fc1_w   = nullptr;  // [d_model, ffn_dim]
    ggml_tensor * ffn_fc1_b   = nullptr;  // [ffn_dim]
    ggml_tensor * ffn_fc2_w   = nullptr;  // [ffn_dim, d_model]
    ggml_tensor * ffn_fc2_b   = nullptr;  // [d_model]
};

// Decoder top: tied token embed + final LN.
struct MoonshineDecTop {
    ggml_tensor * token_embd_w = nullptr;  // [d_model, vocab_size]
    ggml_tensor * final_norm_w = nullptr;  // [d_model]   (no bias)
};

// One decoder transformer block. attention_bias=false on both self- and
// cross-attn. Decoder MLP is SwiGLU: fc1 emits 2·ffn_dim, fc2 takes ffn_dim.
struct MoonshineDecBlock {
    // Self-attention (causal, partial RoPE on q/k).
    ggml_tensor * norm_self_w  = nullptr;
    ggml_tensor * self_q_w     = nullptr;
    ggml_tensor * self_k_w     = nullptr;
    ggml_tensor * self_v_w     = nullptr;
    ggml_tensor * self_out_w   = nullptr;

    // Cross-attention (queries decoder state against encoder; no RoPE).
    ggml_tensor * norm_cross_w = nullptr;
    ggml_tensor * cross_q_w    = nullptr;
    ggml_tensor * cross_k_w    = nullptr;
    ggml_tensor * cross_v_w    = nullptr;
    ggml_tensor * cross_out_w  = nullptr;

    // SwiGLU MLP.
    ggml_tensor * norm_ffn_w   = nullptr;
    ggml_tensor * ffn_fc1_w    = nullptr;  // [d_model, 2·ffn_dim]
    ggml_tensor * ffn_fc1_b    = nullptr;  // [2·ffn_dim]
    ggml_tensor * ffn_fc2_w    = nullptr;  // [ffn_dim, d_model]
    ggml_tensor * ffn_fc2_b    = nullptr;  // [d_model]
};

struct MoonshineWeights {
    MoonshineEncStem                 enc_stem;
    MoonshineEncTop                  enc_top;
    std::vector<MoonshineEncBlock>   enc_blocks;
    MoonshineDecTop                  dec_top;
    std::vector<MoonshineDecBlock>   dec_blocks;
};

transcribe_status build_moonshine_weights(ggml_context *           ctx_meta,
                                          const MoonshineHParams & hp,
                                          MoonshineWeights &       weights);

} // namespace transcribe::moonshine