transcribe-cpp-sys 0.0.8

Native FFI bindings for transcribe.cpp, a C/C++ speech-to-text library built on ggml
Documentation
// arch/gigaam/weights.cpp - read_gigaam_hparams + build_gigaam_weights.
//
// Pattern follows arch/parakeet/weights.cpp: read every required KV
// from the GGUF, then validate the tensor catalog by name + shape using
// transcribe::weights::find_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 <cstring>

namespace transcribe::gigaam {

namespace {

constexpr const char * kFamilyTag = "gigaam";

} // namespace

transcribe_status read_gigaam_hparams(const gguf_context * gguf,
                                      GigaamHParams &      hp)
{
    if (gguf == nullptr) return TRANSCRIBE_ERR_INVALID_ARG;

    // Head discriminator (drives RNN-T vs CTC reads below).
    {
        std::string head_kind_str;
        if (auto st = read_required_string_kv(gguf, "stt.gigaam.head_kind",
                                              kFamilyTag, head_kind_str);
            st != TRANSCRIBE_OK) return st;
        if (head_kind_str == "rnnt") {
            hp.head_kind = HeadKind::RNNT;
        } else if (head_kind_str == "ctc") {
            hp.head_kind = HeadKind::CTC;
        } else {
            log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                         "gigaam: unsupported stt.gigaam.head_kind=%s",
                         head_kind_str.c_str());
            return TRANSCRIBE_ERR_GGUF;
        }
    }

    // Encoder.
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.n_layers",           kFamilyTag, hp.enc_n_layers);           st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.d_model",            kFamilyTag, hp.enc_d_model);            st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.n_heads",            kFamilyTag, hp.enc_n_heads);            st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.d_ff",               kFamilyTag, hp.enc_d_ff);               st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.conv_kernel",        kFamilyTag, hp.enc_conv_kernel);        st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.subsampling_factor", kFamilyTag, hp.enc_subsampling_factor); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.subs_kernel_size",   kFamilyTag, hp.enc_subs_kernel_size);   st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.pos_emb_max_len",    kFamilyTag, hp.enc_pos_emb_max_len);    st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.gigaam.encoder.feat_in",            kFamilyTag, hp.enc_feat_in);            st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.gigaam.encoder.self_attention_model", kFamilyTag, hp.enc_self_attention_model); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.gigaam.encoder.conv_norm_type",  kFamilyTag, hp.enc_conv_norm_type);  st != TRANSCRIBE_OK) return st;

    if (hp.enc_n_heads <= 0 || hp.enc_d_model % hp.enc_n_heads != 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "gigaam: invariant d_model %% n_heads != 0 "
                     "(d_model=%d, n_heads=%d)",
                     hp.enc_d_model, hp.enc_n_heads);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_self_attention_model != "rotary") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "gigaam: unsupported self_attention_model=%s (expected rotary)",
                     hp.enc_self_attention_model.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_conv_norm_type != "layer_norm") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "gigaam: unsupported conv_norm_type=%s (expected layer_norm)",
                     hp.enc_conv_norm_type.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }

    if (hp.head_kind == HeadKind::RNNT) {
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.predictor.hidden",   kFamilyTag, hp.pred_hidden);   st != TRANSCRIBE_OK) return st;
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.predictor.n_layers", kFamilyTag, hp.pred_n_layers); st != TRANSCRIBE_OK) return st;
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.predictor.vocab",    kFamilyTag, hp.pred_vocab);    st != TRANSCRIBE_OK) return st;
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.joint.hidden",       kFamilyTag, hp.joint_hidden);  st != TRANSCRIBE_OK) return st;
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.joint.num_classes",  kFamilyTag, hp.joint_n_classes); st != TRANSCRIBE_OK) return st;
        if (auto st = read_required_string_kv(gguf, "stt.gigaam.joint.activation", kFamilyTag, hp.joint_activation); st != TRANSCRIBE_OK) return st;
        if (hp.pred_vocab != hp.joint_n_classes) {
            log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                         "gigaam: pred_vocab (%d) != joint_n_classes (%d)",
                         hp.pred_vocab, hp.joint_n_classes);
            return TRANSCRIBE_ERR_GGUF;
        }
    } else { // CTC
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.head.feat_in",     kFamilyTag, hp.head_feat_in);   st != TRANSCRIBE_OK) return st;
        if (auto st = read_required_u32_kv(gguf, "stt.gigaam.head.num_classes", kFamilyTag, hp.head_n_classes); st != TRANSCRIBE_OK) return st;
        if (hp.head_feat_in != hp.enc_d_model) {
            log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                         "gigaam: ctc head.feat_in (%d) != encoder.d_model (%d)",
                         hp.head_feat_in, hp.enc_d_model);
            return TRANSCRIBE_ERR_GGUF;
        }
    }

    // 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_f32_kv   (gguf, "stt.frontend.dither",      kFamilyTag, hp.fe_dither);      st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv   (gguf, "stt.frontend.pre_emphasis",kFamilyTag, hp.fe_pre_emphasis);st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv   (gguf, "stt.frontend.f_min",       kFamilyTag, hp.fe_f_min);       st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv   (gguf, "stt.frontend.f_max",       kFamilyTag, hp.fe_f_max);       st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_bool_kv  (gguf, "stt.frontend.center",      kFamilyTag, false, hp.fe_center); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.frontend.mel_norm",    kFamilyTag, hp.fe_mel_norm);    st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv   (gguf, "stt.frontend.log_clamp_min", kFamilyTag, hp.fe_log_clamp_min); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv   (gguf, "stt.frontend.log_clamp_max", kFamilyTag, hp.fe_log_clamp_max); st != TRANSCRIBE_OK) return st;

    return TRANSCRIBE_OK;
}

// ---------------------------------------------------------------------------
// Weights catalog
// ---------------------------------------------------------------------------

#define GET_LN(slot, name) do { \
    ggml_tensor * _t = transcribe::weights::find_tensor( \
        ctx_meta, (name), {GGML_TYPE_F32}, \
        {hp.enc_d_model}, kFamilyTag); \
    if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
    (slot) = _t; \
} while (0)

#define GET_LIN_W(slot, name, in_dim, out_dim) do { \
    ggml_tensor * _t = transcribe::weights::find_tensor( \
        ctx_meta, (name), {TRANSCRIBE_QUANT_LINEAR_TYPES}, \
        {(in_dim), (out_dim)}, kFamilyTag); \
    if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
    (slot) = _t; \
} while (0)

#define GET_LIN_B(slot, name, out_dim) do { \
    ggml_tensor * _t = transcribe::weights::find_tensor( \
        ctx_meta, (name), {GGML_TYPE_F32}, \
        {(out_dim)}, kFamilyTag); \
    if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
    (slot) = _t; \
} while (0)

#define GET_CONV_W(slot, name, k, in_dim, out_dim) do { \
    ggml_tensor * _t = transcribe::weights::find_tensor( \
        ctx_meta, (name), {TRANSCRIBE_QUANT_CONV_TYPES}, \
        {(k), (in_dim), (out_dim)}, kFamilyTag); \
    if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
    (slot) = _t; \
} while (0)

#define GET_F32_1D(slot, name, n) do { \
    ggml_tensor * _t = transcribe::weights::find_tensor( \
        ctx_meta, (name), {GGML_TYPE_F32}, \
        {(n)}, kFamilyTag); \
    if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
    (slot) = _t; \
} while (0)

#define GET_F32_2D(slot, name, n0, n1) do { \
    ggml_tensor * _t = transcribe::weights::find_tensor( \
        ctx_meta, (name), {GGML_TYPE_F32}, \
        {(n0), (n1)}, kFamilyTag); \
    if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
    (slot) = _t; \
} while (0)

using transcribe::weights::lname;

namespace {

transcribe_status build_block(ggml_context *        ctx_meta,
                              const GigaamHParams & hp,
                              int                   i,
                              GigaamBlock &         b)
{
    const int d = hp.enc_d_model;
    const int dff = hp.enc_d_ff;
    const int dw_k = hp.enc_conv_kernel;

    // Macaron FF1.
    GET_LN  (b.norm_ff1_w, lname("enc.blocks.%d.norm_ff1.weight", i));
    GET_LN  (b.norm_ff1_b, lname("enc.blocks.%d.norm_ff1.bias",   i));
    GET_LIN_W(b.ff1_lin1_w, lname("enc.blocks.%d.ff1.linear1.weight", i), d,   dff);
    GET_LIN_B(b.ff1_lin1_b, lname("enc.blocks.%d.ff1.linear1.bias",   i), dff);
    GET_LIN_W(b.ff1_lin2_w, lname("enc.blocks.%d.ff1.linear2.weight", i), dff, d);
    GET_LIN_B(b.ff1_lin2_b, lname("enc.blocks.%d.ff1.linear2.bias",   i), d);

    // Conv module.
    GET_LN  (b.norm_conv_w, lname("enc.blocks.%d.norm_conv.weight", i));
    GET_LN  (b.norm_conv_b, lname("enc.blocks.%d.norm_conv.bias",   i));
    GET_CONV_W(b.conv_pw1_w, lname("enc.blocks.%d.conv.pointwise1.weight", i), 1, d, 2 * d);
    GET_LIN_B (b.conv_pw1_b, lname("enc.blocks.%d.conv.pointwise1.bias",   i), 2 * d);
    GET_CONV_W(b.conv_dw_w,  lname("enc.blocks.%d.conv.depthwise.weight", i), dw_k, 1, d);
    GET_LIN_B (b.conv_dw_b,  lname("enc.blocks.%d.conv.depthwise.bias",   i), d);
    GET_LN  (b.conv_ln_w,    lname("enc.blocks.%d.conv.ln.weight", i));
    GET_LN  (b.conv_ln_b,    lname("enc.blocks.%d.conv.ln.bias",   i));
    GET_CONV_W(b.conv_pw2_w, lname("enc.blocks.%d.conv.pointwise2.weight", i), 1, d, d);
    GET_LIN_B (b.conv_pw2_b, lname("enc.blocks.%d.conv.pointwise2.bias",   i), d);

    // Self-attention (rotary, no linear_pos / pos_bias_*).
    GET_LN  (b.norm_attn_w, lname("enc.blocks.%d.norm_attn.weight", i));
    GET_LN  (b.norm_attn_b, lname("enc.blocks.%d.norm_attn.bias",   i));
    GET_LIN_W(b.attn_q_w,   lname("enc.blocks.%d.attn.linear_q.weight",   i), d, d);
    GET_LIN_B(b.attn_q_b,   lname("enc.blocks.%d.attn.linear_q.bias",     i), d);
    GET_LIN_W(b.attn_k_w,   lname("enc.blocks.%d.attn.linear_k.weight",   i), d, d);
    GET_LIN_B(b.attn_k_b,   lname("enc.blocks.%d.attn.linear_k.bias",     i), d);
    GET_LIN_W(b.attn_v_w,   lname("enc.blocks.%d.attn.linear_v.weight",   i), d, d);
    GET_LIN_B(b.attn_v_b,   lname("enc.blocks.%d.attn.linear_v.bias",     i), d);
    GET_LIN_W(b.attn_out_w, lname("enc.blocks.%d.attn.linear_out.weight", i), d, d);
    GET_LIN_B(b.attn_out_b, lname("enc.blocks.%d.attn.linear_out.bias",   i), d);

    // Macaron FF2.
    GET_LN  (b.norm_ff2_w, lname("enc.blocks.%d.norm_ff2.weight", i));
    GET_LN  (b.norm_ff2_b, lname("enc.blocks.%d.norm_ff2.bias",   i));
    GET_LIN_W(b.ff2_lin1_w, lname("enc.blocks.%d.ff2.linear1.weight", i), d,   dff);
    GET_LIN_B(b.ff2_lin1_b, lname("enc.blocks.%d.ff2.linear1.bias",   i), dff);
    GET_LIN_W(b.ff2_lin2_w, lname("enc.blocks.%d.ff2.linear2.weight", i), dff, d);
    GET_LIN_B(b.ff2_lin2_b, lname("enc.blocks.%d.ff2.linear2.bias",   i), d);

    GET_LN  (b.norm_out_w, lname("enc.blocks.%d.norm_out.weight", i));
    GET_LN  (b.norm_out_b, lname("enc.blocks.%d.norm_out.bias",   i));

    return TRANSCRIBE_OK;
}

} // namespace

transcribe_status build_gigaam_weights(ggml_context *        ctx_meta,
                                       const GigaamHParams & hp,
                                       GigaamWeights &       weights)
{
    if (ctx_meta == nullptr) return TRANSCRIBE_ERR_INVALID_ARG;

    const int d   = hp.enc_d_model;
    const int feat_in = hp.enc_feat_in;
    const int sub_k = hp.enc_subs_kernel_size;
    const int n_freq = hp.fe_n_fft / 2 + 1;

    // Frontend buffers (baked into the GGUF by the converter).
    GET_F32_2D(weights.frontend_mel_filterbank,
               "frontend.mel_filterbank", n_freq, hp.fe_num_mels);
    GET_F32_1D(weights.frontend_window, "frontend.window", hp.fe_win_length);

    // Pre-encode (2 stride-2 conv1d).
    GET_CONV_W(weights.pre_encode.conv0_w, "enc.pre_encode.conv.0.weight",
               sub_k, feat_in, d);
    GET_LIN_B (weights.pre_encode.conv0_b, "enc.pre_encode.conv.0.bias", d);
    GET_CONV_W(weights.pre_encode.conv2_w, "enc.pre_encode.conv.2.weight",
               sub_k, d, d);
    GET_LIN_B (weights.pre_encode.conv2_b, "enc.pre_encode.conv.2.bias", d);

    // Conformer blocks.
    weights.blocks.resize(hp.enc_n_layers);
    for (int i = 0; i < hp.enc_n_layers; ++i) {
        if (auto st = build_block(ctx_meta, hp, i, weights.blocks[i]);
            st != TRANSCRIBE_OK) return st;
    }

    // Head.
    if (hp.head_kind == HeadKind::RNNT) {
        // pred.embed.weight is a token-id lookup table consumed by the
        // host RNN-T decoder. The quantizer treats it as a linear-weight
        // tile (TRANSCRIBE_QUANT_LINEAR_TYPES: F32/F16/Q*), so the loader
        // must too. decoder.cpp's readback_f32 dequantizes once at load
        // time via ggml type traits.
        GET_LIN_W(weights.predictor.embed_w, "pred.embed.weight",
                  hp.pred_hidden, hp.pred_vocab);

        weights.predictor.lstm.resize(hp.pred_n_layers);
        for (int i = 0; i < hp.pred_n_layers; ++i) {
            auto & L = weights.predictor.lstm[i];
            // PyTorch nn.LSTM packs 4 gates along the row dim.
            GET_LIN_W(L.Wx, lname("pred.lstm.%d.Wx", i),
                      hp.pred_hidden, 4 * hp.pred_hidden);
            GET_LIN_W(L.Wh, lname("pred.lstm.%d.Wh", i),
                      hp.pred_hidden, 4 * hp.pred_hidden);
            GET_LIN_B(L.b,  lname("pred.lstm.%d.bias", i),
                      4 * hp.pred_hidden);
        }
        GET_LIN_W(weights.joint.enc_w,  "joint.enc.weight",  d,   hp.joint_hidden);
        GET_LIN_B(weights.joint.enc_b,  "joint.enc.bias",         hp.joint_hidden);
        GET_LIN_W(weights.joint.pred_w, "joint.pred.weight", hp.pred_hidden, hp.joint_hidden);
        GET_LIN_B(weights.joint.pred_b, "joint.pred.bias",         hp.joint_hidden);
        GET_LIN_W(weights.joint.out_w,  "joint.out.weight",  hp.joint_hidden, hp.joint_n_classes);
        GET_LIN_B(weights.joint.out_b,  "joint.out.bias",         hp.joint_n_classes);
    } else { // CTC
        GET_CONV_W(weights.ctc_head.weight, "head.ctc.weight",
                   1, d, hp.head_n_classes);
        GET_LIN_B (weights.ctc_head.bias,   "head.ctc.bias", hp.head_n_classes);
    }

    return TRANSCRIBE_OK;
}

} // namespace transcribe::gigaam