transcribe-cpp-sys 0.0.4

// arch/cohere/encoder.cpp - Cohere ASR Conformer encoder graph builder.
//
// Thin glue over the shared Conformer helpers in src/conformer/. The
// per-op helpers (layer_norm, f32-friendly conv helpers, rel_shift,
// conv_module, rel_pos_mhsa, build_conformer_block, build_pre_encode)
// now live in transcribe::conformer. This file owns only Cohere-
// specific glue:
//
//   - detect_direct_dw (different policy from Parakeet: Cohere uses
//     the im2col path on Metal and CPU, direct only on Vulkan/CUDA)
//   - to_view() helpers that project CoherePreEncode / CohereBlock
//     onto the shared nullable-pointer views (FFN, attention, and
//     conv biases are all present in Cohere)
//   - build_encoder_graph orchestration, including the post-encoder
//     enc_dec_proj projection and debug dump-preservation calls

#include "encoder.h"

#include "weights.h"

#include "conformer/conformer.h"
#include "transcribe-debug.h"
#include "transcribe-log.h"

#include "ggml.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>

namespace transcribe::cohere {

namespace {

namespace conf = transcribe::conformer;

// ----- Per-family depthwise dispatch policy -----------------------
//
// Cohere uses the direct ggml_conv_2d_dw_direct path on Vulkan and
// CUDA (both have native kernels), and the im2col + mul_mat path on
// Metal and CPU. This is the opposite of Parakeet on Metal/CPU — the
// divergence is preserved as-is during the extraction to avoid any
// quiet behavior changes. See the ConvPolicy comment in conformer.h
// for why the two sites (in-block dw and pre_encode dw) share the
// same detect result in Cohere.
bool detect_direct_dw(const char * backend) {
    const char * env = std::getenv("TRANSCRIBE_CONV_DIRECT_DW");
    if (env != nullptr) return true;
    env = std::getenv("TRANSCRIBE_CONV_NO_DIRECT_DW");
    if (env != nullptr) return false;
    if (backend == nullptr) return false;
    if (std::strstr(backend, "Vulkan") != nullptr) return true;
    if (std::strstr(backend, "CUDA")   != nullptr) return true;
    return false;
}

// ----- Views ------------------------------------------------------

conf::PreEncodeView to_view(const CoherePreEncode & pe) {
    conf::PreEncodeView v;
    v.conv0_w = pe.conv0_w; v.conv0_b = pe.conv0_b;
    v.conv2_w = pe.conv2_w; v.conv2_b = pe.conv2_b;
    v.conv3_w = pe.conv3_w; v.conv3_b = pe.conv3_b;
    v.conv5_w = pe.conv5_w; v.conv5_b = pe.conv5_b;
    v.conv6_w = pe.conv6_w; v.conv6_b = pe.conv6_b;
    v.out_w   = pe.out_w;   v.out_b   = pe.out_b;
    return v;
}

conf::BlockView to_view(const CohereBlock & b) {
    conf::BlockView v;

    // FF1 (biases present).
    v.norm_ff1_w = b.norm_ff1_w;
    v.norm_ff1_b = b.norm_ff1_b;
    v.ff1_lin1_w = b.ff1_lin1_w; v.ff1_lin1_b = b.ff1_lin1_b;
    v.ff1_lin2_w = b.ff1_lin2_w; v.ff1_lin2_b = b.ff1_lin2_b;

    // Self-attention (biases on Q/K/V/out; pos projection is bias-free).
    v.norm_attn_w = b.norm_attn_w;
    v.norm_attn_b = b.norm_attn_b;
    v.attn_q_w    = b.attn_q_w;   v.attn_q_b   = b.attn_q_b;
    v.attn_k_w    = b.attn_k_w;   v.attn_k_b   = b.attn_k_b;
    v.attn_v_w    = b.attn_v_w;   v.attn_v_b   = b.attn_v_b;
    v.attn_out_w  = b.attn_out_w; v.attn_out_b = b.attn_out_b;
    v.attn_pos_w  = b.attn_pos_w;
    v.attn_pos_u  = b.attn_pos_u;
    v.attn_pos_v  = b.attn_pos_v;

    // Conv module (biases on pw1/dw/pw2).
    v.norm_conv_w         = b.norm_conv_w;
    v.norm_conv_b         = b.norm_conv_b;
    v.conv_pw1_w          = b.conv_pw1_w; v.conv_pw1_b = b.conv_pw1_b;
    v.conv_dw_w           = b.conv_dw_w;  v.conv_dw_b  = b.conv_dw_b;
    v.conv_pw2_w          = b.conv_pw2_w; v.conv_pw2_b = b.conv_pw2_b;
    v.conv_bn_fused_scale = b.conv_bn_fused_scale;
    v.conv_bn_fused_bias  = b.conv_bn_fused_bias;

    // FF2 (biases present).
    v.norm_ff2_w = b.norm_ff2_w;
    v.norm_ff2_b = b.norm_ff2_b;
    v.ff2_lin1_w = b.ff2_lin1_w; v.ff2_lin1_b = b.ff2_lin1_b;
    v.ff2_lin2_w = b.ff2_lin2_w; v.ff2_lin2_b = b.ff2_lin2_b;

    v.norm_out_w = b.norm_out_w;
    v.norm_out_b = b.norm_out_b;
    return v;
}

} // namespace

EncoderBuild build_encoder_graph(ggml_context *         ctx,
                                 const CohereWeights &  w,
                                 const CohereHParams &  hp,
                                 int                    n_mel_frames,
                                 ggml_type              kv_type,
                                 bool                   use_flash,
                                 const char *           backend_name)
{
    conf::ConvPolicy policy {};
    policy.direct_pw               = conf::detect_direct_pw(backend_name);
    const bool direct_dw           = detect_direct_dw(backend_name);
    policy.direct_dw_in_block      = direct_dw;
    policy.direct_dw_in_pre_encode = direct_dw;

    EncoderBuild eb {};

    if (ctx == nullptr || n_mel_frames <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "cohere encoder: invalid arg "
                     "(ctx=%p, n_mel_frames=%d)",
                     static_cast<void *>(ctx), n_mel_frames);
        return eb;
    }

    if (hp.enc_subsampling_factor != 8 || hp.fe_num_mels != 128) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "cohere encoder: unsupported geometry "
                     "subsampling_factor=%d num_mels=%d "
                     "(only 8/128 implemented)",
                     hp.enc_subsampling_factor, hp.fe_num_mels);
        return eb;
    }

    // Mel input handle.
    eb.mel_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32,
                                   n_mel_frames, hp.fe_num_mels);
    if (eb.mel_in == nullptr) return eb;
    ggml_set_name(eb.mel_in, "mel.in");
    ggml_set_input(eb.mel_in);
    eb.dumps.mel_in = eb.mel_in;

    // Pre-encode.
    ggml_tensor * x = conf::build_pre_encode(ctx, to_view(w.pre_encode),
                                             eb.mel_in, policy,
                                             /*name_prefix=*/"enc.pre_encode",
                                             /*error_tag=*/"cohere");
    if (x == nullptr) return eb;
    eb.dumps.pre_encode_out = x;
    transcribe::debug::mark_tensor_for_dump(x);

    // Conformer blocks.
    if (!w.blocks.empty()) {
        const int64_t T_enc = x->ne[1];

        // Positional embedding input.
        const int64_t pos_len = 2 * T_enc - 1;
        eb.pos_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32,
                                           hp.enc_d_model, pos_len);
        ggml_set_name(eb.pos_emb_in, "pos_emb.in");
        ggml_set_input(eb.pos_emb_in);
        eb.dumps.pos_emb = eb.pos_emb_in;

        conf::BlockParams bparams {};
        bparams.d_model     = hp.enc_d_model;
        bparams.n_head      = hp.enc_n_heads;
        bparams.conv_kernel = hp.enc_conv_kernel;
        bparams.kv_type     = kv_type;
        bparams.use_flash   = use_flash;
        bparams.policy      = policy;

        // Block 0 (named for the dump stream).
        {
            x = conf::build_conformer_block(ctx, x, eb.pos_emb_in,
                                            to_view(w.blocks[0]), bparams);
            x = conf::named(x, "enc.block.0.out");
            eb.dumps.block0_out = x;
            transcribe::debug::mark_tensor_for_dump(x);
        }

        // Blocks 1..N-1.
        for (size_t i = 1; i < w.blocks.size(); ++i) {
            x = conf::build_conformer_block(ctx, x, eb.pos_emb_in,
                                            to_view(w.blocks[i]), bparams);
            // Spot-check at the middle block (n_layers/2 - 1 to match
            // the reference dump naming).
            if (i == static_cast<size_t>(hp.enc_n_layers / 2 - 1)) {
                char bname[64];
                std::snprintf(bname, sizeof(bname), "enc.block.%zu.out", i);
                x = conf::named(x, bname);
                eb.dumps.block_mid_out = x;
                transcribe::debug::mark_tensor_for_dump(x);
            }
            // Spot-check last block.
            if (i == w.blocks.size() - 1) {
                char bname[64];
                std::snprintf(bname, sizeof(bname), "enc.block.%zu.out", i);
                x = conf::named(x, bname);
                eb.dumps.block_last_out = x;
                transcribe::debug::mark_tensor_for_dump(x);
            }
        }
    }

    eb.dumps.final_out = x;
    x = conf::named(x, "enc.final");
    transcribe::debug::mark_tensor_for_dump(x);

    // Encoder-decoder projection. Cohere-specific — not part of the
    // shared Conformer block, so it stays here.
    {
        ggml_tensor * proj = ggml_mul_mat(ctx, w.enc_dec_proj.weight, x);
        proj = ggml_add(ctx, proj, w.enc_dec_proj.bias);
        proj = conf::named(proj, "enc_dec_proj.out");
        eb.dumps.enc_dec_proj_out = proj;
        transcribe::debug::mark_tensor_for_dump(proj);
        x = proj;
    }

    eb.out = x;
    ggml_set_output(eb.out);

    // Build the forward cgraph. 48 blocks + enc_dec_proj means a
    // large graph; 16384 leaves headroom.
    eb.graph = ggml_new_graph_custom(ctx, 16384, false);
    if (eb.graph == nullptr) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "cohere encoder: ggml_new_graph_custom failed");
        return eb;
    }
    ggml_build_forward_expand(eb.graph, eb.out);

    return eb;
}

} // namespace transcribe::cohere