mirror of
https://github.com/Tencent/ncnn.git
synced 2026-10-17 04:41:47 +08:00
nihui
941 lines
28 KiB
C++
941 lines
28 KiB
C++
// Copyright 2025 Tencent
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
// whisper speech recognition implemented with ncnn library
|
|
|
|
// convert openai-whisper checkpoints to ncnn models
|
|
// 1. install pnnx via pip install -U pnnx
|
|
// 2. obtain export_ncnn.py script from https://github.com/nihui/ncnn-android-whisper
|
|
// 3. edit export_ncnn.py for changing the models among tiny/base/small/medium/large-v3-turbo
|
|
// 4. make sure you have good internet connection
|
|
// python export_ncnn.py
|
|
|
|
// convert vocab.json to simple whisper_vocab.txt
|
|
// 1. obtain vocab.json file from https://huggingface.co/openai/whisper-tiny/blob/main/vocab.json
|
|
// 2. convert json dict into plain list, save to whisper_vocab.txt
|
|
|
|
// NOTE large-v3-turbo has special token ids from others, one more language(yue) and does not support translation
|
|
|
|
#include "net.h"
|
|
#include "layer.h"
|
|
#include "layer_type.h"
|
|
|
|
#include <float.h>
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <algorithm>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
// https://huggingface.co/openai/whisper-tiny/blob/main/tokenizer_config.json
|
|
static const int token_endoftext = 50257;
|
|
static const int token_startoftranscript = 50258;
|
|
static const int token_lang_first = 50259;
|
|
static const int token_lang_last = 50357;
|
|
static const int token_lang_count = token_lang_last - token_lang_first + 1;
|
|
// clang-format off
|
|
// *INDENT-OFF*
|
|
static const char* token_langs[] = {
|
|
"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"
|
|
};
|
|
// *INDENT-ON*
|
|
// clang-format on
|
|
static const int token_translate = 50358;
|
|
static const int token_transcribe = 50359;
|
|
static const int token_startoflm = 50360;
|
|
static const int token_startofprev = 50361;
|
|
static const int token_nocaptions = 50362;
|
|
static const int token_notimestamps = 50363;
|
|
static const int token_timestamp_first = 50364;
|
|
static const int token_timestamp_last = 51864;
|
|
|
|
// https://huggingface.co/openai/whisper-large-v3-turbo/blob/main/tokenizer_config.json
|
|
// static const int token_endoftext = 50257;
|
|
// static const int token_startoftranscript = 50258;
|
|
// static const int token_lang_first = 50259;
|
|
// static const int token_lang_last = 50357;
|
|
// static const int token_lang_count = token_lang_last - token_lang_first + 1;
|
|
// // clang-format off
|
|
// // *INDENT-OFF*
|
|
// static const char* token_langs[] = {
|
|
// "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"
|
|
// };
|
|
// // *INDENT-ON*
|
|
// // clang-format on
|
|
// static const int token_translate = 50359;
|
|
// static const int token_transcribe = 50360;
|
|
// static const int token_startoflm = 50361;
|
|
// static const int token_startofprev = 50362;
|
|
// static const int token_nospeech = 50363;
|
|
// static const int token_notimestamps = 50364;
|
|
// static const int token_timestamp_first = 50365;
|
|
// static const int token_timestamp_last = 51865;
|
|
|
|
// tokenizer for handling text tokens
|
|
class Tokenizer
|
|
{
|
|
public:
|
|
std::vector<std::string> reverse_vocab;
|
|
|
|
uint8_t byte_decoder[512]; // unicode code point to byte value
|
|
|
|
// generate byte decoder for tokenization
|
|
void generate_byte_decoder()
|
|
{
|
|
// initialize array to 0
|
|
memset(byte_decoder, 0, 512 * sizeof(uint8_t));
|
|
|
|
// define function to check if char is in "printable" range
|
|
auto is_printable = [](int b) {
|
|
return (b >= '!' && b <= '~') // '!' to '~'
|
|
|| (b >= 161 && b <= 172) // '¡' to '¬'
|
|
|| (b >= 174 && b <= 255); // '®' to 'ÿ'
|
|
};
|
|
|
|
// handle "printable" characters
|
|
// for these chars, key and value are the same
|
|
for (int b = 0; b < 256; ++b)
|
|
{
|
|
if (is_printable(b))
|
|
{
|
|
byte_decoder[b] = static_cast<uint8_t>(b);
|
|
}
|
|
}
|
|
|
|
// handle remaining characters
|
|
// for these chars, key starts from 256 and increments
|
|
int n = 0;
|
|
for (int b = 0; b < 256; ++b)
|
|
{
|
|
if (!is_printable(b))
|
|
{
|
|
byte_decoder[256 + n] = static_cast<uint8_t>(b);
|
|
n++;
|
|
}
|
|
}
|
|
}
|
|
|
|
// convert utf-8 string to code points
|
|
std::vector<uint32_t> utf8_to_codepoints(const std::string& s) const
|
|
{
|
|
std::vector<uint32_t> codepoints;
|
|
for (size_t i = 0; i < s.length();)
|
|
{
|
|
uint32_t cp = 0;
|
|
int len = 0;
|
|
unsigned char c = s[i];
|
|
|
|
if (c < 0x80) // 1-byte
|
|
{
|
|
cp = c;
|
|
len = 1;
|
|
}
|
|
else if ((c & 0xE0) == 0xC0) // 2-byte
|
|
{
|
|
cp = ((s[i] & 0x1F) << 6) | (s[i + 1] & 0x3F);
|
|
len = 2;
|
|
}
|
|
else if ((c & 0xF0) == 0xE0) // 3-byte
|
|
{
|
|
cp = ((s[i] & 0x0F) << 12) | ((s[i + 1] & 0x3F) << 6) | (s[i + 2] & 0x3F);
|
|
len = 3;
|
|
}
|
|
else if ((c & 0xF8) == 0xF0) // 4-byte
|
|
{
|
|
cp = ((s[i] & 0x07) << 18) | ((s[i + 1] & 0x3F) << 12) | ((s[i + 2] & 0x3F) << 6) | (s[i + 3] & 0x3F);
|
|
len = 4;
|
|
}
|
|
else
|
|
{
|
|
// invalid utf-8 start byte, skip
|
|
i++;
|
|
continue;
|
|
}
|
|
codepoints.push_back(cp);
|
|
i += len;
|
|
}
|
|
return codepoints;
|
|
}
|
|
|
|
bool load(const char* vocab_path)
|
|
{
|
|
// generate decoder when loading
|
|
generate_byte_decoder();
|
|
|
|
{
|
|
FILE* fp = fopen(vocab_path, "rb");
|
|
if (!fp)
|
|
{
|
|
fprintf(stderr, "fopen %s failed\n", vocab_path);
|
|
return false;
|
|
}
|
|
|
|
char line[256];
|
|
while (!feof(fp))
|
|
{
|
|
char* s = fgets(line, 255, fp);
|
|
if (!s)
|
|
break;
|
|
|
|
int vocab_len = strlen(line);
|
|
if (vocab_len > 1)
|
|
{
|
|
// drop the tail newline
|
|
vocab_len -= 1;
|
|
}
|
|
|
|
reverse_vocab.push_back(std::string(line, vocab_len));
|
|
}
|
|
|
|
fclose(fp);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
// decode token ids to text
|
|
std::string decode(const std::vector<int>& tokens) const
|
|
{
|
|
std::string outstring;
|
|
bool in_timestamp = false;
|
|
|
|
// step 1: concatenate token ids to a string with special unicode characters
|
|
std::string text_buffer;
|
|
for (int token_id : tokens)
|
|
{
|
|
if (token_id < token_endoftext)
|
|
{
|
|
text_buffer += reverse_vocab[token_id];
|
|
continue;
|
|
}
|
|
|
|
// handle timestamp tokens
|
|
if (token_id >= token_timestamp_first && token_id <= token_timestamp_last)
|
|
{
|
|
int timestamp = (token_id - token_timestamp_first) * 2;
|
|
|
|
char tmp[256];
|
|
sprintf(tmp, " [%d.%02d] ", timestamp / 100, timestamp % 100);
|
|
|
|
if (in_timestamp)
|
|
{
|
|
// step 2: translate the special string back to original byte stream
|
|
std::vector<uint32_t> codepoints = utf8_to_codepoints(text_buffer);
|
|
|
|
std::vector<uint8_t> byte_sequence;
|
|
for (uint32_t cp : codepoints)
|
|
{
|
|
byte_sequence.push_back(byte_decoder[cp]);
|
|
}
|
|
|
|
std::string s(byte_sequence.begin(), byte_sequence.end());
|
|
|
|
text_buffer.clear();
|
|
|
|
outstring += s;
|
|
outstring += tmp;
|
|
outstring += "\n";
|
|
|
|
in_timestamp = false;
|
|
}
|
|
else
|
|
{
|
|
outstring += tmp;
|
|
in_timestamp = true;
|
|
}
|
|
}
|
|
|
|
// ignore functional/special tokens
|
|
}
|
|
|
|
if (!text_buffer.empty())
|
|
{
|
|
// step 2: translate the special string back to original byte stream
|
|
std::vector<uint32_t> codepoints = utf8_to_codepoints(text_buffer);
|
|
|
|
std::vector<uint8_t> byte_sequence;
|
|
for (uint32_t cp : codepoints)
|
|
{
|
|
byte_sequence.push_back(byte_decoder[cp]);
|
|
}
|
|
|
|
std::string s(byte_sequence.begin(), byte_sequence.end());
|
|
|
|
outstring += s;
|
|
}
|
|
|
|
return outstring;
|
|
}
|
|
};
|
|
|
|
// result class for beam search
|
|
class Result
|
|
{
|
|
public:
|
|
std::vector<int> ids;
|
|
float score;
|
|
|
|
std::vector<ncnn::Mat> kvcache;
|
|
};
|
|
|
|
// main whisper implementation class
|
|
class Whisper
|
|
{
|
|
public:
|
|
int load();
|
|
|
|
int detect_lang(const std::vector<short>& samples, std::string& lang) const;
|
|
int transcribe(const std::vector<short>& samples, const char* lang, std::string& text) const;
|
|
|
|
protected:
|
|
int extract_fbank_feature(const std::vector<short>& samples, ncnn::Mat& input_features) const;
|
|
int run_encoder(const ncnn::Mat& input_features, ncnn::Mat& encoder_states) const;
|
|
int run_decoder_prefill(const std::vector<int>& tokens, const ncnn::Mat& encoder_states, ncnn::Mat& last_logits, std::vector<ncnn::Mat>& out_kvcache) const;
|
|
int run_decoder_step(const std::vector<int>& tokens, const ncnn::Mat& encoder_states, ncnn::Mat& last_logits, const std::vector<ncnn::Mat>& kvcache, std::vector<ncnn::Mat>& out_kvcache) const;
|
|
|
|
protected:
|
|
ncnn::Net fbank;
|
|
|
|
ncnn::Net encoder;
|
|
|
|
ncnn::Net embed_token;
|
|
ncnn::Net embed_position;
|
|
ncnn::Net decoder;
|
|
|
|
ncnn::Net proj_out;
|
|
|
|
Tokenizer tokenizer;
|
|
|
|
protected:
|
|
std::vector<int> kv_cache_indexes;
|
|
std::vector<int> out_kv_cache_indexes;
|
|
};
|
|
|
|
int Whisper::load()
|
|
{
|
|
// whisper models could be found at
|
|
// https://github.com/nihui/ncnn-android-whisper/releases
|
|
// https://github.com/nihui/ncnn-android-whisper/tree/master/app/src/main/assets
|
|
|
|
fbank.opt.use_vulkan_compute = true;
|
|
fbank.opt.use_fp16_packed = false;
|
|
fbank.opt.use_fp16_storage = false;
|
|
fbank.opt.use_fp16_arithmetic = false;
|
|
|
|
encoder.opt.use_vulkan_compute = true;
|
|
encoder.opt.use_fp16_packed = false;
|
|
encoder.opt.use_fp16_storage = false;
|
|
encoder.opt.use_fp16_arithmetic = false;
|
|
|
|
decoder.opt.use_vulkan_compute = true;
|
|
decoder.opt.use_fp16_packed = false;
|
|
decoder.opt.use_fp16_storage = false;
|
|
decoder.opt.use_fp16_arithmetic = false;
|
|
|
|
proj_out.opt.use_vulkan_compute = true;
|
|
proj_out.opt.use_fp16_packed = false;
|
|
proj_out.opt.use_fp16_storage = false;
|
|
proj_out.opt.use_fp16_arithmetic = false;
|
|
|
|
fbank.load_param("whisper_tiny_fbank.ncnn.param");
|
|
fbank.load_model("whisper_tiny_fbank.ncnn.bin");
|
|
|
|
encoder.load_param("whisper_tiny_encoder.ncnn.param");
|
|
encoder.load_model("whisper_tiny_encoder.ncnn.bin");
|
|
|
|
embed_token.load_param("whisper_tiny_embed_token.ncnn.param");
|
|
embed_token.load_model("whisper_tiny_embed_token.ncnn.bin");
|
|
|
|
embed_position.load_param("whisper_tiny_embed_position.ncnn.param");
|
|
embed_position.load_model("whisper_tiny_embed_position.ncnn.bin");
|
|
|
|
decoder.load_param("whisper_tiny_decoder.ncnn.param");
|
|
decoder.load_model("whisper_tiny_decoder.ncnn.bin");
|
|
|
|
proj_out.load_param("whisper_tiny_proj_out.ncnn.param");
|
|
proj_out.load_model("whisper_tiny_proj_out.ncnn.bin");
|
|
|
|
// fbank.load_param("whisper_large_v3_turbo_fbank.ncnn.param");
|
|
// fbank.load_model("whisper_large_v3_turbo_fbank.ncnn.bin");
|
|
//
|
|
// encoder.load_param("whisper_large_v3_turbo_encoder.ncnn.param");
|
|
// encoder.load_model("whisper_large_v3_turbo_encoder.ncnn.bin");
|
|
//
|
|
// embed_token.load_param("whisper_large_v3_turbo_embed_token.ncnn.param");
|
|
// embed_token.load_model("whisper_large_v3_turbo_embed_token.ncnn.bin");
|
|
//
|
|
// embed_position.load_param("whisper_large_v3_turbo_embed_position.ncnn.param");
|
|
// embed_position.load_model("whisper_large_v3_turbo_embed_position.ncnn.bin");
|
|
//
|
|
// decoder.load_param("whisper_large_v3_turbo_decoder.ncnn.param");
|
|
// decoder.load_model("whisper_large_v3_turbo_decoder.ncnn.bin");
|
|
//
|
|
// proj_out.load_param("whisper_large_v3_turbo_proj_out.ncnn.param");
|
|
// proj_out.load_model("whisper_large_v3_turbo_proj_out.ncnn.bin");
|
|
|
|
tokenizer.load("whisper_vocab.txt");
|
|
|
|
// resolve kv cache blob indexes
|
|
for (size_t i = 0; i < decoder.layers().size(); i++)
|
|
{
|
|
const ncnn::Layer* mha = decoder.layers()[i];
|
|
if (mha->typeindex != ncnn::LayerType::MultiHeadAttention)
|
|
continue;
|
|
|
|
const size_t input_count = mha->bottoms.size();
|
|
const size_t output_count = mha->tops.size();
|
|
|
|
if (output_count == 3)
|
|
{
|
|
kv_cache_indexes.push_back(mha->bottoms[input_count - 2]);
|
|
kv_cache_indexes.push_back(mha->bottoms[input_count - 1]);
|
|
out_kv_cache_indexes.push_back(mha->tops[output_count - 2]);
|
|
out_kv_cache_indexes.push_back(mha->tops[output_count - 1]);
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
// apply log_softmax in-place
|
|
static void log_softmax_inplace(ncnn::Mat& m)
|
|
{
|
|
ncnn::Option opt;
|
|
opt.use_packing_layout = false;
|
|
opt.use_fp16_storage = false;
|
|
|
|
{
|
|
ncnn::Layer* softmax = ncnn::create_layer_cpu("Softmax");
|
|
ncnn::ParamDict pd;
|
|
pd.set(0, 0); // axis
|
|
softmax->load_param(pd);
|
|
softmax->forward_inplace(m, opt);
|
|
delete softmax;
|
|
}
|
|
|
|
{
|
|
ncnn::Layer* log = ncnn::create_layer_cpu("UnaryOp");
|
|
ncnn::ParamDict pd;
|
|
pd.set(0, 8); // log
|
|
log->load_param(pd);
|
|
log->forward_inplace(m, opt);
|
|
delete log;
|
|
}
|
|
}
|
|
|
|
int Whisper::detect_lang(const std::vector<short>& samples, std::string& lang) const
|
|
{
|
|
std::vector<int> ids(1);
|
|
ids[0] = token_startoftranscript;
|
|
|
|
ncnn::Mat input_features;
|
|
extract_fbank_feature(samples, input_features);
|
|
|
|
ncnn::Mat encoder_states;
|
|
run_encoder(input_features, encoder_states);
|
|
|
|
ncnn::Mat logits;
|
|
std::vector<ncnn::Mat> out_kvcache;
|
|
run_decoder_prefill(ids, encoder_states, logits, out_kvcache);
|
|
|
|
// find the lang token with highest prob
|
|
// we are only interested in lang part and no_speech
|
|
int lang_id = token_lang_first;
|
|
float max_prob = logits[token_lang_first];
|
|
for (int i = token_lang_first; i <= token_lang_last; i++)
|
|
{
|
|
float prob = logits[i];
|
|
if (prob > max_prob)
|
|
{
|
|
max_prob = prob;
|
|
lang_id = i;
|
|
}
|
|
}
|
|
|
|
lang = token_langs[lang_id - token_lang_first];
|
|
|
|
return 0;
|
|
}
|
|
|
|
int Whisper::transcribe(const std::vector<short>& samples, const char* lang, std::string& text) const
|
|
{
|
|
// find lang token id by lang string
|
|
int token_lang = -1;
|
|
for (int i = 0; i < token_lang_count; i++)
|
|
{
|
|
if (strcmp(token_langs[i], lang) == 0)
|
|
{
|
|
token_lang = token_lang_first + i;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (token_lang == -1)
|
|
{
|
|
fprintf(stderr, "language %s not supported\n", lang);
|
|
return -1;
|
|
}
|
|
|
|
// initialize with prompt tokens
|
|
std::vector<int> ids(4);
|
|
ids[0] = token_startoftranscript;
|
|
ids[1] = token_lang;
|
|
ids[2] = token_transcribe;
|
|
ids[3] = token_notimestamps;
|
|
|
|
ncnn::Mat input_features;
|
|
extract_fbank_feature(samples, input_features);
|
|
|
|
ncnn::Mat encoder_states;
|
|
run_encoder(input_features, encoder_states);
|
|
|
|
const int beam_size = 5;
|
|
const int max_candidates = 5;
|
|
|
|
std::vector<Result> finished_beams;
|
|
|
|
std::vector<Result> beams(1);
|
|
beams[0].ids = ids;
|
|
beams[0].score = 0.f;
|
|
|
|
int step = 0;
|
|
|
|
// beam search loop
|
|
for (;;)
|
|
{
|
|
std::vector<Result> candidates;
|
|
|
|
for (size_t i = 0; i < beams.size(); i++)
|
|
{
|
|
const Result& beam = beams[i];
|
|
|
|
ncnn::Mat logits;
|
|
std::vector<ncnn::Mat> out_kvcache;
|
|
if (step == 0)
|
|
{
|
|
run_decoder_prefill(beam.ids, encoder_states, logits, out_kvcache);
|
|
}
|
|
else
|
|
{
|
|
run_decoder_step(beam.ids, encoder_states, logits, beam.kvcache, out_kvcache);
|
|
}
|
|
|
|
log_softmax_inplace(logits);
|
|
|
|
// get topk candidates
|
|
const int topk = 5;
|
|
std::vector<std::pair<float, int> > vec(logits.w);
|
|
for (int j = 0; j < logits.w; j++)
|
|
{
|
|
vec[j] = std::make_pair(logits[j], j);
|
|
}
|
|
std::partial_sort(vec.begin(), vec.begin() + topk, vec.end(), std::greater<std::pair<float, int> >());
|
|
|
|
for (int j = 0; j < topk; j++)
|
|
{
|
|
int next_id = vec[j].second;
|
|
float next_id_score = vec[j].first;
|
|
|
|
Result candidate;
|
|
candidate.ids = beam.ids;
|
|
candidate.ids.push_back(next_id);
|
|
candidate.score = beam.score + next_id_score;
|
|
candidate.kvcache = out_kvcache;
|
|
|
|
candidates.push_back(candidate);
|
|
}
|
|
}
|
|
|
|
// sort candidates by score
|
|
std::sort(candidates.begin(), candidates.end(), [](const Result& a, const Result& b) {
|
|
return a.score > b.score;
|
|
});
|
|
|
|
beams.clear();
|
|
for (size_t i = 0; i < candidates.size(); i++)
|
|
{
|
|
const Result& candidate = candidates[i];
|
|
|
|
if (candidate.ids.back() == token_endoftext)
|
|
{
|
|
finished_beams.push_back(candidate);
|
|
}
|
|
else
|
|
{
|
|
beams.push_back(candidate);
|
|
}
|
|
}
|
|
|
|
if (beams.size() > beam_size)
|
|
{
|
|
beams.resize(beam_size);
|
|
}
|
|
|
|
step++;
|
|
|
|
if (beams.empty())
|
|
{
|
|
break;
|
|
}
|
|
|
|
if (finished_beams.size() >= max_candidates)
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (finished_beams.empty())
|
|
{
|
|
// no results
|
|
return 0;
|
|
}
|
|
|
|
// find the best result based on average score
|
|
int max_avg_score_index = 0;
|
|
float max_avg_score = -FLT_MAX;
|
|
for (size_t i = 0; i < finished_beams.size(); i++)
|
|
{
|
|
const Result& result = finished_beams[i];
|
|
float avg_score = result.score / result.ids.size();
|
|
if (avg_score > max_avg_score)
|
|
{
|
|
max_avg_score_index = (int)i;
|
|
max_avg_score = avg_score;
|
|
}
|
|
}
|
|
|
|
const Result& best_result = finished_beams[max_avg_score_index];
|
|
|
|
text = tokenizer.decode(best_result.ids);
|
|
|
|
return 0;
|
|
}
|
|
|
|
int Whisper::extract_fbank_feature(const std::vector<short>& samples, ncnn::Mat& input_features) const
|
|
{
|
|
const int samples_size = (int)samples.size();
|
|
|
|
// pad to 480000, normalize samples to -1~1
|
|
ncnn::Mat waveform(480000);
|
|
waveform.fill(0.f);
|
|
{
|
|
for (int i = 0; i < samples_size; i++)
|
|
{
|
|
waveform[i] = samples[i] / 32768.0f;
|
|
}
|
|
}
|
|
|
|
ncnn::Extractor ex = fbank.create_extractor();
|
|
|
|
ex.input("in0", waveform);
|
|
|
|
ex.extract("out0", input_features);
|
|
|
|
// drop the last frame
|
|
{
|
|
ncnn::Mat input_features_3k(input_features.w - 1, input_features.h);
|
|
for (int i = 0; i < input_features.h; i++)
|
|
{
|
|
memcpy(input_features_3k.row(i), input_features.row(i), (input_features.w - 1) * sizeof(float));
|
|
}
|
|
input_features = input_features_3k;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int Whisper::run_encoder(const ncnn::Mat& input_features, ncnn::Mat& encoder_states) const
|
|
{
|
|
ncnn::Extractor ex = encoder.create_extractor();
|
|
|
|
ex.input("in0", input_features);
|
|
|
|
ex.extract("out0", encoder_states);
|
|
|
|
return 0;
|
|
}
|
|
|
|
int Whisper::run_decoder_prefill(const std::vector<int>& tokens, const ncnn::Mat& encoder_states, ncnn::Mat& last_logits, std::vector<ncnn::Mat>& out_kvcache) const
|
|
{
|
|
const int dst_seqlen = tokens.size();
|
|
|
|
// token embedding
|
|
ncnn::Mat token_embeds;
|
|
{
|
|
ncnn::Mat input_tokens(dst_seqlen);
|
|
int* p = input_tokens;
|
|
memcpy(p, tokens.data(), tokens.size() * sizeof(int));
|
|
|
|
ncnn::Extractor ex = embed_token.create_extractor();
|
|
ex.input("in0", input_tokens);
|
|
ex.extract("out0", token_embeds);
|
|
}
|
|
|
|
// position embedding
|
|
ncnn::Mat position_embeds;
|
|
{
|
|
ncnn::Mat input_positions(dst_seqlen);
|
|
int* p = input_positions;
|
|
for (int i = 0; i < dst_seqlen; i++)
|
|
{
|
|
p[i] = i;
|
|
}
|
|
|
|
ncnn::Extractor ex = embed_position.create_extractor();
|
|
ex.input("in0", input_positions);
|
|
ex.extract("out0", position_embeds);
|
|
}
|
|
|
|
// input embedding = token + position
|
|
ncnn::Mat input_embeds;
|
|
{
|
|
input_embeds.create_like(token_embeds);
|
|
for (int i = 0; i < input_embeds.total(); i++)
|
|
{
|
|
input_embeds[i] = token_embeds[i] + position_embeds[i];
|
|
}
|
|
}
|
|
|
|
// create attention mask (causal mask)
|
|
ncnn::Mat attention_mask(dst_seqlen, dst_seqlen);
|
|
attention_mask.fill(0.f);
|
|
for (int i = 0; i < dst_seqlen; i++)
|
|
{
|
|
for (int j = i + 1; j < dst_seqlen; j++)
|
|
{
|
|
attention_mask.row(i)[j] = -INFINITY;
|
|
}
|
|
}
|
|
|
|
ncnn::Mat output_states;
|
|
{
|
|
ncnn::Extractor ex = decoder.create_extractor();
|
|
ex.input("in0", input_embeds);
|
|
ex.input("in1", encoder_states);
|
|
ex.input("in2", attention_mask);
|
|
|
|
out_kvcache.resize(out_kv_cache_indexes.size());
|
|
for (size_t i = 0; i < out_kv_cache_indexes.size(); i++)
|
|
{
|
|
ex.extract(out_kv_cache_indexes[i], out_kvcache[i], 1);
|
|
}
|
|
|
|
ex.extract("out0", output_states);
|
|
}
|
|
|
|
// get last token's state for next token prediction
|
|
ncnn::Mat last_state = output_states.row_range(dst_seqlen - 1, 1).clone();
|
|
{
|
|
ncnn::Extractor ex = proj_out.create_extractor();
|
|
ex.input("in0", last_state);
|
|
ex.extract("out0", last_logits);
|
|
}
|
|
|
|
last_logits = last_logits.reshape(last_logits.w);
|
|
|
|
return 0;
|
|
}
|
|
|
|
int Whisper::run_decoder_step(const std::vector<int>& tokens, const ncnn::Mat& encoder_states, ncnn::Mat& last_logits, const std::vector<ncnn::Mat>& kvcache, std::vector<ncnn::Mat>& out_kvcache) const
|
|
{
|
|
const int token_id = tokens.back();
|
|
const int dst_seqlen = 1;
|
|
|
|
// token embedding
|
|
ncnn::Mat token_embeds;
|
|
{
|
|
ncnn::Mat input_tokens(dst_seqlen);
|
|
((int*)input_tokens)[0] = token_id;
|
|
|
|
ncnn::Extractor ex = embed_token.create_extractor();
|
|
ex.input("in0", input_tokens);
|
|
ex.extract("out0", token_embeds);
|
|
}
|
|
|
|
// position embedding
|
|
ncnn::Mat position_embeds;
|
|
{
|
|
ncnn::Mat input_positions(dst_seqlen);
|
|
((int*)input_positions)[0] = tokens.size() - 1;
|
|
|
|
ncnn::Extractor ex = embed_position.create_extractor();
|
|
ex.input("in0", input_positions);
|
|
ex.extract("out0", position_embeds);
|
|
}
|
|
|
|
// input embedding = token + position
|
|
ncnn::Mat input_embeds;
|
|
{
|
|
input_embeds.create_like(token_embeds);
|
|
for (int i = 0; i < input_embeds.total(); i++)
|
|
{
|
|
input_embeds[i] = token_embeds[i] + position_embeds[i];
|
|
}
|
|
}
|
|
|
|
// single token doesn't need attention mask
|
|
ncnn::Mat attention_mask(dst_seqlen, dst_seqlen);
|
|
attention_mask.fill(0.f);
|
|
|
|
ncnn::Mat output_states;
|
|
{
|
|
ncnn::Extractor ex = decoder.create_extractor();
|
|
ex.input("in0", input_embeds);
|
|
ex.input("in1", encoder_states);
|
|
ex.input("in2", attention_mask);
|
|
|
|
// pass in kv cache from previous steps
|
|
for (size_t i = 0; i < kv_cache_indexes.size(); i++)
|
|
{
|
|
ex.input(kv_cache_indexes[i], kvcache[i]);
|
|
}
|
|
|
|
// extract updated kv cache
|
|
out_kvcache.resize(out_kv_cache_indexes.size());
|
|
for (size_t i = 0; i < out_kv_cache_indexes.size(); i++)
|
|
{
|
|
ex.extract(out_kv_cache_indexes[i], out_kvcache[i], 1);
|
|
}
|
|
|
|
ex.extract("out0", output_states);
|
|
}
|
|
|
|
// get last token's state for prediction
|
|
ncnn::Mat last_state = output_states.row_range(dst_seqlen - 1, 1).clone();
|
|
{
|
|
ncnn::Extractor ex = proj_out.create_extractor();
|
|
ex.input("in0", last_state);
|
|
ex.extract("out0", last_logits);
|
|
}
|
|
|
|
last_logits = last_logits.reshape(last_logits.w);
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int load_wav_samples(const char* wavpath, std::vector<short>& samples)
|
|
{
|
|
FILE* fp = fopen(wavpath, "rb");
|
|
if (!fp)
|
|
{
|
|
fprintf(stderr, "open %s failed\n", wavpath);
|
|
return -1;
|
|
}
|
|
|
|
// https://stackoverflow.com/questions/1537964/visual-c-equivalent-of-gccs-attribute-packed
|
|
#ifdef _MSC_VER
|
|
#define PACK(__Declaration__) __pragma(pack(push, 1)) __Declaration__ __pragma(pack(pop))
|
|
#else
|
|
#define PACK(__Declaration__) __Declaration__ __attribute__((__packed__))
|
|
#endif
|
|
|
|
PACK(struct wav_header {
|
|
char riff[4];
|
|
uint32_t chunk_size;
|
|
char wave[4];
|
|
char fmt[4];
|
|
uint32_t subchunk1_size;
|
|
uint16_t audio_format;
|
|
uint16_t num_channels;
|
|
uint32_t sample_rate;
|
|
uint32_t byte_rate;
|
|
uint16_t block_align;
|
|
uint16_t bits_per_sample;
|
|
char data[4];
|
|
uint32_t data_size;
|
|
});
|
|
|
|
wav_header header;
|
|
if (fread(&header, sizeof(wav_header), 1, fp) != 1)
|
|
{
|
|
fprintf(stderr, "failed to read wav header from %s\n", wavpath);
|
|
fclose(fp);
|
|
return -1;
|
|
}
|
|
|
|
if (memcmp(header.riff, "RIFF", 4) != 0 || memcmp(header.wave, "WAVE", 4) != 0
|
|
|| memcmp(header.fmt, "fmt ", 4) != 0 || memcmp(header.data, "data", 4) != 0)
|
|
{
|
|
fprintf(stderr, "%s is not a valid wav file\n", wavpath);
|
|
fclose(fp);
|
|
return -1;
|
|
}
|
|
|
|
if (header.subchunk1_size != 16 || header.audio_format != 1 || header.num_channels != 1
|
|
|| header.sample_rate != 16000 || header.bits_per_sample != 16)
|
|
{
|
|
fprintf(stderr, "%s is not pcm s16le 16k wav\n", wavpath);
|
|
fprintf(stderr, "ffmpeg -i input.xxx -vn -c:a pcm_s16le -ac 1 -ar 16000 -fflags bitexact output.wav\n");
|
|
fclose(fp);
|
|
return -1;
|
|
}
|
|
|
|
fseek(fp, 0, SEEK_END);
|
|
long len = ftell(fp);
|
|
|
|
samples.resize((len - sizeof(wav_header)) / sizeof(short));
|
|
|
|
rewind(fp);
|
|
|
|
fseek(fp, sizeof(wav_header), SEEK_SET);
|
|
|
|
fread(samples.data(), 1, len - sizeof(wav_header), fp);
|
|
|
|
fclose(fp);
|
|
|
|
return 0;
|
|
}
|
|
|
|
int main(int argc, char** argv)
|
|
{
|
|
if (argc != 2)
|
|
{
|
|
fprintf(stderr, "Usage: %s [wavpath]\n", argv[0]);
|
|
return -1;
|
|
}
|
|
|
|
const char* wavpath = argv[1];
|
|
|
|
std::vector<short> samples;
|
|
int ret = load_wav_samples(wavpath, samples);
|
|
if (ret != 0)
|
|
{
|
|
fprintf(stderr, "load wav failed\n");
|
|
return -1;
|
|
}
|
|
|
|
if (samples.size() > 480000)
|
|
{
|
|
fprintf(stderr, "audio duration too long, truncate to 30s\n");
|
|
samples.resize(480000);
|
|
}
|
|
|
|
Whisper whisper;
|
|
whisper.load();
|
|
|
|
// detect language first
|
|
std::string lang;
|
|
whisper.detect_lang(samples, lang);
|
|
fprintf(stderr, "lang = %s\n", lang.c_str());
|
|
|
|
// transcribe audio to text
|
|
std::string text;
|
|
whisper.transcribe(samples, lang.c_str(), text);
|
|
fprintf(stderr, "text = %s\n", text.c_str());
|
|
|
|
return 0;
|
|
}
|