transcribe-cpp-sys 0.4.1

Native FFI bindings for transcribe.cpp, a C/C++ speech-to-text library built on ggml
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

// arch/granite_nar/encoder.h - NLE Conformer encoder + BPE CTC head.
//
// Reference: GraniteSpeechNarCTCEncoder in modeling_granite_speech_nar.py
// from the IBM Granite Speech NAR HF repo. Structurally identical to the
// AR granite encoder (block-local Shaw self-attention, conv_expansion=2
// GLU, mid-layer self-conditioned CTC bypass) plus two NAR-only additions:
//
//   1. A BPE CTC head (1024 → 100352) over a posterior-weighted
//      window=4 pool of valid frames. We expose the bypass-step char-CTC
//      mid_logits (1024 → 348) as `enc.ctc_logits` — the exact tensor
//      the reference model computes at self_conditioning_layer for the
//      self-conditioning residual. The pooled BPE head runs in-graph on
//      window-pooled hidden states and only per-window argmax ids leave the
//      device (see build_encoder_graph).
//   2. All-hidden-states capture: the projector consumes 4 encoder
//      hidden states (post-LN, pre-bypass at the chosen layer
//      boundaries; indices [4, 8, 12, -1] 1-indexed → layer outputs
//      after blocks 3, 7, 11, 15 in 0-indexed). The encoder graph
//      emits these concatenated along the channel axis so the
//      projector consumes one [4096, T_enc] tensor.
//
// This header is INTERNAL to src/arch/granite_nar/.

#pragma once

#include "ggml.h"
#include "weights.h"

#include <cstdint>
#include <vector>

struct ggml_context;
struct ggml_tensor;
struct ggml_cgraph;

namespace transcribe {
class MelFrontend;
}

namespace transcribe::granite_nar {

// Host-side mel + 2-frame stack. Reuses the AR granite implementation.
transcribe_status compute_mel_encoder_input(const transcribe::MelFrontend & mel,
                                            const float *                   pcm,
                                            int                             n_samples,
                                            int                             n_threads,
                                            std::vector<float> &            out_mel,
                                            int &                           out_t_enc);

struct EncoderBuild {
    // Graph inputs (caller uploads at compute time).
    ggml_tensor * mel_in          = nullptr;  // [input_dim, T_enc]
    ggml_tensor * pos_rows        = nullptr;  // [2*ctx-1] i32
    ggml_tensor * last_block_mask = nullptr;  // [ctx, ctx, n_blocks_local]
    ggml_tensor * zero_pad        = nullptr;  // optional [hidden, T_pad-T_enc]

    // Outputs.
    ggml_tensor *              cat_out         = nullptr;  // [num_enc_layers * hidden, T_enc]
                                                           //  the projector input
    ggml_tensor *              ctc_logits      = nullptr;  // [enc_out_dim=348, T_enc]
    ggml_tensor *              mid_blank_probs = nullptr;  // [T_enc] — softmax(mid_ctc)[blank].
                                                           //  importance = 1 - blank_prob_mid is
                                                           //  the BPE pool's per-frame weight.
    // Pooled BPE CTC head: argmax over bpe_output_dim per pool window, in
    // consecutive tiles of windows (tile k covers windows
    // [k * kBpeTileWindows, ...)), plus each window's total importance so the
    // host can treat all-blank windows exactly like the reference.
    std::vector<ggml_tensor *> bpe_window_ids;             // each [n_tile] i32
    ggml_tensor *              bpe_window_mass = nullptr;  // [1, n_windows] f32
    int                        n_bpe_windows   = 0;

    ggml_cgraph * graph = nullptr;

    struct Dumps {
        ggml_tensor * input_linear_out  = nullptr;
        ggml_tensor * block_0_out       = nullptr;
        ggml_tensor * block_mid_pre     = nullptr;  // block (N/2 - 1) post-LN
                                                    // (== `enc.block.7.out`)
        ggml_tensor * block_mid_post    = nullptr;  // block (N/2)   post-LN
                                                    // (== `enc.block.8.out`)
        ggml_tensor * block_last_out    = nullptr;  // block (N-1)
        ggml_tensor * ctc_logits        = nullptr;  // == ctc_logits, named
        // Block-0 sub-step taps.
        ggml_tensor * block_0_post_ff1  = nullptr;
        ggml_tensor * block_0_post_attn = nullptr;
        ggml_tensor * block_0_post_conv = nullptr;
        ggml_tensor * block_0_post_ff2  = nullptr;
    } dumps;

    int n_blocks_local = 0;
    int last_block_rem = 0;
};

EncoderBuild build_encoder_graph(ggml_context *            ctx,
                                 const GraniteNarWeights & weights,
                                 const GraniteNarHParams & hp,
                                 int                       T_enc,
                                 bool                      use_flash);

// Shaw bookkeeping helpers (identical to AR granite).
std::vector<int32_t> precompute_pos_rows(int context_size, int max_pos_emb);
std::vector<float>   precompute_last_block_mask(int context_size, int t_enc_remainder);

// Pool windows per BPE-head tile: bounds the [bpe_output_dim, tile] f32
// logits the head materialises (~103 MB at 100352 x 256).
constexpr int kBpeTileWindows = 256;

// Greedy collapse of the in-graph pooled BPE head: one argmax id per pool
// window and the window's total importance. Windows whose importance sum is
// <= 1e-9 emit blank, exactly as compute_bpe_ctc_initial_hypothesis does.
// Returns false (and leaves out_token_ids empty) if the two vectors disagree
// in length, which can only be a graph/caller bug: the tiles cover every
// window and both are read from tensors sized eb.n_bpe_windows.
bool collapse_bpe_window_ids(const std::vector<int32_t> & window_ids,
                             const std::vector<float> &   window_mass,
                             int                          blank_id,
                             std::vector<int32_t> &       out_token_ids);

// Host-side BPE CTC pool + greedy decode over frame-level logits. This is
// the reference formulation; run() uses the in-graph pooled head above
// (sum_t w_t (W h_t + b) = W (sum_t w_t h_t) + b because the window weights
// sum to 1) and keeps this as the documented semantics.
//
// Input:
//   importance_non_blank  [T_enc] host row of (1 - blank_prob) values used
//                         as per-frame pool weights. Must come from the
//                         *middle* (self-conditioning) CTC head's softmax,
//                         not the final head — see modeling_ctc.py.
//   ctc_bpe      [T_enc, n_bpe_vocab] host-row-major (frame-level BPE
//                logits). n_bpe_vocab is bpe_output_dim: 100353
//                (vocab_size + 1) for the old scheme, 100352 (vocab_size)
//                for the new scheme.
//   pool_window  4 in this family
//   blank_id     selects the decode scheme (see weights.h enc_bpe_blank_id):
//                  - 0   (old): channel 0 is a synthetic blank, channels
//                        1..N hold LLM ids; emitted id = argmax - 1.
//                  - 100257 (new, BOS): channels ARE the LLM ids directly,
//                        blank is the BOS id; emitted id = argmax (no shift).
// Output:
//   token_ids    initial-hypothesis BPE token ids after greedy + collapse +
//                blank removal. Used as the text portion of the NLE LLM
//                forward (each token gets an eos slot inserted around it).
//
// Reference: NLE NARDecoder.compute_text_ctc_preds. Per-frame non-blank
// frames are bucketed into windows of pool_window, the BPE logits over each
// window are weighted by the (softmaxed) CTC non-blank posterior, then a
// greedy argmax + collapse-repeats + drop-blanks yields the hypothesis. See
// the implementation in encoder.cpp.
void compute_bpe_ctc_initial_hypothesis(const std::vector<float> & importance_non_blank,
                                        const std::vector<float> & ctc_bpe_logits,
                                        int                        n_bpe_vocab,
                                        int                        T_enc,
                                        int                        pool_window,
                                        int                        blank_id,
                                        std::vector<int32_t> &     out_token_ids);

}  // namespace transcribe::granite_nar