blumia
21d68e2294
Add native streaming ASR: 8s-aligned audio slices are cached in the KV cache, only new audio is encoded, and the previous text is fed back as an assistant prefix with token rollback. Text matches offline results. Streaming keeps each step within the 2s budget for segments up to 20s. 新增原生流式解码：8s 对齐的音频切片存入 KV 复用，只编码新音频，并把 上一轮文本作为 assistant 前缀回喂（带 token 回滚）。转写结果与离线 一致；20s 以内的段落每步计算都在 2s 预算内。
484 lines
17 KiB
C++
484 lines
17 KiB
C++
#include "qwen3asr.h"
|
||
|
||
#include <chrono>
|
||
#include <cstdio>
|
||
#include <vector>
|
||
|
||
#include <llama.h>
|
||
#include <mtmd.h>
|
||
#include <mtmd-helper.h>
|
||
|
||
namespace lc {
|
||
|
||
namespace {
|
||
|
||
std::string defaultQwenChatPrompt(const std::string & content)
|
||
{
|
||
// Qwen3-ASR 的 chat template 就是标准 ChatML（见模型 GGUF 的 tokenizer.chat_template）
|
||
return "<|im_start|>user\n" + content + "<|im_end|>\n<|im_start|>assistant\n";
|
||
}
|
||
|
||
void silentLog(ggml_log_level /*level*/, const char * /*text*/, void * /*user*/) {}
|
||
|
||
} // namespace
|
||
|
||
struct Qwen3Asr::Impl {
|
||
AsrOptions opt;
|
||
llama_model * model = nullptr;
|
||
llama_context * ctx = nullptr;
|
||
mtmd_context * mctx = nullptr;
|
||
const llama_vocab * vocab = nullptr;
|
||
int audioRate = 16000;
|
||
|
||
// 流式（增量）状态
|
||
bool streamReady = false;
|
||
std::vector<llama_token> headTokens; // 音频之前的 prompt 文本
|
||
std::vector<llama_token> tailTokens; // 音频之后的 prompt 文本（含 assistant 头）
|
||
std::vector<float> streamAudio; // 当前这段的全部音频（不滑动）
|
||
size_t stableSamples = 0; // 已固化进 KV 的音频样本数（8s 对齐）
|
||
llama_pos nPastStable = 0; // stable 音频末尾的 KV 位置
|
||
|
||
std::string buildPrompt(const std::string & marker) const
|
||
{
|
||
const char * tmpl = llama_model_chat_template(model, nullptr);
|
||
if (tmpl) {
|
||
llama_chat_message msg{"user", marker.c_str()};
|
||
std::vector<char> buf(marker.size() * 2 + 4096);
|
||
const int32_t n = llama_chat_apply_template(tmpl, &msg, 1, true, buf.data(), (int32_t) buf.size());
|
||
if (n > 0 && (size_t) n <= buf.size()) {
|
||
return std::string(buf.data(), (size_t) n);
|
||
}
|
||
}
|
||
return defaultQwenChatPrompt(marker);
|
||
}
|
||
|
||
// 用显式位置解码一批 token（便于在增量流式里精确控制位置）
|
||
bool decodeTokens(const llama_token * tokens, int32_t n, llama_pos startPos, bool wantLogits)
|
||
{
|
||
if (n <= 0) return true;
|
||
llama_batch b = llama_batch_init(n, 0, 1);
|
||
for (int32_t i = 0; i < n; ++i) {
|
||
b.token[i] = tokens[i];
|
||
b.pos[i] = startPos + i;
|
||
b.n_seq_id[i] = 1;
|
||
b.seq_id[i][0] = 0;
|
||
b.logits[i] = (wantLogits && i == n - 1) ? 1 : 0;
|
||
}
|
||
b.n_tokens = n;
|
||
const int32_t rc = llama_decode(ctx, b);
|
||
llama_batch_free(b);
|
||
return rc == 0;
|
||
}
|
||
|
||
// 把一段音频编码后追加到 KV 的 nPast 处
|
||
bool appendAudio(const float * pcm, size_t nSamples, llama_pos & nPast, std::string * err)
|
||
{
|
||
if (nSamples == 0) return true;
|
||
|
||
mtmd_bitmap * bmp = mtmd_bitmap_init_from_audio(nSamples, pcm);
|
||
if (!bmp) { if (err) *err = "构造音频 bitmap 失败"; return false; }
|
||
|
||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||
const std::string marker = mtmd_default_marker();
|
||
const mtmd_input_text itext{ marker.c_str(), false, true };
|
||
const mtmd_bitmap * bms[1] = { bmp };
|
||
|
||
if (mtmd_tokenize(mctx, chunks, &itext, bms, 1) != 0) {
|
||
mtmd_bitmap_free(bmp);
|
||
mtmd_input_chunks_free(chunks);
|
||
if (err) *err = "音频 mtmd_tokenize 失败";
|
||
return false;
|
||
}
|
||
|
||
const mtmd_input_chunk * audioChunk = nullptr;
|
||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); ++i) {
|
||
const mtmd_input_chunk * c = mtmd_input_chunks_get(chunks, i);
|
||
if (mtmd_input_chunk_get_type(c) == MTMD_INPUT_CHUNK_TYPE_AUDIO) { audioChunk = c; break; }
|
||
}
|
||
if (!audioChunk) {
|
||
mtmd_bitmap_free(bmp);
|
||
mtmd_input_chunks_free(chunks);
|
||
if (err) *err = "tokenize 结果里没有音频块";
|
||
return false;
|
||
}
|
||
|
||
llama_pos np = nPast;
|
||
const int32_t rc = mtmd_helper_eval_chunk_single(mctx, ctx, audioChunk, nPast, /*seq_id*/ 0,
|
||
opt.nBatch, /*logits_last*/ false, &np);
|
||
mtmd_bitmap_free(bmp);
|
||
mtmd_input_chunks_free(chunks);
|
||
if (rc != 0) { if (err) *err = "音频编码/解码失败"; return false; }
|
||
|
||
nPast = np;
|
||
return true;
|
||
}
|
||
};
|
||
|
||
Qwen3Asr::Qwen3Asr() : d(new Impl) {}
|
||
|
||
Qwen3Asr::~Qwen3Asr()
|
||
{
|
||
if (!d) return;
|
||
if (d->mctx) mtmd_free(d->mctx);
|
||
if (d->ctx) llama_free(d->ctx);
|
||
if (d->model) llama_model_free(d->model);
|
||
delete d;
|
||
}
|
||
|
||
int Qwen3Asr::sampleRate() const { return d->audioRate; }
|
||
|
||
bool Qwen3Asr::load(const AsrOptions & opt, std::string * error)
|
||
{
|
||
const auto fail = [error](const std::string & msg) {
|
||
if (error) *error = msg;
|
||
return false;
|
||
};
|
||
|
||
d->opt = opt;
|
||
|
||
if (!opt.verbose) {
|
||
llama_log_set(silentLog, nullptr);
|
||
mtmd_log_set(silentLog, nullptr);
|
||
}
|
||
|
||
llama_backend_init();
|
||
|
||
llama_model_params mparams = llama_model_default_params();
|
||
mparams.n_gpu_layers = 0; // 本机无独显，强制 CPU
|
||
mparams.use_mmap = true;
|
||
|
||
d->model = llama_model_load_from_file(opt.modelPath.c_str(), mparams);
|
||
if (!d->model) {
|
||
return fail("加载主模型失败: " + opt.modelPath);
|
||
}
|
||
d->vocab = llama_model_get_vocab(d->model);
|
||
|
||
llama_context_params cparams = llama_context_default_params();
|
||
cparams.n_ctx = (uint32_t) opt.nCtx;
|
||
cparams.n_batch = (uint32_t) opt.nBatch;
|
||
cparams.n_ubatch = (uint32_t) opt.nUbatch;
|
||
cparams.n_threads = opt.threads;
|
||
cparams.n_threads_batch = opt.threads;
|
||
|
||
d->ctx = llama_init_from_model(d->model, cparams);
|
||
if (!d->ctx) {
|
||
return fail("创建 llama 上下文失败");
|
||
}
|
||
|
||
mtmd_context_params mctxParams = mtmd_context_params_default();
|
||
mctxParams.use_gpu = false;
|
||
mctxParams.n_threads = opt.threads;
|
||
mctxParams.print_timings = false;
|
||
mctxParams.warmup = opt.warmup;
|
||
|
||
d->mctx = mtmd_init_from_file(opt.mmprojPath.c_str(), d->model, mctxParams);
|
||
if (!d->mctx) {
|
||
return fail("加载 mmproj 失败: " + opt.mmprojPath);
|
||
}
|
||
|
||
if (!mtmd_support_audio(d->mctx)) {
|
||
return fail("该 mmproj 不支持音频输入（需要 clip.projector_type = qwen3a 的版本）");
|
||
}
|
||
|
||
const int rate = mtmd_get_audio_sample_rate(d->mctx);
|
||
if (rate > 0) d->audioRate = rate;
|
||
|
||
return true;
|
||
}
|
||
|
||
Qwen3Asr::Result Qwen3Asr::transcribe(const float * pcm, size_t nSamples)
|
||
{
|
||
return transcribeWithPrefix(pcm, nSamples, std::string());
|
||
}
|
||
|
||
std::string Qwen3Asr::dropLastTokens(const std::string & text, int dropLast) const
|
||
{
|
||
if (!d->vocab || dropLast <= 0 || text.empty()) return text;
|
||
|
||
std::vector<llama_token> tokens(text.size() + 16);
|
||
int32_t n = llama_tokenize(d->vocab, text.c_str(), (int32_t) text.size(),
|
||
tokens.data(), (int32_t) tokens.size(),
|
||
/*add_special*/ false, /*parse_special*/ true);
|
||
if (n < 0) { // 缓冲区不够，按返回的所需长度重来
|
||
tokens.resize((size_t) -n);
|
||
n = llama_tokenize(d->vocab, text.c_str(), (int32_t) text.size(),
|
||
tokens.data(), (int32_t) tokens.size(), false, true);
|
||
if (n < 0) return text;
|
||
}
|
||
tokens.resize((size_t) n);
|
||
if ((int) tokens.size() <= dropLast) return {};
|
||
|
||
tokens.resize(tokens.size() - (size_t) dropLast);
|
||
|
||
// special=true：让 <asr_text> 这类特殊标记以字面形式保留，前缀才能原样回喂
|
||
std::string out;
|
||
char buf[512];
|
||
for (llama_token id : tokens) {
|
||
const int32_t m = llama_token_to_piece(d->vocab, id, buf, sizeof(buf), 0, /*special*/ true);
|
||
if (m > 0) out.append(buf, (size_t) m);
|
||
}
|
||
return out;
|
||
}
|
||
|
||
Qwen3Asr::Result Qwen3Asr::transcribeWithPrefix(const float * pcm, size_t nSamples,
|
||
const std::string & prefix)
|
||
{
|
||
Result res;
|
||
|
||
if (!d->model || !d->ctx || !d->mctx) {
|
||
res.error = "模型未加载";
|
||
return res;
|
||
}
|
||
if (nSamples == 0) {
|
||
res.error = "音频为空";
|
||
return res;
|
||
}
|
||
|
||
const auto t0 = std::chrono::steady_clock::now();
|
||
|
||
// 每段独立转写，先清空 KV cache
|
||
llama_memory_clear(llama_get_memory(d->ctx), true);
|
||
|
||
const std::string marker = mtmd_default_marker();
|
||
// 关键：assistant 头之后接上"已生成文本"，让模型从那里继续生成（前缀回滚）
|
||
const std::string prompt = d->buildPrompt(marker) + prefix;
|
||
|
||
mtmd_bitmap * bitmap = mtmd_bitmap_init_from_audio(nSamples, pcm);
|
||
if (!bitmap) {
|
||
res.error = "构造音频 bitmap 失败";
|
||
return res;
|
||
}
|
||
|
||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||
const mtmd_input_text text{ prompt.c_str(), /*add_special*/ true, /*parse_special*/ true };
|
||
const mtmd_bitmap * bitmaps[1] = { bitmap };
|
||
|
||
const int32_t tokRes = mtmd_tokenize(d->mctx, chunks, &text, bitmaps, 1);
|
||
if (tokRes != 0) {
|
||
mtmd_bitmap_free(bitmap);
|
||
mtmd_input_chunks_free(chunks);
|
||
res.error = "mtmd_tokenize 失败 (code " + std::to_string(tokRes) + ")";
|
||
return res;
|
||
}
|
||
|
||
llama_pos nPast = 0;
|
||
const int32_t evalRes = mtmd_helper_eval_chunks(
|
||
d->mctx, d->ctx, chunks,
|
||
/*n_past*/ 0, /*seq_id*/ 0, /*n_batch*/ d->opt.nBatch,
|
||
/*logits_last*/ true, &nPast);
|
||
mtmd_bitmap_free(bitmap);
|
||
if (evalRes != 0) {
|
||
mtmd_input_chunks_free(chunks);
|
||
res.error = "mtmd_helper_eval_chunks 失败 (code " + std::to_string(evalRes) + ")";
|
||
return res;
|
||
}
|
||
mtmd_input_chunks_free(chunks);
|
||
|
||
// 贪心解码
|
||
llama_sampler * smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
|
||
llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
|
||
|
||
std::string out;
|
||
for (int i = 0; i < d->opt.maxTokens; ++i) {
|
||
llama_token id = llama_sampler_sample(smpl, d->ctx, -1);
|
||
if (llama_vocab_is_eog(d->vocab, id)) break;
|
||
|
||
char buf[256];
|
||
const int32_t n = llama_token_to_piece(d->vocab, id, buf, sizeof(buf), 0, /*special*/ true);
|
||
if (n > 0) out.append(buf, (size_t) n);
|
||
++res.tokenCount;
|
||
|
||
llama_batch batch = llama_batch_get_one(&id, 1);
|
||
if (llama_decode(d->ctx, batch) != 0) {
|
||
res.error = "llama_decode 失败";
|
||
break;
|
||
}
|
||
}
|
||
llama_sampler_free(smpl);
|
||
|
||
const auto t1 = std::chrono::steady_clock::now();
|
||
res.inferSeconds = std::chrono::duration<double>(t1 - t0).count();
|
||
res.raw = out;
|
||
res.transcript = parseAsrOutput(out);
|
||
res.ok = res.error.empty();
|
||
return res;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 流式（增量 KV 复用）
|
||
// ---------------------------------------------------------------------------
|
||
|
||
void Qwen3Asr::streamReset()
|
||
{
|
||
if (!d->ctx) return;
|
||
llama_memory_clear(llama_get_memory(d->ctx), true);
|
||
d->streamReady = false;
|
||
d->headTokens.clear();
|
||
d->tailTokens.clear();
|
||
d->streamAudio.clear();
|
||
d->stableSamples = 0;
|
||
d->nPastStable = 0;
|
||
}
|
||
|
||
bool Qwen3Asr::prepareStream(std::string * error)
|
||
{
|
||
if (d->streamReady) return true;
|
||
|
||
// 用一个占位音频 tokenize 一次，只为切出「音频之前」和「音频之后」的 prompt 文本 token。
|
||
// 占位音频本身不会被解码。
|
||
std::vector<float> dummy(1600, 0.0f); // 0.1s
|
||
mtmd_bitmap * bmp = mtmd_bitmap_init_from_audio(dummy.size(), dummy.data());
|
||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||
|
||
const std::string marker = mtmd_default_marker();
|
||
const std::string prompt = d->buildPrompt(marker);
|
||
const mtmd_input_text itext{ prompt.c_str(), true, true };
|
||
const mtmd_bitmap * bms[1] = { bmp };
|
||
|
||
if (mtmd_tokenize(d->mctx, chunks, &itext, bms, 1) != 0) {
|
||
mtmd_bitmap_free(bmp);
|
||
mtmd_input_chunks_free(chunks);
|
||
if (error) *error = "prepareStream: mtmd_tokenize 失败";
|
||
return false;
|
||
}
|
||
mtmd_bitmap_free(bmp);
|
||
|
||
bool seenAudio = false;
|
||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); ++i) {
|
||
const mtmd_input_chunk * c = mtmd_input_chunks_get(chunks, i);
|
||
if (mtmd_input_chunk_get_type(c) == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||
seenAudio = true;
|
||
continue;
|
||
}
|
||
size_t n = 0;
|
||
const llama_token * toks = mtmd_input_chunk_get_tokens_text(c, &n);
|
||
if (!toks || n == 0) continue;
|
||
std::vector<llama_token> & dst = seenAudio ? d->tailTokens : d->headTokens;
|
||
dst.insert(dst.end(), toks, toks + n);
|
||
}
|
||
mtmd_input_chunks_free(chunks);
|
||
|
||
if (d->headTokens.empty() || d->tailTokens.empty()) {
|
||
if (error) *error = "prepareStream: 没能切出音频前后的 prompt 文本";
|
||
return false;
|
||
}
|
||
|
||
llama_memory_clear(llama_get_memory(d->ctx), true);
|
||
if (!d->decodeTokens(d->headTokens.data(), (int32_t) d->headTokens.size(), 0, false)) {
|
||
if (error) *error = "prepareStream: 解码 prompt 头部失败";
|
||
return false;
|
||
}
|
||
d->nPastStable = (llama_pos) d->headTokens.size();
|
||
d->streamReady = true;
|
||
return true;
|
||
}
|
||
|
||
// 8 秒对齐的音频切片：这是编码器的注意力窗口（n_window_infer = 800 帧 @100Hz）。
|
||
// 按这个粒度切，编码结果才和「整段音频一次编码」一致；切小了会明显掉准确率。
|
||
static constexpr double kSliceSeconds = 8.0;
|
||
|
||
Qwen3Asr::StreamStep Qwen3Asr::streamAppend(const float * pcm, size_t nSamples,
|
||
const std::string & prefix, int maxTokens)
|
||
{
|
||
StreamStep res;
|
||
if (!d->model || !d->ctx || !d->mctx) { res.error = "模型未加载"; return res; }
|
||
|
||
std::string err;
|
||
if (!prepareStream(&err)) { res.error = err; return res; }
|
||
|
||
const auto t0 = std::chrono::steady_clock::now();
|
||
|
||
d->streamAudio.insert(d->streamAudio.end(), pcm, pcm + nSamples);
|
||
|
||
const size_t sliceSamples = (size_t) (kSliceSeconds * d->audioRate);
|
||
const size_t targetStable = (d->streamAudio.size() / sliceSamples) * sliceSamples;
|
||
|
||
const llama_pos nPastBefore = d->nPastStable;
|
||
|
||
// 1) 已排满 8s 的切片固化进 KV，之后永远复用
|
||
if (targetStable > d->stableSamples) {
|
||
llama_memory_seq_rm(llama_get_memory(d->ctx), 0, d->nPastStable, -1);
|
||
llama_pos np = d->nPastStable;
|
||
if (!d->appendAudio(d->streamAudio.data() + d->stableSamples,
|
||
targetStable - d->stableSamples, np, &err)) {
|
||
res.error = "固化音频切片失败: " + err;
|
||
return res;
|
||
}
|
||
d->nPastStable = np;
|
||
d->stableSamples = targetStable;
|
||
}
|
||
|
||
// 2) 尾部不满 8s 的部分：每次重编码，但只是临时的（下一步会被丢掉）
|
||
llama_memory_seq_rm(llama_get_memory(d->ctx), 0, d->nPastStable, -1);
|
||
llama_pos nPast = d->nPastStable;
|
||
if (d->streamAudio.size() > d->stableSamples) {
|
||
if (!d->appendAudio(d->streamAudio.data() + d->stableSamples,
|
||
d->streamAudio.size() - d->stableSamples, nPast, &err)) {
|
||
res.error = "编码尾部音频失败: " + err;
|
||
return res;
|
||
}
|
||
}
|
||
res.audioTokens = (int) (nPast - nPastBefore);
|
||
const auto tEnc = std::chrono::steady_clock::now();
|
||
|
||
// 3) 音频之后的 prompt 文本 + 前缀
|
||
std::vector<llama_token> prefixToks;
|
||
if (!prefix.empty()) {
|
||
prefixToks.resize(prefix.size() + 16);
|
||
int32_t n = llama_tokenize(d->vocab, prefix.c_str(), (int32_t) prefix.size(),
|
||
prefixToks.data(), (int32_t) prefixToks.size(),
|
||
/*add_special*/ false, /*parse_special*/ true);
|
||
if (n < 0) {
|
||
prefixToks.resize((size_t) -n);
|
||
n = llama_tokenize(d->vocab, prefix.c_str(), (int32_t) prefix.size(),
|
||
prefixToks.data(), (int32_t) prefixToks.size(), false, true);
|
||
}
|
||
if (n < 0) { res.error = "前缀分词失败"; return res; }
|
||
prefixToks.resize((size_t) n);
|
||
}
|
||
|
||
llama_pos p = nPast;
|
||
const bool prefixLast = prefixToks.empty();
|
||
if (!d->decodeTokens(d->tailTokens.data(), (int32_t) d->tailTokens.size(), p, prefixLast)) {
|
||
res.error = "解码 tail 失败";
|
||
return res;
|
||
}
|
||
p += (llama_pos) d->tailTokens.size();
|
||
|
||
if (!prefixToks.empty()) {
|
||
if (!d->decodeTokens(prefixToks.data(), (int32_t) prefixToks.size(), p, true)) {
|
||
res.error = "解码前缀失败";
|
||
return res;
|
||
}
|
||
p += (llama_pos) prefixToks.size();
|
||
}
|
||
|
||
// 4) 贪心生成（每个 token 的 logits 都要留给下一次采样）
|
||
llama_sampler * smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
|
||
llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
|
||
|
||
std::string out;
|
||
for (int i = 0; i < maxTokens; ++i) {
|
||
llama_token id = llama_sampler_sample(smpl, d->ctx, -1);
|
||
if (llama_vocab_is_eog(d->vocab, id)) break;
|
||
|
||
char buf[256];
|
||
const int32_t n = llama_token_to_piece(d->vocab, id, buf, sizeof(buf), 0, /*special*/ true);
|
||
if (n > 0) out.append(buf, (size_t) n);
|
||
++res.newTokens;
|
||
|
||
if (!d->decodeTokens(&id, 1, p, true)) { res.error = "生成时解码失败"; break; }
|
||
++p;
|
||
}
|
||
llama_sampler_free(smpl);
|
||
|
||
const auto t1 = std::chrono::steady_clock::now();
|
||
res.encodeSeconds = std::chrono::duration<double>(tEnc - t0).count();
|
||
res.inferSeconds = std::chrono::duration<double>(t1 - t0).count();
|
||
res.newText = out;
|
||
res.ok = res.error.empty();
|
||
return res;
|
||
}
|
||
|
||
} // namespace lc
|