live-caption/src/qwen3asr.cpp

#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;

    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);
    }
};

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)
{
    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();
    const std::string prompt = d->buildPrompt(marker);

    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;
}

} // namespace lc