760 lines
32 KiB
C++
760 lines
32 KiB
C++
// 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 VOCAB_SIZE = 51865;
|
||
constexpr int MEAN_DECODE_LEN = 200;
|
||
constexpr int HEAD_DIM = 64;
|
||
constexpr float MASK_NEG = -100.0f;
|
||
|
||
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.
|
||
auto make_opts = [&](Ort::SessionOptions & opts) {
|
||
if (cfg.use_htp) {
|
||
// QNN ExecutionProvider via AppendExecutionProvider("QNN", {...})
|
||
std::vector<const char *> keys = { "backend_path" };
|
||
std::vector<const char *> vals = { "libQnnHtp.so" };
|
||
opts.AppendExecutionProvider("QNN",
|
||
{{"backend_path", "libQnnHtp.so"}});
|
||
}
|
||
opts.SetIntraOpNumThreads(cfg.n_threads);
|
||
};
|
||
|
||
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 exact decodeHfKvCache.kt)
|
||
double td0 = now_s();
|
||
const int kv_slots = MEAN_DECODE_LEN - 1; // 199
|
||
// Self KV layout : k [H, 1, head_dim, slots] ; v [H, 1, slots, head_dim].
|
||
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));
|
||
}
|
||
|
||
std::vector<float> mask_f32((size_t)MEAN_DECODE_LEN, MASK_NEG);
|
||
std::vector<uint16_t> mask_fp16((size_t)MEAN_DECODE_LEN);
|
||
|
||
std::vector<int> generated;
|
||
generated.reserve(MEAN_DECODE_LEN);
|
||
|
||
int current_token = SOT;
|
||
int position_id = 0;
|
||
|
||
// Préalloue les noms decoder (input_ids, attention_mask, k/v_cache_self_*_in,
|
||
// k/v_cache_cross_*, position_ids). On les chercher dans dec_in_names_owned.
|
||
// Cache des indices pour Run() : on passe les Value dans l'ordre des noms cachés.
|
||
|
||
for (int step = 0; step < MEAN_DECODE_LEN - 1; ++step) {
|
||
// Mask R-to-L
|
||
mask_f32[MEAN_DECODE_LEN - step - 1] = 0.0f;
|
||
for (int i = 0; i < MEAN_DECODE_LEN; ++i) mask_fp16[i] = fp32_to_fp16(mask_f32[i]);
|
||
|
||
// Build Values aligned avec dec_in_names_owned[i] ordering
|
||
tensor_buffers.clear();
|
||
std::vector<std::vector<int32_t>> int_buffers;
|
||
|
||
std::vector<Ort::Value> dec_inputs;
|
||
dec_inputs.reserve(eng->dec_in_names.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(make_int32_tensor(*eng, current_token, {1, 1}, int_buffers));
|
||
} else if (nm == "attention_mask") {
|
||
dec_inputs.emplace_back(make_fp16_tensor(*eng, mask_fp16,
|
||
{1, 1, 1, (int64_t)MEAN_DECODE_LEN}, tensor_buffers));
|
||
} else if (nm == "position_ids") {
|
||
dec_inputs.emplace_back(make_int32_tensor(*eng, position_id, {1}, int_buffers));
|
||
} 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(make_fp16_tensor(*eng, self_k[idx].data, self_k[idx].shape, tensor_buffers));
|
||
} 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(make_fp16_tensor(*eng, self_v[idx].data, self_v[idx].shape, tensor_buffers));
|
||
} else if (nm.rfind("k_cache_cross_", 0) == 0) {
|
||
int idx = std::stoi(nm.substr(14));
|
||
dec_inputs.emplace_back(make_fp16_tensor(*eng, cross_k[idx].data, cross_k[idx].shape, tensor_buffers));
|
||
} else if (nm.rfind("v_cache_cross_", 0) == 0) {
|
||
int idx = std::stoi(nm.substr(14));
|
||
dec_inputs.emplace_back(make_fp16_tensor(*eng, cross_v[idx].data, cross_v[idx].shape, tensor_buffers));
|
||
} else {
|
||
fprintf(stderr, "stt_engine_transcribe: input decoder inconnu : %s\n", nm.c_str());
|
||
R.err = -5; return R;
|
||
}
|
||
}
|
||
|
||
// Run decoder
|
||
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;
|
||
}
|
||
|
||
// Parse outputs : logits + updated self KV
|
||
int logits_idx = -1;
|
||
for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
|
||
if (eng->dec_out_names_owned[i] == "logits") { logits_idx = (int)i; break; }
|
||
}
|
||
if (logits_idx < 0) {
|
||
fprintf(stderr, "stt_engine_transcribe: pas de sortie 'logits' au decoder\n");
|
||
R.err = -7; return R;
|
||
}
|
||
const uint16_t * logits_data = dec_outputs[logits_idx].GetTensorMutableData<uint16_t>();
|
||
int token = argmax_fp16_logits(logits_data, VOCAB_SIZE);
|
||
|
||
// Override translate -> transcribe
|
||
if (cfg.force_transcribe && token == TRANSLATE_TOK) token = TRANSCRIBE_TOK;
|
||
|
||
// Update self KV from outputs
|
||
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) {
|
||
int idx = std::stoi(nm.substr(13));
|
||
auto & val = dec_outputs[i];
|
||
size_t n = 1; for (auto d : val.GetTensorTypeAndShapeInfo().GetShape()) n *= (size_t)d;
|
||
if (n == self_k[idx].data.size()) {
|
||
std::memcpy(self_k[idx].data.data(), val.GetTensorMutableData<uint16_t>(), n * 2);
|
||
}
|
||
} else if (nm.rfind("v_cache_self_", 0) == 0 && nm.find("_out") != std::string::npos) {
|
||
int idx = std::stoi(nm.substr(13));
|
||
auto & val = dec_outputs[i];
|
||
size_t n = 1; for (auto d : val.GetTensorTypeAndShapeInfo().GetShape()) n *= (size_t)d;
|
||
if (n == self_v[idx].data.size()) {
|
||
std::memcpy(self_v[idx].data.data(), val.GetTensorMutableData<uint16_t>(), n * 2);
|
||
}
|
||
}
|
||
}
|
||
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;
|
||
}
|