wehub-resource-sync
276 lines
14 KiB
C++
276 lines
14 KiB
C++
// funasr-encoder: ggml C++ forward pass for the Fun-ASR-Nano audio encoder
|
|
// (SenseVoice SAN-M, 50+20 layers) + Transformer adaptor.
|
|
//
|
|
// Input : fbank.bin (T x 560 f32, the encoder input features)
|
|
// Output: out.bin (T' x 1024 f32, audio embeddings for the LLM)
|
|
// Weights: funasr-encoder.gguf (exported by export_encoder_gguf.py)
|
|
//
|
|
// Validated layer-by-layer against PyTorch golden dumps. fbank is currently
|
|
// produced in Python; porting the fbank frontend to C++ is the remaining piece.
|
|
|
|
#include "ggml.h"
|
|
#include "ggml-cpu.h"
|
|
#include "ggml-alloc.h"
|
|
#include "ggml-backend.h"
|
|
#include "gguf.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <map>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
struct cfg {
|
|
int input_size = 560, d_model = 512, n_head = 4, ffn = 2048;
|
|
int num_blocks = 50, tp_blocks = 20, kernel = 11;
|
|
int adp_llm = 1024, adp_ffn = 2048, adp_layers = 2, adp_head = 8;
|
|
};
|
|
|
|
static const float LN_EPS = 1e-5f;
|
|
|
|
struct funasr_model {
|
|
cfg c;
|
|
struct ggml_context * ctx_w = nullptr; // weights (CPU malloc, data set)
|
|
std::map<std::string, struct ggml_tensor *> t;
|
|
struct ggml_tensor * get(const std::string & n) {
|
|
auto it = t.find(n);
|
|
if (it == t.end()) { fprintf(stderr, "missing tensor: %s\n", n.c_str()); exit(1); }
|
|
return it->second;
|
|
}
|
|
};
|
|
|
|
static bool load_model(const char * path, funasr_model & m) {
|
|
struct gguf_init_params p = { /*no_alloc=*/false, /*ctx=*/&m.ctx_w };
|
|
struct gguf_context * gguf = gguf_init_from_file(path, p);
|
|
if (!gguf) { fprintf(stderr, "failed to load gguf %s\n", path); return false; }
|
|
auto rd = [&](const char * k, int def) {
|
|
int i = gguf_find_key(gguf, k); return i < 0 ? def : (int) gguf_get_val_u32(gguf, i);
|
|
};
|
|
m.c.input_size = rd("funasr.enc.input_size", 560);
|
|
m.c.d_model = rd("funasr.enc.output_size", 512);
|
|
m.c.n_head = rd("funasr.enc.attention_heads", 4);
|
|
m.c.ffn = rd("funasr.enc.linear_units", 2048);
|
|
m.c.num_blocks = rd("funasr.enc.num_blocks", 50);
|
|
m.c.tp_blocks = rd("funasr.enc.tp_blocks", 20);
|
|
m.c.kernel = rd("funasr.enc.kernel_size", 11);
|
|
m.c.adp_llm = rd("funasr.adp.llm_dim", 1024);
|
|
m.c.adp_ffn = rd("funasr.adp.ffn_dim", 2048);
|
|
m.c.adp_layers = rd("funasr.adp.n_layer", 2);
|
|
m.c.adp_head = rd("funasr.adp.attention_heads", 8);
|
|
int n = gguf_get_n_tensors(gguf);
|
|
for (int i = 0; i < n; i++) {
|
|
const char * name = gguf_get_tensor_name(gguf, i);
|
|
m.t[name] = ggml_get_tensor(m.ctx_w, name);
|
|
}
|
|
fprintf(stderr, "loaded %d tensors; cfg: d_model=%d heads=%d blocks=%d tp=%d kernel=%d adp_llm=%d\n",
|
|
n, m.c.d_model, m.c.n_head, m.c.num_blocks, m.c.tp_blocks, m.c.kernel, m.c.adp_llm);
|
|
gguf_free(gguf);
|
|
return true;
|
|
}
|
|
|
|
// helpers ---------------------------------------------------------------
|
|
static struct ggml_tensor * linear(ggml_context * ctx, ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) {
|
|
struct ggml_tensor * y = ggml_mul_mat(ctx, w, x); // [out, T]
|
|
if (b) y = ggml_add(ctx, y, b);
|
|
return y;
|
|
}
|
|
static struct ggml_tensor * layernorm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * g, ggml_tensor * b) {
|
|
x = ggml_norm(ctx, x, LN_EPS);
|
|
x = ggml_mul(ctx, x, g);
|
|
x = ggml_add(ctx, x, b);
|
|
return x;
|
|
}
|
|
|
|
// SAN-M self-attention + FSMN. x:[D_in,T] -> [d_model,T]
|
|
static struct ggml_tensor * sanm_attn(ggml_context * ctx, funasr_model & m, const std::string & pfx,
|
|
ggml_tensor * x, int T) {
|
|
const int D = m.c.d_model, H = m.c.n_head, dk = D / H, K = m.c.kernel;
|
|
struct ggml_tensor * qkv = linear(ctx, m.get(pfx + "linear_q_k_v.weight"),
|
|
m.get(pfx + "linear_q_k_v.bias"), x); // [3D, T]
|
|
size_t nb1 = qkv->nb[1];
|
|
struct ggml_tensor * q = ggml_cont(ctx, ggml_view_2d(ctx, qkv, D, T, nb1, 0));
|
|
struct ggml_tensor * k = ggml_cont(ctx, ggml_view_2d(ctx, qkv, D, T, nb1, (size_t) D * sizeof(float)));
|
|
struct ggml_tensor * v = ggml_cont(ctx, ggml_view_2d(ctx, qkv, D, T, nb1, (size_t) 2 * D * sizeof(float)));
|
|
|
|
// FSMN: depthwise conv1d along time (per-channel kernel K, "same" padding),
|
|
// plus residual v. Implemented as an exact f32 shift-accumulate to avoid the
|
|
// F16-only ggml_conv_1d_dw path. fsmn kernel stored as [D, K] (ne0=D, ne1=K).
|
|
const int pad = (K - 1) / 2;
|
|
struct ggml_tensor * fk = m.get(pfx + "fsmn_block.weight"); // [D, K]
|
|
struct ggml_tensor * vpad = ggml_pad_ext(ctx, v, 0, 0, pad, pad, 0, 0, 0, 0); // [D, T+2*pad]
|
|
struct ggml_tensor * fsmn = v; // residual
|
|
for (int j = 0; j < K; j++) {
|
|
struct ggml_tensor * sl = ggml_view_2d(ctx, vpad, D, T, vpad->nb[1], (size_t) j * vpad->nb[1]);
|
|
struct ggml_tensor * wj = ggml_view_1d(ctx, fk, D, (size_t) j * fk->nb[1]);
|
|
fsmn = ggml_add(ctx, fsmn, ggml_mul(ctx, ggml_cont(ctx, sl), wj));
|
|
}
|
|
|
|
// multi-head attention
|
|
q = ggml_reshape_3d(ctx, q, dk, H, T);
|
|
k = ggml_reshape_3d(ctx, k, dk, H, T);
|
|
struct ggml_tensor * vh = ggml_reshape_3d(ctx, v, dk, H, T);
|
|
q = ggml_permute(ctx, q, 0, 2, 1, 3); // [dk, T, H]
|
|
k = ggml_permute(ctx, k, 0, 2, 1, 3); // [dk, T, H]
|
|
vh = ggml_cont(ctx, ggml_permute(ctx, vh, 1, 2, 0, 3)); // [T, dk, H]
|
|
struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q); // [T, T, H]
|
|
kq = ggml_scale(ctx, kq, 1.0f / sqrtf((float) dk));
|
|
kq = ggml_soft_max(ctx, kq);
|
|
struct ggml_tensor * kqv = ggml_mul_mat(ctx, vh, kq); // [dk, T, H]
|
|
kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3); // [dk, H, T]
|
|
kqv = ggml_cont_2d(ctx, kqv, D, T); // [D, T]
|
|
struct ggml_tensor * att = linear(ctx, m.get(pfx + "linear_out.weight"),
|
|
m.get(pfx + "linear_out.bias"), kqv);
|
|
return ggml_add(ctx, att, fsmn);
|
|
}
|
|
|
|
// one SAN-M encoder layer. in_size may differ from d_model (first layer)
|
|
static struct ggml_tensor * sanm_layer(ggml_context * ctx, funasr_model & m, const std::string & pfx,
|
|
ggml_tensor * x, int T, bool residual_attn) {
|
|
struct ggml_tensor * res = x;
|
|
struct ggml_tensor * h = layernorm(ctx, x, m.get(pfx + "norm1.weight"), m.get(pfx + "norm1.bias"));
|
|
struct ggml_tensor * sa = sanm_attn(ctx, m, pfx + "self_attn.", h, T);
|
|
x = residual_attn ? ggml_add(ctx, res, sa) : sa;
|
|
res = x;
|
|
h = layernorm(ctx, x, m.get(pfx + "norm2.weight"), m.get(pfx + "norm2.bias"));
|
|
h = linear(ctx, m.get(pfx + "feed_forward.w_1.weight"), m.get(pfx + "feed_forward.w_1.bias"), h);
|
|
h = ggml_relu(ctx, h);
|
|
h = linear(ctx, m.get(pfx + "feed_forward.w_2.weight"), m.get(pfx + "feed_forward.w_2.bias"), h);
|
|
return ggml_add(ctx, res, h);
|
|
}
|
|
|
|
// standard transformer layer (adaptor). d=adp_llm
|
|
static struct ggml_tensor * adp_layer(ggml_context * ctx, funasr_model & m, const std::string & pfx,
|
|
ggml_tensor * x, int T) {
|
|
const int D = m.c.adp_llm, H = m.c.adp_head, dk = D / H;
|
|
struct ggml_tensor * res = x;
|
|
struct ggml_tensor * h = layernorm(ctx, x, m.get(pfx + "norm1.weight"), m.get(pfx + "norm1.bias"));
|
|
struct ggml_tensor * q = linear(ctx, m.get(pfx + "self_attn.linear_q.weight"), m.get(pfx + "self_attn.linear_q.bias"), h);
|
|
struct ggml_tensor * k = linear(ctx, m.get(pfx + "self_attn.linear_k.weight"), m.get(pfx + "self_attn.linear_k.bias"), h);
|
|
struct ggml_tensor * v = linear(ctx, m.get(pfx + "self_attn.linear_v.weight"), m.get(pfx + "self_attn.linear_v.bias"), h);
|
|
q = ggml_permute(ctx, ggml_reshape_3d(ctx, q, dk, H, T), 0, 2, 1, 3);
|
|
k = ggml_permute(ctx, ggml_reshape_3d(ctx, k, dk, H, T), 0, 2, 1, 3);
|
|
struct ggml_tensor * vh = ggml_cont(ctx, ggml_permute(ctx, ggml_reshape_3d(ctx, v, dk, H, T), 1, 2, 0, 3));
|
|
struct ggml_tensor * kq = ggml_soft_max(ctx, ggml_scale(ctx, ggml_mul_mat(ctx, k, q), 1.0f / sqrtf((float) dk)));
|
|
struct ggml_tensor * kqv = ggml_cont_2d(ctx, ggml_permute(ctx, ggml_mul_mat(ctx, vh, kq), 0, 2, 1, 3), D, T);
|
|
struct ggml_tensor * att = linear(ctx, m.get(pfx + "self_attn.linear_out.weight"), m.get(pfx + "self_attn.linear_out.bias"), kqv);
|
|
x = ggml_add(ctx, res, att);
|
|
res = x;
|
|
h = layernorm(ctx, x, m.get(pfx + "norm2.weight"), m.get(pfx + "norm2.bias"));
|
|
h = linear(ctx, m.get(pfx + "feed_forward.w_1.weight"), m.get(pfx + "feed_forward.w_1.bias"), h);
|
|
h = ggml_relu(ctx, h);
|
|
h = linear(ctx, m.get(pfx + "feed_forward.w_2.weight"), m.get(pfx + "feed_forward.w_2.bias"), h);
|
|
return ggml_add(ctx, res, h);
|
|
}
|
|
|
|
// sinusoidal position encoding, depth = input feature dim, positions 1..T
|
|
static void add_posenc(std::vector<float> & x, int T, int depth) {
|
|
double inc = log(10000.0) / (depth / 2.0 - 1.0);
|
|
for (int t = 0; t < T; t++) {
|
|
double pos = t + 1; // positions start at 1
|
|
for (int i = 0; i < depth / 2; i++) {
|
|
double its = exp(i * -inc);
|
|
double st = pos * its;
|
|
x[(size_t) t * depth + i] += (float) sin(st);
|
|
x[(size_t) t * depth + depth / 2 + i] += (float) cos(st);
|
|
}
|
|
}
|
|
}
|
|
|
|
int main(int argc, char ** argv) {
|
|
std::string gguf_path, fbank_path, out_path = "out.bin";
|
|
int limit = -1; // -L: run only first N (encoders0+encoders) layers, dump running x
|
|
bool run_adaptor = true;
|
|
for (int i = 1; i < argc; i++) {
|
|
if (!strcmp(argv[i], "-m") && i+1 < argc) gguf_path = argv[++i];
|
|
else if (!strcmp(argv[i], "-f") && i+1 < argc) fbank_path = argv[++i];
|
|
else if (!strcmp(argv[i], "-o") && i+1 < argc) out_path = argv[++i];
|
|
else if (!strcmp(argv[i], "-L") && i+1 < argc) { limit = atoi(argv[++i]); run_adaptor = false; }
|
|
else { fprintf(stderr, "usage: %s -m enc.gguf -f fbank.bin [-o out.bin] [-L nlayers]\n", argv[0]); return 1; }
|
|
}
|
|
|
|
funasr_model m;
|
|
if (!load_model(gguf_path.c_str(), m)) return 1;
|
|
|
|
// read fbank.bin (T x F)
|
|
FILE * f = fopen(fbank_path.c_str(), "rb");
|
|
if (!f) { fprintf(stderr, "cannot open %s\n", fbank_path.c_str()); return 1; }
|
|
int32_t T, F; if (fread(&T, 4, 1, f) != 1 || fread(&F, 4, 1, f) != 1) { fclose(f); return 1; }
|
|
std::vector<float> fbank((size_t) T * F);
|
|
if (fread(fbank.data(), sizeof(float), fbank.size(), f) != fbank.size()) { fclose(f); return 1; }
|
|
fclose(f);
|
|
fprintf(stderr, "fbank: T=%d F=%d\n", T, F);
|
|
|
|
// pre-scale (*sqrt(d_model)) and add position encoding on the host
|
|
float scale = sqrtf((float) m.c.d_model);
|
|
for (auto & v : fbank) v *= scale;
|
|
add_posenc(fbank, T, F);
|
|
|
|
// backend + compute context
|
|
ggml_backend_t backend = ggml_backend_cpu_init();
|
|
size_t ctx_size = (size_t) 1024*1024*1024; // graph metadata
|
|
struct ggml_init_params cp = { ctx_size, nullptr, true };
|
|
struct ggml_context * ctx = ggml_init(cp);
|
|
|
|
struct ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, F, T);
|
|
ggml_set_name(inp, "inp");
|
|
ggml_set_input(inp);
|
|
|
|
struct ggml_tensor * x = inp;
|
|
int done = 0;
|
|
bool stop = false;
|
|
// encoders0 (1 layer, in_size=input_size != d_model -> no attn residual)
|
|
x = sanm_layer(ctx, m, "audio_encoder.encoders0.0.", x, T, /*residual_attn=*/false);
|
|
done++;
|
|
if (limit >= 0 && done >= limit) stop = true;
|
|
// encoders (num_blocks-1 layers)
|
|
for (int i = 0; i < m.c.num_blocks - 1 && !stop; i++) {
|
|
x = sanm_layer(ctx, m, "audio_encoder.encoders." + std::to_string(i) + ".", x, T, true);
|
|
done++;
|
|
if (limit >= 0 && done >= limit) stop = true;
|
|
}
|
|
if (!stop) {
|
|
x = layernorm(ctx, x, m.get("audio_encoder.after_norm.weight"), m.get("audio_encoder.after_norm.bias"));
|
|
for (int i = 0; i < m.c.tp_blocks; i++)
|
|
x = sanm_layer(ctx, m, "audio_encoder.tp_encoders." + std::to_string(i) + ".", x, T, true);
|
|
x = layernorm(ctx, x, m.get("audio_encoder.tp_norm.weight"), m.get("audio_encoder.tp_norm.bias"));
|
|
// adaptor: downsample_rate=1 -> linear1(relu)linear2 then blocks
|
|
if (run_adaptor) {
|
|
x = linear(ctx, m.get("audio_adaptor.linear1.weight"), m.get("audio_adaptor.linear1.bias"), x);
|
|
x = ggml_relu(ctx, x);
|
|
x = linear(ctx, m.get("audio_adaptor.linear2.weight"), m.get("audio_adaptor.linear2.bias"), x);
|
|
for (int i = 0; i < m.c.adp_layers; i++)
|
|
x = adp_layer(ctx, m, "audio_adaptor.blocks." + std::to_string(i) + ".", x, T);
|
|
}
|
|
}
|
|
ggml_set_output(x);
|
|
|
|
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false);
|
|
ggml_build_forward_expand(gf, x);
|
|
|
|
ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type());
|
|
ggml_gallocr_alloc_graph(galloc, gf);
|
|
ggml_backend_tensor_set(inp, fbank.data(), 0, ggml_nbytes(inp));
|
|
|
|
ggml_backend_cpu_set_n_threads(backend, 8);
|
|
int64_t t0 = ggml_time_us();
|
|
if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) {
|
|
fprintf(stderr, "compute failed\n"); return 1;
|
|
}
|
|
int64_t t1 = ggml_time_us();
|
|
|
|
int D = (int) x->ne[0];
|
|
std::vector<float> out((size_t) D * T);
|
|
ggml_backend_tensor_get(x, out.data(), 0, ggml_nbytes(x));
|
|
FILE * fo = fopen(out_path.c_str(), "wb");
|
|
if (!fo) { fprintf(stderr, "failed to open output file %s\n", out_path.c_str()); return 1; }
|
|
fwrite(&T, 4, 1, fo); fwrite(&D, 4, 1, fo); fwrite(out.data(), sizeof(float), out.size(), fo);
|
|
fclose(fo);
|
|
fprintf(stderr, "done: wrote %s [%d x %d] in %.2f s (layers run=%d, adaptor=%d)\n",
|
|
out_path.c_str(), T, D, (t1 - t0)/1e6, done, run_adaptor && !stop);
|
|
ggml_gallocr_free(galloc); ggml_free(ctx); ggml_backend_free(backend);
|
|
if (m.ctx_w) ggml_free(m.ctx_w);
|
|
return 0;
|
|
}
|