#include "bin_load.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "transcribe-arch.h"
#include "transcribe-bin-loader.h"
#include "transcribe-load-common.h"
#include "transcribe-log.h"
#include "transcribe-tokenizer.h"
#include "weights.h"
#include "whisper.h"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
namespace transcribe::whisper {
extern const Arch arch;
namespace {
constexpr const char * kTag = "whisper-bin";
std::string detect_variant(const transcribe::bin_loader::WhisperBinModel & bm) {
const auto & h = bm.hp;
std::string base;
switch (h.n_audio_layer) {
case 4:
base = "whisper-tiny";
break;
case 6:
base = "whisper-base";
break;
case 12:
base = "whisper-small";
break;
case 24:
base = "whisper-medium";
break;
case 32:
if (h.n_text_layer == 4) {
base = "whisper-large-v3-turbo";
} else if (h.n_mels == 128) {
base = "whisper-large-v3";
} else {
base = "whisper-large";
}
break;
default:
base = "whisper";
break;
}
if (!bm.is_multilingual) {
base += ".en";
}
return base;
}
const char * const k_whisper_lang_codes[] = {
"en", "zh", "de", "es", "ru", "ko", "fr", "ja", "pt", "tr", "pl", "ca", "nl", "ar", "sv", "it", "id",
"hi", "fi", "vi", "he", "uk", "el", "ms", "cs", "ro", "da", "hu", "ta", "no", "th", "ur", "hr", "bg",
"lt", "la", "mi", "ml", "cy", "sk", "te", "fa", "lv", "bn", "sr", "az", "sl", "kn", "et", "mk", "br",
"eu", "is", "hy", "ne", "mn", "bs", "kk", "sq", "sw", "gl", "mr", "pa", "si", "km", "sn", "yo", "so",
"af", "oc", "ka", "be", "tg", "sd", "gu", "am", "yi", "lo", "uz", "fo", "ht", "ps", "tk", "nn", "mt",
"sa", "lb", "my", "bo", "tl", "mg", "as", "tt", "haw", "ln", "ha", "ba", "jw", "su", "yue",
};
constexpr int k_whisper_n_lang_codes = static_cast<int>(sizeof(k_whisper_lang_codes) / sizeof(k_whisper_lang_codes[0]));
struct WhisperSpecials {
int eot = 50256;
int sot = 50257;
int translate = 50357;
int transcribe = 50358;
int solm = 50359;
int prev = 50360;
int nosp = 50361;
int notimestamps = 50362; int beg = 50363; };
WhisperSpecials compute_specials(int n_vocab) {
WhisperSpecials s;
const bool is_multilingual = n_vocab >= 51865;
if (is_multilingual) {
s.eot += 1;
s.sot += 1;
const int num_languages = n_vocab - 51765 - 1;
const int dt = num_languages - 98;
s.translate += dt;
s.transcribe += dt;
s.solm += dt;
s.prev += dt;
s.nosp += dt;
s.notimestamps += dt;
s.beg += dt;
}
return s;
}
const int32_t k_whisper_suppress_tokens[] = {
1, 2, 7, 8, 9, 10, 14, 25, 26, 27, 28, 29, 31, 58, 59,
60, 61, 62, 63, 90, 91, 92, 93, 359, 503, 522, 542, 873, 893, 902,
918, 922, 931, 1350, 1853, 1982, 2460, 2627, 3246, 3253, 3268, 3536, 3846, 3961, 4183,
4667, 6585, 6647, 7273, 9061, 9383, 10428, 10929, 11938, 12033, 12331, 12562, 13793, 14157, 14635,
15265, 15618, 16553, 16604, 18362, 18956, 20075, 21675, 22520, 26130, 26161, 26435, 28279, 29464, 31650,
32302, 32470, 36865, 42863, 47425, 49870, 50254, 50258, 50358, 50359, 50360, 50361, 50362,
};
constexpr int k_whisper_n_suppress =
static_cast<int>(sizeof(k_whisper_suppress_tokens) / sizeof(k_whisper_suppress_tokens[0]));
struct StaticRename {
const char * legacy;
const char * canonical;
bool collapse_to_1d; };
const StaticRename k_static_renames[] = {
{ "encoder.positional_embedding", "enc.pos_emb.weight", false },
{ "encoder.conv1.weight", "enc.conv.0.weight", false },
{ "encoder.conv1.bias", "enc.conv.0.bias", true },
{ "encoder.conv2.weight", "enc.conv.1.weight", false },
{ "encoder.conv2.bias", "enc.conv.1.bias", true },
{ "encoder.ln_post.weight", "enc.final_norm.weight", false },
{ "encoder.ln_post.bias", "enc.final_norm.bias", false },
{ "decoder.positional_embedding", "dec.pos_emb.weight", false },
{ "decoder.token_embedding.weight", "dec.token_embd.weight", false },
{ "decoder.ln.weight", "dec.final_norm.weight", false },
{ "decoder.ln.bias", "dec.final_norm.bias", false },
};
struct LayeredRename {
const char * legacy_fmt;
const char * canonical_fmt;
};
const LayeredRename k_enc_block_renames[] = {
{ "encoder.blocks.%d.attn_ln.weight", "enc.blocks.%d.norm_attn.weight" },
{ "encoder.blocks.%d.attn_ln.bias", "enc.blocks.%d.norm_attn.bias" },
{ "encoder.blocks.%d.attn.query.weight", "enc.blocks.%d.attn.q.weight" },
{ "encoder.blocks.%d.attn.query.bias", "enc.blocks.%d.attn.q.bias" },
{ "encoder.blocks.%d.attn.key.weight", "enc.blocks.%d.attn.k.weight" }, { "encoder.blocks.%d.attn.value.weight", "enc.blocks.%d.attn.v.weight" },
{ "encoder.blocks.%d.attn.value.bias", "enc.blocks.%d.attn.v.bias" },
{ "encoder.blocks.%d.attn.out.weight", "enc.blocks.%d.attn.out.weight" },
{ "encoder.blocks.%d.attn.out.bias", "enc.blocks.%d.attn.out.bias" },
{ "encoder.blocks.%d.mlp_ln.weight", "enc.blocks.%d.norm_ffn.weight" },
{ "encoder.blocks.%d.mlp_ln.bias", "enc.blocks.%d.norm_ffn.bias" },
{ "encoder.blocks.%d.mlp.0.weight", "enc.blocks.%d.ffn.fc1.weight" },
{ "encoder.blocks.%d.mlp.0.bias", "enc.blocks.%d.ffn.fc1.bias" },
{ "encoder.blocks.%d.mlp.2.weight", "enc.blocks.%d.ffn.fc2.weight" },
{ "encoder.blocks.%d.mlp.2.bias", "enc.blocks.%d.ffn.fc2.bias" },
};
const LayeredRename k_dec_block_renames[] = {
{ "decoder.blocks.%d.attn_ln.weight", "dec.blocks.%d.norm_self.weight" },
{ "decoder.blocks.%d.attn_ln.bias", "dec.blocks.%d.norm_self.bias" },
{ "decoder.blocks.%d.attn.query.weight", "dec.blocks.%d.self_attn.q.weight" },
{ "decoder.blocks.%d.attn.query.bias", "dec.blocks.%d.self_attn.q.bias" },
{ "decoder.blocks.%d.attn.key.weight", "dec.blocks.%d.self_attn.k.weight" },
{ "decoder.blocks.%d.attn.value.weight", "dec.blocks.%d.self_attn.v.weight" },
{ "decoder.blocks.%d.attn.value.bias", "dec.blocks.%d.self_attn.v.bias" },
{ "decoder.blocks.%d.attn.out.weight", "dec.blocks.%d.self_attn.out.weight" },
{ "decoder.blocks.%d.attn.out.bias", "dec.blocks.%d.self_attn.out.bias" },
{ "decoder.blocks.%d.cross_attn_ln.weight", "dec.blocks.%d.norm_cross.weight" },
{ "decoder.blocks.%d.cross_attn_ln.bias", "dec.blocks.%d.norm_cross.bias" },
{ "decoder.blocks.%d.cross_attn.query.weight", "dec.blocks.%d.cross_attn.q.weight" },
{ "decoder.blocks.%d.cross_attn.query.bias", "dec.blocks.%d.cross_attn.q.bias" },
{ "decoder.blocks.%d.cross_attn.key.weight", "dec.blocks.%d.cross_attn.k.weight" },
{ "decoder.blocks.%d.cross_attn.value.weight", "dec.blocks.%d.cross_attn.v.weight" },
{ "decoder.blocks.%d.cross_attn.value.bias", "dec.blocks.%d.cross_attn.v.bias" },
{ "decoder.blocks.%d.cross_attn.out.weight", "dec.blocks.%d.cross_attn.out.weight" },
{ "decoder.blocks.%d.cross_attn.out.bias", "dec.blocks.%d.cross_attn.out.bias" },
{ "decoder.blocks.%d.mlp_ln.weight", "dec.blocks.%d.norm_ffn.weight" },
{ "decoder.blocks.%d.mlp_ln.bias", "dec.blocks.%d.norm_ffn.bias" },
{ "decoder.blocks.%d.mlp.0.weight", "dec.blocks.%d.ffn.fc1.weight" },
{ "decoder.blocks.%d.mlp.0.bias", "dec.blocks.%d.ffn.fc1.bias" },
{ "decoder.blocks.%d.mlp.2.weight", "dec.blocks.%d.ffn.fc2.weight" },
{ "decoder.blocks.%d.mlp.2.bias", "dec.blocks.%d.ffn.fc2.bias" },
};
std::string fmt1(const char * fmt, int idx) {
char buf[160];
std::snprintf(buf, sizeof(buf), fmt, idx);
return std::string(buf);
}
std::vector<float> hann_periodic(int n_fft) {
std::vector<float> w(static_cast<size_t>(n_fft));
const double two_pi_over_n = 2.0 * M_PI / static_cast<double>(n_fft);
for (int k = 0; k < n_fft; ++k) {
w[static_cast<size_t>(k)] = static_cast<float>(0.5 * (1.0 - std::cos(two_pi_over_n * k)));
}
return w;
}
void fill_whisper_hparams(const transcribe::bin_loader::WhisperBinModel & bm, WhisperHParams & hp) {
hp = WhisperHParams{};
hp.enc_n_layers = bm.hp.n_audio_layer;
hp.enc_d_model = bm.hp.n_audio_state;
hp.enc_n_heads = bm.hp.n_audio_head;
hp.enc_ffn_dim = 4 * bm.hp.n_audio_state;
hp.enc_num_mel_bins = bm.hp.n_mels;
hp.enc_max_source_positions = bm.hp.n_audio_ctx;
hp.enc_activation = "gelu";
hp.dec_n_layers = bm.hp.n_text_layer;
hp.dec_d_model = bm.hp.n_text_state;
hp.dec_n_heads = bm.hp.n_text_head;
hp.dec_ffn_dim = 4 * bm.hp.n_text_state;
hp.dec_max_target_positions = bm.hp.n_text_ctx;
hp.dec_vocab_size = bm.hp.n_vocab;
hp.dec_activation = "gelu";
hp.dec_tie_word_embeddings = true;
hp.dec_scale_embedding = false;
const WhisperSpecials sp = compute_specials(bm.hp.n_vocab);
hp.decoder_start_token_id = sp.sot;
hp.sot_token_id = sp.sot;
hp.no_timestamps_token_id = sp.notimestamps;
hp.transcribe_token_id = bm.is_multilingual ? sp.transcribe : -1;
hp.translate_token_id = bm.is_multilingual ? sp.translate : -1;
hp.prev_sot_token_id = sp.prev;
hp.suppress_tokens.assign(k_whisper_suppress_tokens, k_whisper_suppress_tokens + k_whisper_n_suppress);
hp.begin_suppress_tokens.assign({ 220, sp.eot });
hp.fe_type = "mel";
hp.fe_num_mels = bm.hp.n_mels;
hp.fe_sample_rate = 16000;
hp.fe_n_fft = 400;
hp.fe_win_length = 400;
hp.fe_hop_length = 160;
hp.fe_window = "hann_periodic";
hp.fe_normalize = "whisper_logmel";
hp.fe_dither = 0.0f;
hp.fe_pre_emphasis = 0.0f;
hp.fe_f_min = 0.0f;
hp.fe_f_max = 8000.0f;
hp.fe_pad_mode = "reflect";
hp.fe_center = true;
hp.fe_mel_norm = "slaney";
hp.fe_chunk_length = 30;
hp.fe_n_samples = 480000;
hp.fe_nb_max_frames = 3000;
hp.cap_lang_detect = bm.is_multilingual;
hp.cap_translate = bm.is_multilingual;
hp.cap_timestamps = true;
}
ggml_tensor * create_canonical_tensor(ggml_context * ctx_meta,
const std::string & canonical_name,
const transcribe::bin_loader::BinTensorEntry & src,
bool collapse_to_1d) {
int n_dims = src.n_dims;
int64_t ne[GGML_MAX_DIMS];
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
ne[i] = 1;
}
for (int i = 0; i < n_dims; ++i) {
ne[i] = src.ne[i];
}
if (collapse_to_1d) {
int64_t total = 1;
for (int i = 0; i < n_dims; ++i) {
total *= ne[i];
}
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
ne[i] = 1;
}
ne[0] = total;
n_dims = 1;
}
ggml_tensor * t = ggml_new_tensor(ctx_meta, src.type, n_dims, ne);
if (t == nullptr) {
log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "%s: ggml_new_tensor failed for \"%s\"", kTag, canonical_name.c_str());
return nullptr;
}
ggml_set_name(t, canonical_name.c_str());
return t;
}
transcribe_status resolve_tensors(const transcribe::bin_loader::WhisperBinModel & bm,
const WhisperHParams & hp,
ggml_context * ctx_meta,
std::vector<transcribe::bin_loader::BinStreamSlot> & out_slots) {
std::unordered_map<std::string, const transcribe::bin_loader::BinTensorEntry *> by_name;
by_name.reserve(bm.tensors.size() * 2);
for (const auto & e : bm.tensors) {
if (!by_name.emplace(e.name, &e).second) {
log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "%s: duplicate tensor \"%s\" in .bin", kTag, e.name.c_str());
return TRANSCRIBE_ERR_GGUF;
}
}
std::unordered_set<std::string> consumed;
consumed.reserve(by_name.size() * 2);
{
const int64_t n_fft_half = static_cast<int64_t>(hp.fe_n_fft / 2 + 1);
const int64_t fb_ne[2] = { n_fft_half, static_cast<int64_t>(hp.fe_num_mels) };
ggml_tensor * fb = ggml_new_tensor(ctx_meta, GGML_TYPE_F32, 2, fb_ne);
if (fb == nullptr) {
return TRANSCRIBE_ERR_GGUF;
}
ggml_set_name(fb, "frontend.mel_filterbank");
const int64_t win_ne[1] = { static_cast<int64_t>(hp.fe_n_fft) };
ggml_tensor * win = ggml_new_tensor(ctx_meta, GGML_TYPE_F32, 1, win_ne);
if (win == nullptr) {
return TRANSCRIBE_ERR_GGUF;
}
ggml_set_name(win, "frontend.window");
}
auto add_slot = [&](const StaticRename & rule) -> transcribe_status {
auto it = by_name.find(rule.legacy);
if (it == by_name.end()) {
log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
"%s: missing tensor \"%s\" (canonical \"%s\") "
"— .bin is incomplete or not a whisper model",
kTag, rule.legacy, rule.canonical);
return TRANSCRIBE_ERR_GGUF;
}
ggml_tensor * t = create_canonical_tensor(ctx_meta, rule.canonical, *it->second, rule.collapse_to_1d);
if (t == nullptr) {
return TRANSCRIBE_ERR_GGUF;
}
out_slots.push_back({ t, it->second->offset, it->second->nbytes });
consumed.insert(rule.legacy);
return TRANSCRIBE_OK;
};
auto add_slot_layered = [&](const LayeredRename & rule, int idx) -> transcribe_status {
StaticRename s;
const std::string legacy_s = fmt1(rule.legacy_fmt, idx);
const std::string canonical_s = fmt1(rule.canonical_fmt, idx);
s.legacy = legacy_s.c_str();
s.canonical = canonical_s.c_str();
s.collapse_to_1d = false;
return add_slot(s);
};
for (const auto & r : k_static_renames) {
if (auto st = add_slot(r); st != TRANSCRIBE_OK) {
return st;
}
}
for (int i = 0; i < hp.enc_n_layers; ++i) {
for (const auto & r : k_enc_block_renames) {
if (auto st = add_slot_layered(r, i); st != TRANSCRIBE_OK) {
return st;
}
}
}
for (int i = 0; i < hp.dec_n_layers; ++i) {
for (const auto & r : k_dec_block_renames) {
if (auto st = add_slot_layered(r, i); st != TRANSCRIBE_OK) {
return st;
}
}
}
if (consumed.size() != bm.tensors.size()) {
for (const auto & e : bm.tensors) {
if (consumed.find(e.name) == consumed.end()) {
log_msg(TRANSCRIBE_LOG_LEVEL_INFO,
"%s: unconsumed .bin tensor \"%s\" — "
"this loader does not yet support it",
kTag, e.name.c_str());
}
}
return TRANSCRIBE_ERR_GGUF;
}
return TRANSCRIBE_OK;
}
void apply_caps_and_languages(WhisperModel & m, const transcribe::bin_loader::WhisperBinModel & bm) {
apply_family_invariants(m);
if (!bm.is_multilingual) {
m.caps.supports_language_detect = false;
m.caps.supports_translate = false;
}
m.lang_codes.clear();
if (bm.is_multilingual) {
const int n = std::min(bm.num_languages, k_whisper_n_lang_codes);
m.lang_codes.reserve(static_cast<size_t>(n));
for (int i = 0; i < n; ++i) {
m.lang_codes.emplace_back(k_whisper_lang_codes[i]);
}
} else {
m.lang_codes.emplace_back("en");
}
m.lang_token_ids.clear();
if (bm.is_multilingual) {
const WhisperSpecials sp = compute_specials(bm.hp.n_vocab);
m.lang_token_ids.reserve(m.lang_codes.size());
for (size_t i = 0; i < m.lang_codes.size(); ++i) {
m.lang_token_ids.push_back(static_cast<int32_t>(sp.sot + 1 + static_cast<int>(i)));
}
}
m.set_languages(m.lang_codes);
}
transcribe_status install_tokenizer(WhisperModel & m, const transcribe::bin_loader::WhisperBinModel & bm) {
transcribe::Tokenizer::DecodeOnlySpecials sp;
const WhisperSpecials s = compute_specials(bm.hp.n_vocab);
sp.eos_id = s.eot;
sp.bos_id = s.eot;
if (auto st = m.tok.load_decode_only_raw_bytes(bm.vocab_tokens, sp); st != TRANSCRIBE_OK) {
return st;
}
m.tok.add_special_piece("<|endoftext|>", s.eot);
m.tok.add_special_piece("<|startoftranscript|>", s.sot);
m.tok.add_special_piece("<|startofprev|>", s.prev);
m.tok.add_special_piece("<|nospeech|>", s.nosp);
m.tok.add_special_piece("<|notimestamps|>", s.notimestamps);
if (bm.is_multilingual) {
m.tok.add_special_piece("<|translate|>", s.translate);
m.tok.add_special_piece("<|transcribe|>", s.transcribe);
const int n_lang = std::min(bm.num_languages, k_whisper_n_lang_codes);
for (int i = 0; i < n_lang; ++i) {
std::string lit = std::string("<|") + k_whisper_lang_codes[i] + "|>";
m.tok.add_special_piece(lit, s.sot + 1 + i);
}
}
char tsbuf[16];
for (int i = 0; i <= 1500; ++i) {
std::snprintf(tsbuf, sizeof(tsbuf), "<|%.2f|>", static_cast<double>(i) * 0.02);
m.tok.add_special_piece(tsbuf, s.beg + i);
}
return TRANSCRIBE_OK;
}
}
transcribe_status load_from_bin(const char * path,
const struct transcribe_model_load_params * params,
struct transcribe_model ** out_model) {
if (out_model == nullptr || path == nullptr || params == nullptr) {
return TRANSCRIBE_ERR_INVALID_ARG;
}
*out_model = nullptr;
const int64_t t_load_start = ggml_time_us();
transcribe::bin_loader::WhisperBinModel bm;
if (auto st = transcribe::bin_loader::parse_whisper_bin(path, bm); st != TRANSCRIBE_OK) {
return st;
}
auto m = std::make_unique<WhisperModel>();
m->arch = &arch;
m->t_load_us = 0;
m->variant = detect_variant(bm);
m->backend.clear();
fill_whisper_hparams(bm, m->hparams);
if (auto st = install_tokenizer(*m, bm); st != TRANSCRIBE_OK) {
return st;
}
apply_caps_and_languages(*m, bm);
const size_t n_tensors_est =
static_cast<size_t>(15 + 15 * m->hparams.enc_n_layers + 24 * m->hparams.dec_n_layers + 32);
ggml_init_params init_params{};
init_params.mem_size = (n_tensors_est + 64) * ggml_tensor_overhead();
init_params.mem_buffer = nullptr;
init_params.no_alloc = true;
m->ctx_meta = ggml_init(init_params);
if (m->ctx_meta == nullptr) {
return TRANSCRIBE_ERR_GGUF;
}
std::vector<transcribe::bin_loader::BinStreamSlot> stream_slots;
stream_slots.reserve(bm.tensors.size());
if (auto st = resolve_tensors(bm, m->hparams, m->ctx_meta, stream_slots); st != TRANSCRIBE_OK) {
return st;
}
if (auto st = build_whisper_weights(m->ctx_meta, m->hparams, m->weights); st != TRANSCRIBE_OK) {
return st;
}
if (auto st = transcribe::load_common::init_backends(params->backend, params->device, "whisper", m->plan);
st != TRANSCRIBE_OK) {
return st;
}
m->backend = ggml_backend_name(m->plan.primary);
m->primary_backend = m->plan.primary;
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(m->ctx_meta, m->plan.primary);
if (buf == nullptr) {
log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "%s: ggml_backend_alloc_ctx_tensors failed", kTag);
return TRANSCRIBE_ERR_GGUF;
}
m->backend_buffer = buf;
ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
if (auto st = transcribe::bin_loader::stream_tensor_data_from_bin(bm.path, stream_slots, "whisper");
st != TRANSCRIBE_OK) {
return st;
}
{
std::vector<float> filterbank = bm.mel_filterbank; std::vector<float> window = hann_periodic(m->hparams.fe_n_fft);
if (m->weights.frontend.mel_filterbank != nullptr) {
const size_t want = ggml_nbytes(m->weights.frontend.mel_filterbank);
const size_t have = filterbank.size() * sizeof(float);
if (have != want) {
log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
"%s: mel filterbank size mismatch "
"(have %zu bytes, expected %zu)",
kTag, have, want);
return TRANSCRIBE_ERR_GGUF;
}
ggml_backend_tensor_set(m->weights.frontend.mel_filterbank, filterbank.data(), 0, have);
}
if (m->weights.frontend.window != nullptr) {
const size_t want = ggml_nbytes(m->weights.frontend.window);
const size_t have = window.size() * sizeof(float);
if (have != want) {
log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
"%s: hann window size mismatch "
"(have %zu bytes, expected %zu)",
kTag, have, want);
return TRANSCRIBE_ERR_GGUF;
}
ggml_backend_tensor_set(m->weights.frontend.window, window.data(), 0, have);
}
if (auto st = install_mel_from_buffers(m->hparams, std::move(filterbank), std::move(window), m->mel);
st != TRANSCRIBE_OK) {
return st;
}
}
m->t_load_us = ggml_time_us() - t_load_start;
*out_model = m.release();
return TRANSCRIBE_OK;
}
}