Kazeia-engine/dist/jni/stt_engine.cpp

882 lines
39 KiB
C++
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

// Implementation stt_engine — voir stt_engine.h pour l'API publique.
//
// Port C++ complet de WhisperHybridEngine.kt (527L Kotlin -> ce fichier).
// Backend : ONNX Runtime 1.24.3 + QNN ExecutionProvider (HTP V79 sur SM8750).
// Modèle : Qualcomm AI Hub HfWhisper KV-cache decoder.
//
// Pipeline transcribe :
// 1. PCM16 -> mel via kazeia_mel (cfg Whisper) -> [1, 80, 3000] fp16
// 2. Encoder (NPU) : input_features -> N cross KV caches fp16
// 3. Decoder (NPU) autoregressif, KV-cache, mask R-to-L, 200 steps max :
// input_ids[1,1] + attention_mask[1,1,1,200] + N self KV + N cross KV + position_ids
// -> logits[1,51865,1,1] fp16 + N updated self KV
// 4. argmax logits FP16 (override <|translate|> -> <|transcribe|>, EOT -> stop)
// 5. Token decode via BPE byte-level (vocab.json) -> texte UTF-8
//
// Auto-detect dimensions (num_decoder_layers, num_decoder_heads) via session outputs.
// Bench prod ref : Whisper-Small FR 1.6s -> 825 ms total (mel 189 + enc 125 + dec 510), RTF 0.51.
#include "stt_engine.h"
#include "kazeia_mel.h"
#include "onnxruntime/onnxruntime_cxx_api.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <chrono>
#include <fstream>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
// ============================================================================
// FP16 helpers (cohabite avec Ort::Float16_t qui est juste un wrapper uint16_t)
// ============================================================================
// IEEE 754 binary16 conversion (round-to-nearest-even).
static uint16_t fp32_to_fp16(float f) {
uint32_t x; std::memcpy(&x, &f, 4);
uint32_t sign = (x >> 16) & 0x8000u;
int32_t exp = (int32_t)((x >> 23) & 0xFFu) - 127 + 15;
uint32_t mant = x & 0x7FFFFFu;
if (exp >= 31) { // inf / nan
return (uint16_t)(sign | 0x7C00u | (mant ? 1u : 0u));
} else if (exp <= 0) { // subnormal / underflow
if (exp < -10) return (uint16_t)sign;
mant |= 0x800000u;
uint32_t shift = (uint32_t)(14 - exp);
uint32_t round = (mant >> (shift - 1)) & 1u;
return (uint16_t)(sign | ((mant >> shift) + round));
} else {
uint32_t round = (mant >> 12) & 1u;
return (uint16_t)(sign | (uint32_t)(exp << 10) | (mant >> 13) + round);
}
}
static float fp16_to_fp32(uint16_t h) {
uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
uint32_t exp = (h >> 10) & 0x1Fu;
uint32_t mant = h & 0x3FFu;
uint32_t out;
if (exp == 0) {
if (mant == 0) { out = sign; }
else {
// subnormal
int e = -1;
while (!(mant & 0x400u)) { mant <<= 1; e--; }
mant &= 0x3FFu;
out = sign | (uint32_t)((127 + e) << 23) | (mant << 13);
}
} else if (exp == 31) {
out = sign | 0x7F800000u | (mant << 13);
} else {
out = sign | (uint32_t)((127 - 15 + (int)exp) << 23) | (mant << 13);
}
float f; std::memcpy(&f, &out, 4);
return f;
}
// ============================================================================
// VAD RMS énergie (port C++ de VadStage.kt) -- identique à S1.3
// ============================================================================
struct SttVadState {
int sample_rate, frame_size, rms_threshold, min_speech_frames, silence_end_frames;
int speech_count = 0, silence_count = 0;
bool in_speech = false, end_of_speech = false;
};
SttVadState * stt_vad_new(int sample_rate, int frame_size, int rms_threshold,
int min_speech_frames, int silence_end_frames) {
auto * s = new SttVadState();
s->sample_rate = sample_rate; s->frame_size = frame_size;
s->rms_threshold = rms_threshold;
s->min_speech_frames = min_speech_frames; s->silence_end_frames = silence_end_frames;
return s;
}
void stt_vad_reset(SttVadState * s) {
if (!s) return;
s->speech_count = 0; s->silence_count = 0;
s->in_speech = false; s->end_of_speech = false;
}
void stt_vad_free(SttVadState * s) { delete s; }
bool stt_vad_is_speech(const SttVadState * s) { return s && s->in_speech; }
bool stt_vad_is_end_of_speech(const SttVadState * s) { return s && s->end_of_speech; }
static float rms_pcm16(const int16_t * pcm, int n) {
if (n <= 0) return 0.0f;
double acc = 0.0;
for (int i = 0; i < n; ++i) { double v = (double)pcm[i]; acc += v * v; }
return (float)std::sqrt(acc / (double)n);
}
bool stt_vad_push(SttVadState * s, const int16_t * pcm, int n_samples) {
if (!s || !pcm || n_samples <= 0) return false;
s->end_of_speech = false;
int offset = 0;
while (offset + s->frame_size <= n_samples) {
float r = rms_pcm16(pcm + offset, s->frame_size);
if (r >= (float)s->rms_threshold) {
s->speech_count += 1; s->silence_count = 0;
if (!s->in_speech && s->speech_count >= s->min_speech_frames) s->in_speech = true;
} else {
s->silence_count += 1;
if (s->in_speech && s->silence_count >= s->silence_end_frames) {
s->in_speech = false; s->end_of_speech = true; s->speech_count = 0;
} else if (!s->in_speech) {
s->speech_count = 0;
}
}
offset += s->frame_size;
}
return true;
}
// ============================================================================
// JSON minimaliste (mel_filters.json = flat float array, vocab.json = {"k":v,...})
// ============================================================================
// Parse mel_filters.json : "[1.234e-05, 0.0, ...]" (flat array de N floats).
static std::vector<float> parse_mel_filters_json(const char * path) {
std::ifstream f(path); if (!f) return {};
std::stringstream ss; ss << f.rdbuf();
std::string s = ss.str();
std::vector<float> out;
out.reserve(80 * 201);
size_t i = 0;
while (i < s.size()) {
// Trouve prochain digit ou signe
while (i < s.size() && !(std::isdigit((unsigned char)s[i]) || s[i] == '-' || s[i] == '+' || s[i] == '.')) ++i;
if (i >= s.size()) break;
size_t j = i;
// Avance tant que c'est un caractère numérique
while (j < s.size() && (std::isdigit((unsigned char)s[j]) || s[j] == '.' || s[j] == 'e' || s[j] == 'E' || s[j] == '+' || s[j] == '-')) {
// garde-fou : - et + uniquement au début ou après e
if ((s[j] == '-' || s[j] == '+') && j != i && s[j-1] != 'e' && s[j-1] != 'E') break;
++j;
}
if (j > i) {
try { out.push_back(std::stof(s.substr(i, j - i))); } catch (...) {}
}
i = j;
}
return out;
}
// Parse vocab.json (HF whisper format) : {"!":0,"\"":1,...,"<|endoftext|>":50257,...}
// On retourne id -> token (chaîne UTF-8 brute, qui contient parfois des unicodes BPE).
static bool parse_vocab_json(const char * path,
std::unordered_map<int, std::string> & id_to_token) {
std::ifstream f(path); if (!f) return false;
std::stringstream ss; ss << f.rdbuf();
std::string s = ss.str();
// On parcourt en cherchant des paires "key":number
size_t i = 0;
while (i < s.size()) {
// trouve "
while (i < s.size() && s[i] != '"') ++i;
if (i >= s.size()) break;
// Parse string key (handle \\, \", \uXXXX)
++i;
std::string key;
while (i < s.size() && s[i] != '"') {
if (s[i] == '\\' && i + 1 < s.size()) {
char c = s[i+1];
if (c == '"') { key.push_back('"'); i += 2; }
else if (c == '\\') { key.push_back('\\'); i += 2; }
else if (c == '/') { key.push_back('/'); i += 2; }
else if (c == 'n') { key.push_back('\n'); i += 2; }
else if (c == 't') { key.push_back('\t'); i += 2; }
else if (c == 'r') { key.push_back('\r'); i += 2; }
else if (c == 'b') { key.push_back('\b'); i += 2; }
else if (c == 'f') { key.push_back('\f'); i += 2; }
else if (c == 'u' && i + 5 < s.size()) {
// \uXXXX -> codepoint -> UTF-8
unsigned cp = 0;
for (int k = 0; k < 4; ++k) {
char h = s[i+2+k];
cp <<= 4;
if (h >= '0' && h <= '9') cp |= (unsigned)(h - '0');
else if (h >= 'a' && h <= 'f') cp |= (unsigned)(h - 'a' + 10);
else if (h >= 'A' && h <= 'F') cp |= (unsigned)(h - 'A' + 10);
}
// Encode UTF-8
if (cp < 0x80) key.push_back((char)cp);
else if (cp < 0x800) {
key.push_back((char)(0xC0 | (cp >> 6)));
key.push_back((char)(0x80 | (cp & 0x3F)));
} else {
key.push_back((char)(0xE0 | (cp >> 12)));
key.push_back((char)(0x80 | ((cp >> 6) & 0x3F)));
key.push_back((char)(0x80 | (cp & 0x3F)));
}
i += 6;
} else { key.push_back(s[i]); ++i; }
} else {
key.push_back(s[i]); ++i;
}
}
if (i >= s.size()) break;
++i; // skip closing "
// skip whitespace + ':' + whitespace
while (i < s.size() && (s[i] == ' ' || s[i] == ':' || s[i] == '\t' || s[i] == '\n')) ++i;
// Parse integer value
size_t j = i;
while (j < s.size() && std::isdigit((unsigned char)s[j])) ++j;
if (j > i) {
int id = std::stoi(s.substr(i, j - i));
id_to_token[id] = key;
}
i = j;
// skip jusqu'à la prochaine ","
while (i < s.size() && s[i] != ',' && s[i] != '}') ++i;
if (i < s.size() && s[i] == '}') break;
}
return !id_to_token.empty();
}
// BPE byte-level GPT-2/Whisper : map character code (in vocab token strings) -> raw byte
static std::unordered_map<int, int> build_bpe_unicode_to_byte() {
std::unordered_map<int, int> map;
for (int b = 33; b <= 126; ++b) map[b] = b;
for (int b = 161; b <= 172; ++b) map[b] = b;
for (int b = 174; b <= 255; ++b) map[b] = b;
int n = 0;
for (int b = 0; b <= 255; ++b) {
bool found = false;
for (auto & kv : map) { if (kv.second == b) { found = true; break; } }
if (!found) { map[256 + n] = b; ++n; }
}
return map;
}
// Parse UTF-8 codepoint at position i. Renvoie cp et avance i.
static unsigned utf8_next(const std::string & s, size_t & i) {
unsigned char c = (unsigned char)s[i++];
if (c < 0x80) return c;
if ((c >> 5) == 0x6) {
unsigned cp = (c & 0x1F);
cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
return cp;
}
if ((c >> 4) == 0xE) {
unsigned cp = (c & 0x0F);
cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
return cp;
}
if ((c >> 3) == 0x1E) {
unsigned cp = (c & 0x07);
cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
return cp;
}
return c;
}
// Décode une liste de token ids -> texte UTF-8 via BPE byte-level Whisper.
// Filtre les tokens spéciaux <|...|>.
static std::string decode_tokens(const std::vector<int> & ids,
const std::unordered_map<int, std::string> & vocab,
const std::unordered_map<int, int> & b2u_byte_map) {
std::string bytes;
bytes.reserve(ids.size() * 2);
for (int id : ids) {
auto it = vocab.find(id); if (it == vocab.end()) continue;
const std::string & w = it->second;
if (w.size() >= 4 && w[0] == '<' && w[1] == '|' &&
w[w.size()-2] == '|' && w[w.size()-1] == '>') continue;
size_t i = 0;
while (i < w.size()) {
unsigned cp = utf8_next(w, i);
auto bit = b2u_byte_map.find((int)cp);
if (bit != b2u_byte_map.end()) bytes.push_back((char)bit->second);
}
}
return bytes;
}
// ============================================================================
// SttEngine — ORT QNN HfWhisper KV-cache
// ============================================================================
namespace {
// Constantes Whisper (multilingue) — confirmées dans le code prod Kotlin.
constexpr int SOT = 50258;
constexpr int EOT = 50257;
constexpr int TRANSLATE_TOK = 50358;
constexpr int TRANSCRIBE_TOK = 50359;
constexpr int NOTIMESTAMPS = 50362;
constexpr int LANG_FR = 50265;
constexpr int LANG_EN = 50259;
constexpr int LANG_AUTO = -1; // sentinel: laisse le modèle prédire la langue au step 1
constexpr int VOCAB_SIZE = 51865;
constexpr int MEAN_DECODE_LEN = 200;
constexpr int HEAD_DIM = 64;
constexpr float MASK_NEG = -100.0f;
// Map ISO -> Whisper language token id (tous les codes Whisper officiels du modèle multilingue)
static int lang_to_token(const char * lang) {
if (!lang || !*lang) return LANG_FR;
std::string s = lang;
if (s == "auto") return LANG_AUTO;
// Mapping principal (les 99 langues Whisper utilisent les ids 50259..50357)
if (s == "en") return 50259;
if (s == "zh") return 50260;
if (s == "de") return 50261;
if (s == "es") return 50262;
if (s == "ru") return 50263;
if (s == "ko") return 50264;
if (s == "fr") return 50265;
if (s == "ja") return 50266;
if (s == "pt") return 50267;
if (s == "it") return 50274;
if (s == "nl") return 50271;
if (s == "ar") return 50272;
return LANG_FR; // default raisonnable pour Kazeia
}
double now_s() {
using clk = std::chrono::steady_clock;
return std::chrono::duration<double>(clk::now().time_since_epoch()).count();
}
} // namespace
struct SttEngine {
SttEngineLoadCfg cfg;
std::vector<float> mel_basis;
std::unique_ptr<Ort::Env> ort_env;
std::unique_ptr<Ort::SessionOptions> enc_opts, dec_opts;
std::unique_ptr<Ort::Session> enc_sess, dec_sess;
std::unique_ptr<Ort::MemoryInfo> mem_info;
std::unique_ptr<Ort::AllocatorWithDefaultOptions> allocator;
// Auto-detected dims (Whisper-Base: 6/8, Small: 12/12, Medium: 24/16)
int num_decoder_layers = 0;
int num_decoder_heads = 0;
// Cached I/O names (for the decoder loop tight allocation budget)
std::vector<std::string> enc_in_names_owned;
std::vector<const char *> enc_in_names;
std::vector<std::string> enc_out_names_owned;
std::vector<const char *> enc_out_names;
std::vector<std::string> dec_in_names_owned;
std::vector<const char *> dec_in_names;
std::vector<std::string> dec_out_names_owned;
std::vector<const char *> dec_out_names;
// Vocab + BPE
std::unordered_map<int, std::string> vocab;
std::unordered_map<int, int> bpe_unicode_to_byte;
bool loaded = false;
};
// Charge un fichier binaire ou JSON pour mel filters et retourne le tableau.
// Tente d'abord mel_filters.bin, sinon mel_filters.json.
static bool load_mel_basis_any(const std::string & model_dir,
const kazeia_mel::MelConfig & cfg,
std::vector<float> & out) {
if (kazeia_mel::load_mel_basis((model_dir + "/mel_filters.bin").c_str(), cfg, out)) return true;
auto v = parse_mel_filters_json((model_dir + "/mel_filters.json").c_str());
const size_t expected = (size_t)cfg.n_mels * (cfg.n_fft/2+1);
if (v.size() == expected) { out = std::move(v); return true; }
fprintf(stderr, "stt_engine: mel basis introuvable (ni .bin ni .json valide) à %s\n",
model_dir.c_str());
return false;
}
// Détecte num_decoder_layers (en comptant k_cache_cross_*) et num_decoder_heads (shape[0] du 1er).
static void detect_dims(SttEngine & e) {
Ort::AllocatorWithDefaultOptions a;
size_t n_out = e.enc_sess->GetOutputCount();
int n_layers = 0;
int n_heads = 0;
for (size_t i = 0; i < n_out; ++i) {
auto name_ptr = e.enc_sess->GetOutputNameAllocated(i, a);
std::string name = name_ptr.get();
if (name.rfind("k_cache_cross_", 0) == 0) {
++n_layers;
if (n_heads == 0) {
auto info = e.enc_sess->GetOutputTypeInfo(i);
auto tinfo = info.GetTensorTypeAndShapeInfo();
auto shape = tinfo.GetShape();
if (!shape.empty() && shape[0] > 0) n_heads = (int)shape[0];
}
}
}
// Défauts si détection ratée (Whisper-Base = 6/8, Small = 12/12)
if (n_layers == 0) n_layers = 12;
if (n_heads == 0) n_heads = 12;
e.num_decoder_layers = n_layers;
e.num_decoder_heads = n_heads;
}
// Cache les input/output names (les sessions ORT exigent const char**)
static void cache_io_names(Ort::Session & sess,
std::vector<std::string> & in_owned,
std::vector<const char *> & in_ptrs,
std::vector<std::string> & out_owned,
std::vector<const char *> & out_ptrs) {
Ort::AllocatorWithDefaultOptions a;
size_t ni = sess.GetInputCount(), no = sess.GetOutputCount();
in_owned.reserve(ni); in_ptrs.reserve(ni);
for (size_t i = 0; i < ni; ++i) {
auto p = sess.GetInputNameAllocated(i, a);
in_owned.emplace_back(p.get());
}
for (auto & s : in_owned) in_ptrs.push_back(s.c_str());
out_owned.reserve(no); out_ptrs.reserve(no);
for (size_t i = 0; i < no; ++i) {
auto p = sess.GetOutputNameAllocated(i, a);
out_owned.emplace_back(p.get());
}
for (auto & s : out_owned) out_ptrs.push_back(s.c_str());
}
SttEngine * stt_engine_load(const SttEngineLoadCfg & cfg) {
if (!cfg.model_dir) {
fprintf(stderr, "stt_engine_load: model_dir requis\n"); return nullptr;
}
auto * eng = new SttEngine();
eng->cfg = cfg;
std::string D = cfg.model_dir;
// 1) Mel basis (binaire optimal, JSON fallback)
auto mel_cfg = kazeia_mel::config_whisper();
if (!load_mel_basis_any(D, mel_cfg, eng->mel_basis)) {
delete eng; return nullptr;
}
// 2) ORT env
try {
eng->ort_env = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "kazeia_stt");
eng->mem_info = std::make_unique<Ort::MemoryInfo>(
Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault));
eng->allocator = std::make_unique<Ort::AllocatorWithDefaultOptions>();
} catch (const std::exception & e) {
fprintf(stderr, "stt_engine_load: Ort env FAIL: %s\n", e.what());
delete eng; return nullptr;
}
// 3) Sessions encoder + decoder. QNN EP si use_htp.
// Options HTP V79 réglables via env (debug bug "empty output sur certains audios") :
// KZSTT_QNN_BURST : "1" (def) = htp_performance_mode=burst, "0" = laisse default
// KZSTT_QNN_FP16 : "1" (def) = enable_htp_fp16_precision=1, "0" = laisse default
// KZSTT_QNN_ARCH : "1" (def) = htp_arch=79 explicite, "0" = auto-detect
// KZSTT_QNN_OPTALL : "1" (def) = SetGraphOptimizationLevel(ALL_OPT), "0" = default
auto envb = [](const char * k, bool dflt) {
const char * v = std::getenv(k); return v ? (atoi(v) != 0) : dflt;
};
const bool opt_burst = envb("KZSTT_QNN_BURST", true);
const bool opt_fp16 = envb("KZSTT_QNN_FP16", true);
const bool opt_arch = envb("KZSTT_QNN_ARCH", true);
const bool opt_optall = envb("KZSTT_QNN_OPTALL", true);
auto make_opts = [&](Ort::SessionOptions & opts) {
if (cfg.use_htp) {
std::vector<std::pair<std::string, std::string>> qnn_opts = {
{"backend_path", "libQnnHtp.so"},
{"profiling_level", "off"},
{"rpc_control_latency", "100"},
};
if (opt_burst) qnn_opts.push_back({"htp_performance_mode", "burst"});
if (opt_arch) qnn_opts.push_back({"htp_arch", "79"});
if (opt_fp16) qnn_opts.push_back({"enable_htp_fp16_precision", "1"});
std::unordered_map<std::string, std::string> qnn_map(qnn_opts.begin(), qnn_opts.end());
opts.AppendExecutionProvider("QNN", qnn_map);
}
opts.SetIntraOpNumThreads(cfg.n_threads);
if (opt_optall) opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
opts.DisableMemPattern();
};
fprintf(stderr, "stt_engine_load: QNN opts: burst=%d fp16=%d arch79=%d optall=%d\n",
opt_burst, opt_fp16, opt_arch, opt_optall);
std::string enc_path = D + "/HfWhisperEncoder.onnx";
std::string dec_path = D + "/HfWhisperDecoder.onnx";
try {
eng->enc_opts = std::make_unique<Ort::SessionOptions>();
make_opts(*eng->enc_opts);
double t0 = now_s();
eng->enc_sess = std::make_unique<Ort::Session>(*eng->ort_env, enc_path.c_str(), *eng->enc_opts);
fprintf(stderr, "stt_engine_load: encoder loaded in %.0f ms\n", (now_s() - t0) * 1000.0);
eng->dec_opts = std::make_unique<Ort::SessionOptions>();
make_opts(*eng->dec_opts);
double t1 = now_s();
eng->dec_sess = std::make_unique<Ort::Session>(*eng->ort_env, dec_path.c_str(), *eng->dec_opts);
fprintf(stderr, "stt_engine_load: decoder loaded in %.0f ms\n", (now_s() - t1) * 1000.0);
} catch (const std::exception & e) {
fprintf(stderr, "stt_engine_load: ORT session FAIL: %s\n", e.what());
delete eng; return nullptr;
}
// 4) Auto-détection des dims + cache des noms I/O
detect_dims(*eng);
cache_io_names(*eng->enc_sess, eng->enc_in_names_owned, eng->enc_in_names,
eng->enc_out_names_owned, eng->enc_out_names);
cache_io_names(*eng->dec_sess, eng->dec_in_names_owned, eng->dec_in_names,
eng->dec_out_names_owned, eng->dec_out_names);
fprintf(stderr, "stt_engine_load: %d layers / %d heads detected, enc I/O=%zu/%zu dec I/O=%zu/%zu\n",
eng->num_decoder_layers, eng->num_decoder_heads,
eng->enc_in_names.size(), eng->enc_out_names.size(),
eng->dec_in_names.size(), eng->dec_out_names.size());
// 5) Vocab.json + BPE byte map
if (!parse_vocab_json((D + "/vocab.json").c_str(), eng->vocab)) {
fprintf(stderr, "stt_engine_load: vocab.json absent ou parse FAIL à %s\n", D.c_str());
delete eng; return nullptr;
}
eng->bpe_unicode_to_byte = build_bpe_unicode_to_byte();
fprintf(stderr, "stt_engine_load: vocab %zu tokens, BPE map %zu entries\n",
eng->vocab.size(), eng->bpe_unicode_to_byte.size());
eng->loaded = true;
return eng;
}
// Création tensor FP16 owned (heap buffer kept alive jusqu'à la fin de la frame).
// Whisper-decoder loop produit ~200 fwd × (1+1+2L+1) tensors, ~13MB/step au pic. OK CPU.
static Ort::Value make_fp16_tensor(const SttEngine & e,
const std::vector<uint16_t> & data,
const std::vector<int64_t> & shape,
std::vector<std::vector<uint16_t>> & owner) {
owner.emplace_back(data);
auto & ref = owner.back();
return Ort::Value::CreateTensor(
*e.mem_info, ref.data(), ref.size() * sizeof(uint16_t),
shape.data(), shape.size(),
ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
}
static Ort::Value make_int32_tensor(const SttEngine & e, int v,
const std::vector<int64_t> & shape,
std::vector<std::vector<int32_t>> & owner) {
owner.push_back({v});
auto & ref = owner.back();
return Ort::Value::CreateTensor<int32_t>(
*e.mem_info, ref.data(), ref.size(),
shape.data(), shape.size());
}
// Trouve l'index argmax sur logits FP16 [VOCAB_SIZE].
static int argmax_fp16_logits(const uint16_t * data, int n) {
int best = 0; float best_v = -1e30f;
for (int i = 0; i < n; ++i) {
float v = fp16_to_fp32(data[i]);
if (v > best_v) { best_v = v; best = i; }
}
return best;
}
// Convertit FloatArray mel -> FP16 buffer (taille N_MELS * N_FRAMES = 80*3000 = 240000).
static std::vector<uint16_t> mel_to_fp16(const std::vector<float> & mel) {
std::vector<uint16_t> out(mel.size());
for (size_t i = 0; i < mel.size(); ++i) out[i] = fp32_to_fp16(mel[i]);
return out;
}
SttTranscribeResult stt_engine_transcribe(SttEngine * eng, const SttTranscribeCfg & cfg) {
SttTranscribeResult R{};
if (!eng || !eng->loaded || !cfg.pcm16 || cfg.n_samples <= 0) {
R.err = -1; return R;
}
if (cfg.sample_rate != 16000) {
fprintf(stderr, "stt_engine_transcribe: sample_rate=%d ; 16000 attendu\n", cfg.sample_rate);
R.err = -10; return R;
}
double t_start = now_s();
// 1) PCM -> wav f32
std::vector<float> wav((size_t)cfg.n_samples);
for (int i = 0; i < cfg.n_samples; ++i) wav[i] = (float)cfg.pcm16[i] / 32768.0f;
const double audio_s = (double)cfg.n_samples / cfg.sample_rate;
// 2) Mel via kazeia_mel
auto mel_cfg = kazeia_mel::config_whisper();
double tm0 = now_s();
int T = 0;
auto mel = kazeia_mel::compute(wav, eng->mel_basis, mel_cfg, T);
if (mel.empty()) { R.err = -3; return R; }
R.mel_ms = (int)((now_s() - tm0) * 1000.0);
// 3) Encoder — fp16 input -> N cross KV caches fp16
double te0 = now_s();
auto mel_fp16 = mel_to_fp16(mel);
std::vector<std::vector<uint16_t>> tensor_buffers;
tensor_buffers.reserve(eng->num_decoder_layers * 4 + 4);
std::vector<int64_t> mel_shape = {1, 80, 3000};
std::vector<Ort::Value> enc_inputs;
enc_inputs.emplace_back(make_fp16_tensor(*eng, mel_fp16, mel_shape, tensor_buffers));
// ORT expects: const char* const* input_names + Value* + size_t input_count.
std::vector<Ort::Value> enc_outputs;
try {
enc_outputs = eng->enc_sess->Run(
Ort::RunOptions{nullptr},
eng->enc_in_names.data(), enc_inputs.data(), enc_inputs.size(),
eng->enc_out_names.data(), eng->enc_out_names.size());
} catch (const std::exception & e) {
fprintf(stderr, "stt_engine_transcribe: encoder Run FAIL: %s\n", e.what());
R.err = -4; return R;
}
R.encoder_ms = (int)((now_s() - te0) * 1000.0);
// 4) Map outputs (encoder) by name -> get cross KV pointers + shapes
// Owned copies for the decoder loop (will be reused across all steps).
struct KvCache {
std::vector<uint16_t> data;
std::vector<int64_t> shape;
};
std::vector<KvCache> cross_k(eng->num_decoder_layers), cross_v(eng->num_decoder_layers);
for (size_t i = 0; i < eng->enc_out_names.size(); ++i) {
std::string nm = eng->enc_out_names_owned[i];
if (nm.rfind("k_cache_cross_", 0) == 0) {
int idx = std::stoi(nm.substr(14));
if (idx < eng->num_decoder_layers) {
auto & val = enc_outputs[i];
auto shape = val.GetTensorTypeAndShapeInfo().GetShape();
size_t n = 1; for (auto d : shape) n *= (size_t)d;
cross_k[idx].shape = shape;
cross_k[idx].data.resize(n);
std::memcpy(cross_k[idx].data.data(), val.GetTensorMutableData<uint16_t>(), n * 2);
}
} else if (nm.rfind("v_cache_cross_", 0) == 0) {
int idx = std::stoi(nm.substr(14));
if (idx < eng->num_decoder_layers) {
auto & val = enc_outputs[i];
auto shape = val.GetTensorTypeAndShapeInfo().GetShape();
size_t n = 1; for (auto d : shape) n *= (size_t)d;
cross_v[idx].shape = shape;
cross_v[idx].data.resize(n);
std::memcpy(cross_v[idx].data.data(), val.GetTensorMutableData<uint16_t>(), n * 2);
}
}
}
enc_outputs.clear();
// 5) Decoder loop autoregressif KV-cache (port optimisé zero-copy)
//
// ZERO-COPY : on alloue UNE FOIS tous les buffers, on crée UNE FOIS les Ort::Value
// qui pointent dedans, et on met juste à jour les valeurs in-place entre steps.
// Critical path : pas de std::vector::emplace_back(data) qui copie 12 MB/step.
double td0 = now_s();
const int kv_slots = MEAN_DECODE_LEN - 1; // 199
const size_t self_k_n = (size_t)eng->num_decoder_heads * HEAD_DIM * kv_slots;
const size_t self_v_n = (size_t)eng->num_decoder_heads * kv_slots * HEAD_DIM;
const std::vector<int64_t> self_k_shape = {eng->num_decoder_heads, 1, HEAD_DIM, kv_slots};
const std::vector<int64_t> self_v_shape = {eng->num_decoder_heads, 1, kv_slots, HEAD_DIM};
std::vector<KvCache> self_k(eng->num_decoder_layers), self_v(eng->num_decoder_layers);
for (int l = 0; l < eng->num_decoder_layers; ++l) {
self_k[l].shape = self_k_shape; self_k[l].data.assign(self_k_n, fp32_to_fp16(0.0f));
self_v[l].shape = self_v_shape; self_v[l].data.assign(self_v_n, fp32_to_fp16(0.0f));
}
// Buffers persistants pour input_ids (int32 [1,1]), attention_mask (fp16 [1,1,1,200]),
// position_ids (int32 [1])
std::vector<int32_t> input_ids_buf(1);
std::vector<uint16_t> mask_buf((size_t)MEAN_DECODE_LEN, fp32_to_fp16(MASK_NEG));
std::vector<int32_t> pos_ids_buf(1);
const std::vector<int64_t> input_ids_shape = {1, 1};
const std::vector<int64_t> mask_shape = {1, 1, 1, (int64_t)MEAN_DECODE_LEN};
const std::vector<int64_t> pos_ids_shape = {1};
// Pré-création des Ort::Value DECODER inputs UNE FOIS, alignés avec dec_in_names_owned.
// Chaque Value pointe sur son buffer permanent : ORT lit le buffer à chaque Run, on
// n'a qu'à updater le contenu in-place avant Run.
std::vector<Ort::Value> dec_inputs;
dec_inputs.reserve(eng->dec_in_names_owned.size());
for (size_t i = 0; i < eng->dec_in_names_owned.size(); ++i) {
const std::string & nm = eng->dec_in_names_owned[i];
if (nm == "input_ids") {
dec_inputs.emplace_back(Ort::Value::CreateTensor<int32_t>(
*eng->mem_info, input_ids_buf.data(), input_ids_buf.size(),
input_ids_shape.data(), input_ids_shape.size()));
} else if (nm == "attention_mask") {
dec_inputs.emplace_back(Ort::Value::CreateTensor(
*eng->mem_info, mask_buf.data(), mask_buf.size() * sizeof(uint16_t),
mask_shape.data(), mask_shape.size(),
ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
} else if (nm == "position_ids") {
dec_inputs.emplace_back(Ort::Value::CreateTensor<int32_t>(
*eng->mem_info, pos_ids_buf.data(), pos_ids_buf.size(),
pos_ids_shape.data(), pos_ids_shape.size()));
} else if (nm.rfind("k_cache_self_", 0) == 0 && nm.find("_in") != std::string::npos) {
int idx = std::stoi(nm.substr(13));
dec_inputs.emplace_back(Ort::Value::CreateTensor(
*eng->mem_info, self_k[idx].data.data(), self_k[idx].data.size() * sizeof(uint16_t),
self_k[idx].shape.data(), self_k[idx].shape.size(),
ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
} else if (nm.rfind("v_cache_self_", 0) == 0 && nm.find("_in") != std::string::npos) {
int idx = std::stoi(nm.substr(13));
dec_inputs.emplace_back(Ort::Value::CreateTensor(
*eng->mem_info, self_v[idx].data.data(), self_v[idx].data.size() * sizeof(uint16_t),
self_v[idx].shape.data(), self_v[idx].shape.size(),
ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
} else if (nm.rfind("k_cache_cross_", 0) == 0) {
int idx = std::stoi(nm.substr(14));
dec_inputs.emplace_back(Ort::Value::CreateTensor(
*eng->mem_info, cross_k[idx].data.data(), cross_k[idx].data.size() * sizeof(uint16_t),
cross_k[idx].shape.data(), cross_k[idx].shape.size(),
ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
} else if (nm.rfind("v_cache_cross_", 0) == 0) {
int idx = std::stoi(nm.substr(14));
dec_inputs.emplace_back(Ort::Value::CreateTensor(
*eng->mem_info, cross_v[idx].data.data(), cross_v[idx].data.size() * sizeof(uint16_t),
cross_v[idx].shape.data(), cross_v[idx].shape.size(),
ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
} else {
fprintf(stderr, "stt_engine_transcribe: input decoder inconnu : %s\n", nm.c_str());
R.err = -5; return R;
}
}
// Cache l'index du tensor "logits" dans dec_out_names_owned
int logits_out_idx = -1;
for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
if (eng->dec_out_names_owned[i] == "logits") { logits_out_idx = (int)i; break; }
}
if (logits_out_idx < 0) {
fprintf(stderr, "stt_engine_transcribe: pas de sortie 'logits' au decoder\n");
R.err = -7; return R;
}
// Cache la correspondance self_*_out -> idx layer (parsing 1 fois suffit)
std::vector<int> k_self_out_layer(eng->dec_out_names_owned.size(), -1);
std::vector<int> v_self_out_layer(eng->dec_out_names_owned.size(), -1);
for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
const std::string & nm = eng->dec_out_names_owned[i];
if (nm.rfind("k_cache_self_", 0) == 0 && nm.find("_out") != std::string::npos) {
k_self_out_layer[i] = std::stoi(nm.substr(13));
} else if (nm.rfind("v_cache_self_", 0) == 0 && nm.find("_out") != std::string::npos) {
v_self_out_layer[i] = std::stoi(nm.substr(13));
}
}
std::vector<int> generated;
generated.reserve(MEAN_DECODE_LEN);
// Forced decoder prompt : <SOT> <|lang|> <|transcribe|>
// On NE force PAS <|notimestamps|> : le decoder QAIRT préfère le mode timestamps
// (top-1 step 2 = 50363 = timestamp_begin). Forcer notimestamps fait EOT immédiat.
// En laissant le modèle décider timestamps vs notimestamps, on assure la cohérence
// avec le decoder QAIRT.
//
// Pourquoi le forced prompt : sans lui, sur audios borderline (voix aiguës,
// énergie faible), Whisper saute la prédiction de langue (top-1 = 50362
// notimestamps direct) et sort EOT step 1. Forcer lang+task résout 3/6 audios.
const int lang_tok = lang_to_token(cfg.language);
std::vector<int> forced_prompt = { SOT };
if (lang_tok != LANG_AUTO) {
forced_prompt.push_back(lang_tok);
forced_prompt.push_back(TRANSCRIBE_TOK);
}
// Si lang=auto, on garde juste SOT et on laisse Whisper prédire la langue+task.
int forced_idx = 0;
int current_token = forced_prompt[0];
int position_id = 0;
int real_vocab_size = -1;
for (int step = 0; step < MEAN_DECODE_LEN - 1; ++step) {
// Update mask : un seul fp16 à écrire par step (le slot R-to-L)
mask_buf[MEAN_DECODE_LEN - step - 1] = fp32_to_fp16(0.0f);
input_ids_buf[0] = current_token;
pos_ids_buf[0] = position_id;
// Run decoder — les Ort::Value pointent déjà sur les bons buffers (à jour in-place)
std::vector<Ort::Value> dec_outputs;
try {
dec_outputs = eng->dec_sess->Run(
Ort::RunOptions{nullptr},
eng->dec_in_names.data(), dec_inputs.data(), dec_inputs.size(),
eng->dec_out_names.data(), eng->dec_out_names.size());
} catch (const std::exception & e) {
fprintf(stderr, "stt_engine_transcribe: decoder Run step=%d FAIL: %s\n", step, e.what());
R.err = -6; return R;
}
// argmax logits avec shape réelle (détectée au step 0)
const auto logits_shape = dec_outputs[logits_out_idx].GetTensorTypeAndShapeInfo().GetShape();
if (real_vocab_size < 0) {
size_t n = 1; for (auto d : logits_shape) n *= (size_t)d;
real_vocab_size = (int)n;
}
const uint16_t * logits_data = dec_outputs[logits_out_idx].GetTensorMutableData<uint16_t>();
int token = argmax_fp16_logits(logits_data, real_vocab_size);
// Forced prompt : tant qu'on est dans le prompt, on impose le token suivant
// (les self_k/v se construisent normalement, mais on ignore le logit de sortie).
if (forced_idx + 1 < (int)forced_prompt.size()) {
forced_idx += 1;
token = forced_prompt[forced_idx];
}
if (step < 8 && std::getenv("KZSTT_DEBUG_STEPS")) {
std::vector<std::pair<float,int>> scores(real_vocab_size);
for (int i = 0; i < real_vocab_size; ++i)
scores[i] = {fp16_to_fp32(logits_data[i]), i};
std::partial_sort(scores.begin(), scores.begin() + 5, scores.end(),
[](auto & a, auto & b){ return a.first > b.first; });
fprintf(stderr, " step=%d current_token=%d pos=%d top5:", step, current_token, position_id);
for (int i = 0; i < 5; ++i) {
fprintf(stderr, " [%d:%+.3f]", scores[i].second, scores[i].first);
}
fprintf(stderr, " -> picked=%d\n", token);
}
if (cfg.force_transcribe && token == TRANSLATE_TOK) token = TRANSCRIBE_TOK;
// Update self KV from outputs : memcpy direct dans nos buffers persistants
// (les Ort::Value dec_inputs pointent DÉJÀ dessus, prêts pour le step suivant)
for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
int k_idx = k_self_out_layer[i];
int v_idx = v_self_out_layer[i];
if (k_idx >= 0) {
auto & val = dec_outputs[i];
std::memcpy(self_k[k_idx].data.data(),
val.GetTensorMutableData<uint16_t>(),
self_k[k_idx].data.size() * sizeof(uint16_t));
} else if (v_idx >= 0) {
auto & val = dec_outputs[i];
std::memcpy(self_v[v_idx].data.data(),
val.GetTensorMutableData<uint16_t>(),
self_v[v_idx].data.size() * sizeof(uint16_t));
}
}
dec_outputs.clear();
if (token == EOT) break;
if (token < 50257) generated.push_back(token);
current_token = token;
position_id += 1;
}
R.decoder_ms = (int)((now_s() - td0) * 1000.0);
// 6) Token decode -> texte
R.text = decode_tokens(generated, eng->vocab, eng->bpe_unicode_to_byte);
R.n_tokens = (int)generated.size();
R.detected_language = cfg.language;
R.total_ms = (int)((now_s() - t_start) * 1000.0);
R.rtf = (audio_s > 0) ? (float)((double)R.total_ms / (audio_s * 1000.0)) : 0.0f;
return R;
}
void stt_engine_free(SttEngine * eng) {
if (!eng) return;
eng->enc_sess.reset(); eng->dec_sess.reset();
eng->enc_opts.reset(); eng->dec_opts.reset();
eng->ort_env.reset();
delete eng;
}