transcribe-cpp-sys 0.0.9

// arch/sensevoice/weights.cpp - read_sensevoice_hparams +
// build_sensevoice_weights.
//
// Pattern mirrors src/arch/parakeet/weights.cpp:
//   - read every required hparam from KV explicitly,
//   - build the tensor catalog as get_tensor() calls with explicit
//     expected shapes,
//   - missing tensor or shape mismatch → TRANSCRIBE_ERR_GGUF with a
//     log naming the offending tensor.

#include "weights.h"

#include "transcribe-meta.h"
#include "transcribe-weights-util.h"
#include "transcribe-log.h"

#include "ggml.h"
#include "gguf.h"

#include <cstdio>
#include <initializer_list>

namespace transcribe::sensevoice {

namespace {

constexpr const char * kFamilyTag = "sensevoice";

} // namespace

transcribe_status read_sensevoice_hparams(const gguf_context *  gguf,
                                          SenseVoiceHParams &   hp)
{
    if (gguf == nullptr) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }

    // Encoder.
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.n_blocks",   kFamilyTag, hp.enc_n_blocks);    st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.tp_blocks",  kFamilyTag, hp.enc_tp_blocks);   st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.d_model",    kFamilyTag, hp.enc_d_model);     st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.d_input",    kFamilyTag, hp.enc_d_input);     st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.n_heads",    kFamilyTag, hp.enc_n_heads);     st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.d_ff",       kFamilyTag, hp.enc_d_ff);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.kernel_size",kFamilyTag, hp.enc_kernel);      st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.encoder.sanm_shift", kFamilyTag, hp.enc_sanm_shift);  st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.sensevoice.encoder.attention_type", kFamilyTag, hp.enc_attn_type); st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_bool_kv(gguf, "stt.sensevoice.encoder.normalize_before", kFamilyTag, true, hp.enc_normalize_before); st != TRANSCRIBE_OK) return st;

    // Pre-encoder prefix indices. Required so the runtime never has to
    // hard-code lid_dict / textnorm_dict layouts.
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_auto",      kFamilyTag, hp.prefix_lang_auto);      st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_zh",        kFamilyTag, hp.prefix_lang_zh);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_en",        kFamilyTag, hp.prefix_lang_en);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_yue",       kFamilyTag, hp.prefix_lang_yue);       st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_ja",        kFamilyTag, hp.prefix_lang_ja);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_ko",        kFamilyTag, hp.prefix_lang_ko);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.lang_nospeech",  kFamilyTag, hp.prefix_lang_nospeech);  st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.withitn",        kFamilyTag, hp.prefix_withitn);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.sensevoice.special.woitn",          kFamilyTag, hp.prefix_woitn);          st != TRANSCRIBE_OK) return st;
    // event_speech / emotion_neutral are documentation-only KV; the
    // hard-coded literal [1, 2] in the upstream forward is what the
    // C++ encoder uses, so we don't read these keys into hparams.

    // Frontend.
    if (auto st = read_required_string_kv(gguf, "stt.frontend.type",         kFamilyTag, hp.fe_type);         st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.num_mels",     kFamilyTag, hp.fe_num_mels);     st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.sample_rate",  kFamilyTag, hp.fe_sample_rate);  st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.n_fft",        kFamilyTag, hp.fe_n_fft);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.win_length",   kFamilyTag, hp.fe_win_length);   st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.hop_length",   kFamilyTag, hp.fe_hop_length);   st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.frontend.window",       kFamilyTag, hp.fe_window);       st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.frontend.normalize",    kFamilyTag, hp.fe_normalize);    st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.frontend.fbank_style",  kFamilyTag, hp.fe_fbank_style);  st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf,    "stt.frontend.dither",       kFamilyTag, hp.fe_dither);       st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_bool_kv(gguf,   "stt.frontend.upscale_samples", kFamilyTag, true, hp.fe_upscale_samples); st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_bool_kv(gguf,   "stt.frontend.snip_edges",   kFamilyTag, true, hp.fe_snip_edges); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.lfr_m",        kFamilyTag, hp.fe_lfr_m);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf,    "stt.frontend.lfr_n",        kFamilyTag, hp.fe_lfr_n);        st != TRANSCRIBE_OK) return st;

    // Cross-field invariants.
    if (hp.enc_n_blocks <= 1 || hp.enc_tp_blocks < 0 ||
        hp.enc_d_model <= 0  || hp.enc_d_input <= 0 ||
        hp.enc_n_heads <= 0  || hp.enc_d_ff <= 0    ||
        hp.enc_kernel <= 0   || hp.enc_kernel % 2 == 0)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: encoder hparams invalid (n_blocks=%d "
                     "tp_blocks=%d d_model=%d d_input=%d n_heads=%d d_ff=%d "
                     "kernel=%d — kernel must be positive and odd)",
                     hp.enc_n_blocks, hp.enc_tp_blocks, hp.enc_d_model,
                     hp.enc_d_input, hp.enc_n_heads, hp.enc_d_ff, hp.enc_kernel);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_d_model % hp.enc_n_heads != 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: encoder d_model (%d) not divisible by n_heads (%d)",
                     hp.enc_d_model, hp.enc_n_heads);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_num_mels <= 0 || hp.fe_sample_rate <= 0 ||
        hp.fe_n_fft <= 0    || hp.fe_win_length <= 0  ||
        hp.fe_hop_length <= 0 || hp.fe_lfr_m <= 0     ||
        hp.fe_lfr_n <= 0)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: frontend dimensions must be positive");
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_d_input != hp.fe_num_mels * hp.fe_lfr_m) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: encoder d_input (%d) must equal num_mels (%d) "
                     "× lfr_m (%d) = %d",
                     hp.enc_d_input, hp.fe_num_mels, hp.fe_lfr_m,
                     hp.fe_num_mels * hp.fe_lfr_m);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_attn_type != "sanm") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: unsupported attention type \"%s\" "
                     "(only \"sanm\" is implemented)",
                     hp.enc_attn_type.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_type != "kaldi_fbank_lfr") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: unsupported frontend type \"%s\" "
                     "(only \"kaldi_fbank_lfr\" is implemented)",
                     hp.fe_type.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_window != "hamming") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: unsupported frontend window \"%s\" "
                     "(only \"hamming\" is implemented)",
                     hp.fe_window.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_normalize != "per_feature") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: unsupported frontend normalize \"%s\" "
                     "(only \"per_feature\" is implemented)",
                     hp.fe_normalize.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_fbank_style != "kaldi_htk") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: unsupported fbank_style \"%s\" "
                     "(only \"kaldi_htk\" is implemented)",
                     hp.fe_fbank_style.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }

    return TRANSCRIBE_OK;
}

namespace {

using transcribe::weights::find_tensor;
using transcribe::weights::lname;

constexpr const char * kTag = kFamilyTag;

#define GET_F32(slot, name, ...) \
    do { \
        ggml_tensor * _t = find_tensor( \
            ctx_meta, (name), \
            {GGML_TYPE_F32}, {__VA_ARGS__}, kTag); \
        if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
        (slot) = _t; \
    } while (0)

#define GET_CONV(slot, name, ...) \
    do { \
        ggml_tensor * _t = find_tensor( \
            ctx_meta, (name), \
            {TRANSCRIBE_QUANT_CONV_TYPES}, \
            {__VA_ARGS__}, kTag); \
        if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
        (slot) = _t; \
    } while (0)

#define GET_LIN(slot, name, ...) \
    do { \
        ggml_tensor * _t = find_tensor( \
            ctx_meta, (name), \
            {TRANSCRIBE_QUANT_LINEAR_TYPES}, \
            {__VA_ARGS__}, kTag); \
        if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
        (slot) = _t; \
    } while (0)

// Bind one SAN-M block (13 tensors). `prefix` is e.g. "enc.encoders.7"
// or "enc.encoders0.0". `d_in_norm` and `d_in_qkv` are the input-side
// dims (= d_model except for encoders0[0] where they equal d_input).
transcribe_status load_block(ggml_context *      ctx_meta,
                             const char *        prefix,
                             int64_t             d_in,
                             int64_t             d_model,
                             int64_t             d_ff,
                             int64_t             kernel,
                             SenseVoiceBlock &   b)
{
    char buf[256];

    auto namef = [&](const char * suffix) -> const char * {
        std::snprintf(buf, sizeof(buf), "%s.%s", prefix, suffix);
        return buf;
    };
    // Reserve a second buffer for callers that need two names live in
    // the same expression. We dispatch through copies into local
    // std::strings to dodge buffer reuse.
    std::string name_norm_attn_w = std::string(prefix) + ".norm_attn.weight";
    std::string name_norm_attn_b = std::string(prefix) + ".norm_attn.bias";
    std::string name_qkv_w       = std::string(prefix) + ".attn.qkv.weight";
    std::string name_qkv_b       = std::string(prefix) + ".attn.qkv.bias";
    std::string name_out_w       = std::string(prefix) + ".attn.out.weight";
    std::string name_out_b       = std::string(prefix) + ".attn.out.bias";
    std::string name_fsmn_w      = std::string(prefix) + ".attn.fsmn.weight";
    std::string name_norm_ffn_w  = std::string(prefix) + ".norm_ffn.weight";
    std::string name_norm_ffn_b  = std::string(prefix) + ".norm_ffn.bias";
    std::string name_fc1_w       = std::string(prefix) + ".ffn.fc1.weight";
    std::string name_fc1_b       = std::string(prefix) + ".ffn.fc1.bias";
    std::string name_fc2_w       = std::string(prefix) + ".ffn.fc2.weight";
    std::string name_fc2_b       = std::string(prefix) + ".ffn.fc2.bias";

    GET_F32(b.norm_attn_w, name_norm_attn_w.c_str(), d_in);
    GET_F32(b.norm_attn_b, name_norm_attn_b.c_str(), d_in);
    // Fused QKV: PyTorch [out=3·d_model, in=d_in]; ggml ne[0]=in, ne[1]=out.
    GET_LIN(b.attn_qkv_w,  name_qkv_w.c_str(),       d_in, 3 * d_model);
    GET_F32(b.attn_qkv_b,  name_qkv_b.c_str(),       3 * d_model);
    // attn_out: PyTorch [d_model, d_model] always operates in d_model space.
    GET_LIN(b.attn_out_w,  name_out_w.c_str(),       d_model, d_model);
    GET_F32(b.attn_out_b,  name_out_b.c_str(),       d_model);
    // FSMN depthwise conv. PyTorch shape (out=d_model, in_per_group=1, kernel).
    // ggml stores conv1d weights as [kernel, in_per_group, out] in ne order
    // — see parakeet's depthwise convention.
    GET_CONV(b.attn_fsmn_w, name_fsmn_w.c_str(),     kernel, 1, d_model);

    GET_F32(b.norm_ffn_w,  name_norm_ffn_w.c_str(),  d_model);
    GET_F32(b.norm_ffn_b,  name_norm_ffn_b.c_str(),  d_model);
    GET_LIN(b.ffn_fc1_w,   name_fc1_w.c_str(),       d_model, d_ff);
    GET_F32(b.ffn_fc1_b,   name_fc1_b.c_str(),       d_ff);
    GET_LIN(b.ffn_fc2_w,   name_fc2_w.c_str(),       d_ff, d_model);
    GET_F32(b.ffn_fc2_b,   name_fc2_b.c_str(),       d_model);
    (void)namef; // currently unused; left for future per-block helpers.
    return TRANSCRIBE_OK;
}

} // namespace

transcribe_status build_sensevoice_weights(ggml_context *            ctx_meta,
                                           const SenseVoiceHParams & hp,
                                           SenseVoiceWeights &       weights)
{
    if (ctx_meta == nullptr) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }

    const int64_t d_input = hp.enc_d_input;
    const int64_t d_model = hp.enc_d_model;
    const int64_t d_ff    = hp.enc_d_ff;
    const int64_t kernel  = hp.enc_kernel;
    const int64_t vocab   = hp.vocab_size;

    if (vocab <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "sensevoice: build_sensevoice_weights called before "
                     "tokenizer load (vocab_size=%d)", hp.vocab_size);
        return TRANSCRIBE_ERR_GGUF;
    }

    // ----- frontend (per-feature CMVN) -----
    GET_F32(weights.cmvn_shift, "frontend.cmvn.shift", d_input);
    GET_F32(weights.cmvn_scale, "frontend.cmvn.scale", d_input);

    // ----- pre-encoder prefix embedding -----
    // 16 rows × d_input cols. Embed table rows are read with
    // ggml_get_rows; PyTorch shape [16, d_input] → ggml ne[0]=d_input,
    // ne[1]=16.
    GET_F32(weights.embed, "enc.embed.weight", d_input, 16);

    // ----- encoders0[0] (560 → 512 projection) -----
    if (auto st = load_block(ctx_meta, "enc.encoders0.0",
                             /*d_in=*/d_input, d_model, d_ff, kernel,
                             weights.encoders0);
        st != TRANSCRIBE_OK) return st;

    // ----- encoders[0..n_blocks-2] (49 blocks at d_model) -----
    weights.encoders.assign(hp.enc_n_blocks - 1, SenseVoiceBlock{});
    for (int i = 0; i < hp.enc_n_blocks - 1; ++i) {
        char prefix[64];
        std::snprintf(prefix, sizeof(prefix), "enc.encoders.%d", i);
        if (auto st = load_block(ctx_meta, prefix,
                                 d_model, d_model, d_ff, kernel,
                                 weights.encoders[static_cast<size_t>(i)]);
            st != TRANSCRIBE_OK) return st;
    }

    // ----- after_norm -----
    GET_F32(weights.after_norm_w, "enc.after_norm.weight", d_model);
    GET_F32(weights.after_norm_b, "enc.after_norm.bias",   d_model);

    // ----- tp_encoders[0..tp_blocks-1] -----
    weights.tp_encoders.assign(hp.enc_tp_blocks, SenseVoiceBlock{});
    for (int i = 0; i < hp.enc_tp_blocks; ++i) {
        char prefix[64];
        std::snprintf(prefix, sizeof(prefix), "enc.tp_encoders.%d", i);
        if (auto st = load_block(ctx_meta, prefix,
                                 d_model, d_model, d_ff, kernel,
                                 weights.tp_encoders[static_cast<size_t>(i)]);
            st != TRANSCRIBE_OK) return st;
    }

    // ----- tp_norm -----
    GET_F32(weights.tp_norm_w, "enc.tp_norm.weight", d_model);
    GET_F32(weights.tp_norm_b, "enc.tp_norm.bias",   d_model);

    // ----- CTC head -----
    // PyTorch [vocab, d_model] → ggml ne[0]=d_model, ne[1]=vocab.
    GET_LIN(weights.ctc_head_w, "ctc.head.weight", d_model, vocab);
    GET_F32(weights.ctc_head_b, "ctc.head.bias",   vocab);

    return TRANSCRIBE_OK;
}

#undef GET_F32
#undef GET_CONV
#undef GET_LIN

} // namespace transcribe::sensevoice