Sherpa-ONNX-KWS.AXERA / cpp /sherpa_kws_ax.cpp
HY-2012's picture
First commit
dc79b9f verified
Raw
History Blame Contribute Delete
26 kB
#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <limits>
#include <memory>
#include <queue>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "ax_engine_api.h"
#include "ax_sys_api.h"
#include "kaldi-native-fbank/csrc/online-feature.h"
#include "src/engine_wrapper.hpp"
#include "src/wav_reader.hpp"
#ifndef AXERA_TARGET_NAME
#define AXERA_TARGET_NAME "AXERA"
#endif
namespace {
constexpr int kSampleRate = 16000;
constexpr int kFeatureDim = 80;
constexpr int kEncoderDim = 320;
constexpr int kVocabSize = 263;
constexpr int kBlankId = 0;
constexpr int kUnkId = 2;
constexpr int kContextSize = 2;
using Clock = std::chrono::steady_clock;
double ElapsedSeconds(Clock::time_point begin, Clock::time_point end) {
return std::chrono::duration<double>(end - begin).count();
}
struct Args {
std::string models_dir = "models";
std::string tokens = "config/tokens.txt";
std::string keywords = "config/keywords.txt";
std::string initial_decoder = "config/sherpa_decoder_initial.bin";
std::string audio = "audio/sherpa/zh_0.wav";
int chunk_size = 8;
float default_score = 1.0f;
float default_threshold = 0.25f;
int trailing_blanks = 1;
int max_active_paths = 1;
};
void Usage(const char *program) {
std::printf(
"Usage: %s [--models-dir DIR] [--tokens FILE] [--keywords FILE] "
"[--initial-decoder-output FILE] [--audio WAV] [--chunk-size 8|16] "
"[--keywords-score VALUE] [--keywords-threshold VALUE] "
"[--num-trailing-blanks N] [--max-active-paths N]\n",
program);
}
Args ParseArgs(int argc, char **argv) {
Args args;
for (int i = 1; i < argc; ++i) {
const std::string key = argv[i];
auto value = [&]() -> std::string {
if (++i >= argc) throw std::runtime_error("Missing value for " + key);
return argv[i];
};
if (key == "--models-dir") {
args.models_dir = value();
} else if (key == "--tokens") {
args.tokens = value();
} else if (key == "--keywords") {
args.keywords = value();
} else if (key == "--initial-decoder-output") {
args.initial_decoder = value();
} else if (key == "--audio") {
args.audio = value();
} else if (key == "--chunk-size") {
args.chunk_size = std::stoi(value());
} else if (key == "--keywords-score") {
args.default_score = std::stof(value());
} else if (key == "--keywords-threshold") {
args.default_threshold = std::stof(value());
} else if (key == "--num-trailing-blanks") {
args.trailing_blanks = std::stoi(value());
} else if (key == "--max-active-paths") {
args.max_active_paths = std::stoi(value());
} else if (key == "-h" || key == "--help") {
Usage(argv[0]);
std::exit(0);
} else {
throw std::runtime_error("Unknown argument: " + key);
}
}
if (args.chunk_size != 8 && args.chunk_size != 16) {
throw std::runtime_error("--chunk-size must be 8 or 16");
}
if (args.max_active_paths < 1 || args.max_active_paths > 32) {
throw std::runtime_error("--max-active-paths must be in [1, 32]");
}
return args;
}
std::string Join(const std::string &left, const std::string &right) {
return left.empty() || left.back() == '/' ? left + right
: left + "/" + right;
}
std::string ModelPath(const Args &args, const std::string &component) {
return Join(args.models_dir,
"sherpa__" + component + "-epoch-13-avg-2-chunk-" +
std::to_string(args.chunk_size) + "-left-64.axmodel");
}
class AxRuntime {
public:
AxRuntime() {
if (AX_SYS_Init() != 0) throw std::runtime_error("AX_SYS_Init failed");
sys_initialized_ = true;
AX_ENGINE_NPU_ATTR_T attr{};
if (AX_ENGINE_Init(&attr) != 0) {
AX_SYS_Deinit();
sys_initialized_ = false;
throw std::runtime_error("AX_ENGINE_Init failed");
}
engine_initialized_ = true;
}
~AxRuntime() {
if (engine_initialized_) AX_ENGINE_Deinit();
if (sys_initialized_) AX_SYS_Deinit();
}
private:
bool sys_initialized_ = false;
bool engine_initialized_ = false;
};
std::unordered_map<std::string, int32_t> LoadTokens(const std::string &path) {
std::ifstream input(path);
if (!input) throw std::runtime_error("Cannot open tokens: " + path);
std::unordered_map<std::string, int32_t> result;
std::string line;
int line_number = 0;
while (std::getline(input, line)) {
++line_number;
const std::size_t split = line.find_last_of(' ');
if (split == std::string::npos) {
throw std::runtime_error("Invalid token line " +
std::to_string(line_number));
}
result[line.substr(0, split)] = std::stoi(line.substr(split + 1));
}
return result;
}
struct Keyword {
std::vector<int32_t> tokens;
std::string phrase;
float score = 1.0f;
float threshold = 0.25f;
};
std::vector<Keyword> LoadKeywords(
const std::string &path,
const std::unordered_map<std::string, int32_t> &token_table,
float default_score, float default_threshold) {
std::ifstream input(path);
if (!input) throw std::runtime_error("Cannot open keywords: " + path);
std::vector<Keyword> result;
std::string line;
int line_number = 0;
while (std::getline(input, line)) {
++line_number;
std::istringstream stream(line);
std::string part;
Keyword keyword;
keyword.score = default_score;
keyword.threshold = default_threshold;
while (stream >> part) {
if (part.front() == '@') {
keyword.phrase = part.substr(1);
} else if (part.front() == ':') {
keyword.score = std::stof(part.substr(1));
} else if (part.front() == '#') {
keyword.threshold = std::stof(part.substr(1));
} else {
const auto it = token_table.find(part);
if (it == token_table.end()) {
throw std::runtime_error("Unknown keyword token at line " +
std::to_string(line_number) + ": " + part);
}
keyword.tokens.push_back(it->second);
}
}
if (!keyword.tokens.empty()) {
if (keyword.phrase.empty()) keyword.phrase = line;
result.push_back(std::move(keyword));
}
}
if (result.empty()) throw std::runtime_error("No keywords found in " + path);
return result;
}
struct ContextNode {
int32_t token = -1;
int level = 0;
float token_score = 0.0f;
float node_score = 0.0f;
float output_score = 0.0f;
bool is_end = false;
std::string phrase;
float threshold = 0.0f;
std::unordered_map<int32_t, std::unique_ptr<ContextNode>> children;
ContextNode *fail = nullptr;
ContextNode *output = nullptr;
};
class ContextGraph {
public:
struct Transition {
float score = 0.0f;
ContextNode *state = nullptr;
ContextNode *matched = nullptr;
};
explicit ContextGraph(const std::vector<Keyword> &keywords) {
root_.fail = &root_;
for (const Keyword &keyword : keywords) {
ContextNode *node = &root_;
for (std::size_t i = 0; i < keyword.tokens.size(); ++i) {
const int32_t token = keyword.tokens[i];
auto &child = node->children[token];
if (!child) {
child = std::make_unique<ContextNode>();
child->token = token;
child->level = static_cast<int>(i + 1);
child->token_score = keyword.score;
child->node_score = node->node_score + keyword.score;
const bool is_end = i + 1 == keyword.tokens.size();
child->output_score = is_end ? child->node_score : 0.0f;
child->is_end = is_end;
} else {
child->token_score = std::max(child->token_score, keyword.score);
child->node_score = node->node_score + child->token_score;
child->is_end = child->is_end || i + 1 == keyword.tokens.size();
child->output_score = child->is_end ? child->node_score : 0.0f;
}
node = child.get();
}
node->is_end = true;
node->phrase = keyword.phrase;
node->threshold = keyword.threshold;
}
FillFailureLinks();
}
ContextNode *Root() { return &root_; }
Transition Forward(ContextNode *state, int32_t token) {
ContextNode *node = nullptr;
float score = 0.0f;
const auto direct = state->children.find(token);
if (direct != state->children.end()) {
node = direct->second.get();
score = node->token_score;
} else {
node = state->fail;
while (node->children.count(token) == 0) {
node = node->fail;
if (node->token == -1) break;
}
const auto fallback = node->children.find(token);
if (fallback != node->children.end()) node = fallback->second.get();
score = node->node_score - state->node_score;
}
ContextNode *matched = node->is_end ? node : node->output;
return {score + node->output_score, node, matched};
}
ContextNode *Matched(ContextNode *state) {
return state->is_end ? state : state->output;
}
private:
void FillFailureLinks() {
std::queue<ContextNode *> queue;
for (auto &entry : root_.children) {
entry.second->fail = &root_;
queue.push(entry.second.get());
}
while (!queue.empty()) {
ContextNode *current = queue.front();
queue.pop();
for (auto &entry : current->children) {
const int32_t token = entry.first;
ContextNode *child = entry.second.get();
ContextNode *failure = current->fail;
while (failure != &root_ && failure->children.count(token) == 0) {
failure = failure->fail;
}
const auto it = failure->children.find(token);
child->fail = (it != failure->children.end() && it->second.get() != child)
? it->second.get()
: &root_;
ContextNode *output = child->fail;
while (output != &root_ && !output->is_end) output = output->fail;
child->output = output->is_end ? output : nullptr;
if (child->output) child->output_score += child->output->output_score;
queue.push(child);
}
}
}
ContextNode root_;
};
std::array<float, kEncoderDim> LoadInitialDecoder(const std::string &path) {
std::ifstream input(path, std::ios::binary);
if (!input) {
throw std::runtime_error("Cannot open initial decoder output: " + path);
}
char magic[8]{};
input.read(magic, 8);
uint32_t version = 0;
uint32_t count = 0;
input.read(reinterpret_cast<char *>(&version), sizeof(version));
input.read(reinterpret_cast<char *>(&count), sizeof(count));
if (std::memcmp(magic, "SHDEC1", 6) != 0 || version != 1 ||
count != kEncoderDim) {
throw std::runtime_error("Invalid initial decoder output file");
}
std::array<float, kEncoderDim> result{};
input.read(reinterpret_cast<char *>(result.data()),
result.size() * sizeof(float));
if (!input) throw std::runtime_error("Truncated initial decoder output");
return result;
}
std::vector<float> ComputeFbank(const PcmWav &wav) {
if (wav.sample_rate != kSampleRate) {
throw std::runtime_error("Input WAV must use 16 kHz sample rate");
}
knf::FbankOptions options;
options.frame_opts.samp_freq = kSampleRate;
options.frame_opts.dither = 0.0f;
options.frame_opts.frame_length_ms = 25.0f;
options.frame_opts.frame_shift_ms = 10.0f;
options.frame_opts.snip_edges = false;
options.frame_opts.window_type = "povey";
options.mel_opts.num_bins = kFeatureDim;
options.mel_opts.low_freq = 20.0f;
options.mel_opts.high_freq = -400.0f;
options.energy_floor = 0.0f;
knf::OnlineFbank fbank(options);
std::vector<float> waveform(wav.samples.size());
for (std::size_t i = 0; i < wav.samples.size(); ++i) {
waveform[i] = static_cast<float>(wav.samples[i]) / 32768.0f;
}
fbank.AcceptWaveform(kSampleRate, waveform.data(),
static_cast<int32_t>(waveform.size()));
std::vector<float> tail(static_cast<std::size_t>(0.8f * kSampleRate));
fbank.AcceptWaveform(kSampleRate, tail.data(),
static_cast<int32_t>(tail.size()));
fbank.InputFinished();
const int frames = fbank.NumFramesReady();
std::vector<float> result(static_cast<std::size_t>(frames) * kFeatureDim);
for (int i = 0; i < frames; ++i) {
std::memcpy(result.data() + static_cast<std::size_t>(i) * kFeatureDim,
fbank.GetFrame(i), kFeatureDim * sizeof(float));
}
return result;
}
class KeywordDecoder {
public:
KeywordDecoder(EngineWrapper *decoder, EngineWrapper *joiner,
ContextGraph *graph,
const std::array<float, kEncoderDim> &initial_decoder,
int trailing_blanks, int max_active_paths)
: decoder_(decoder),
joiner_(joiner),
graph_(graph),
initial_decoder_(initial_decoder),
required_trailing_blanks_(trailing_blanks),
max_active_paths_(max_active_paths) {
Reset();
}
void Reset() {
hypotheses_.clear();
Hypothesis initial;
initial.history = {-1, kBlankId};
initial.context_state = graph_->Root();
hypotheses_.push_back(std::move(initial));
}
int trailing_blanks() const { return BestHypothesis().trailing_blanks; }
std::string DecodeFrame(const float *encoder_frame) {
std::vector<Candidate> candidates;
candidates.reserve(hypotheses_.size() * kVocabSize);
for (std::size_t hyp_index = 0; hyp_index < hypotheses_.size();
++hyp_index) {
const Hypothesis &hypothesis = hypotheses_[hyp_index];
const std::array<float, kEncoderDim> decoder_output =
DecoderOutput(hypothesis);
if (joiner_->SetInputByName("encoder_out", encoder_frame,
kEncoderDim * sizeof(float)) != 0 ||
joiner_->SetInputByName("decoder_out", decoder_output.data(),
kEncoderDim * sizeof(float)) != 0 ||
joiner_->RunSync() != 0) {
throw std::runtime_error("Joiner inference failed");
}
std::array<float, kVocabSize> logits{};
if (joiner_->GetOutputByName("logit", logits.data(),
logits.size() * sizeof(float)) != 0) {
throw std::runtime_error("Failed to read joiner output");
}
const float maximum = *std::max_element(logits.begin(), logits.end());
double sum = 0.0;
for (float value : logits) sum += std::exp(value - maximum);
const double log_normalizer = maximum + std::log(sum);
for (int32_t token = 0; token < kVocabSize; ++token) {
Candidate candidate;
candidate.hypothesis = hyp_index;
candidate.token = token;
candidate.acoustic_probability =
static_cast<float>(std::exp(logits[token] - log_normalizer));
candidate.selection_score =
hypothesis.log_probability + logits[token] - log_normalizer;
candidate.score = candidate.selection_score;
candidate.context_state = hypothesis.context_state;
if (token != kBlankId && token != kUnkId) {
const auto transition = graph_->Forward(hypothesis.context_state, token);
candidate.score += transition.score;
candidate.context_state = transition.state;
}
candidates.push_back(candidate);
}
}
const std::size_t keep = std::min<std::size_t>(
static_cast<std::size_t>(max_active_paths_), candidates.size());
std::partial_sort(candidates.begin(), candidates.begin() + keep,
candidates.end(),
[](const Candidate &left, const Candidate &right) {
return left.selection_score > right.selection_score;
});
std::unordered_map<std::string, Hypothesis> merged;
for (std::size_t i = 0; i < keep; ++i) {
const Candidate &candidate = candidates[i];
Hypothesis next = hypotheses_[candidate.hypothesis];
next.log_probability = candidate.score;
if (candidate.token != kBlankId && candidate.token != kUnkId) {
next.history.push_back(candidate.token);
next.probabilities.push_back(candidate.acoustic_probability);
next.trailing_blanks = 0;
next.context_state = candidate.context_state;
if (next.context_state == graph_->Root()) {
next.history = {-1, kBlankId};
next.probabilities.clear();
}
} else {
++next.trailing_blanks;
}
const std::string key = HistoryKey(next.history);
const auto existing = merged.find(key);
if (existing == merged.end()) {
merged.emplace(key, std::move(next));
} else {
existing->second.log_probability =
LogAdd(existing->second.log_probability, next.log_probability);
}
}
hypotheses_.clear();
hypotheses_.reserve(merged.size());
for (auto &entry : merged) hypotheses_.push_back(std::move(entry.second));
const Hypothesis &best = BestHypothesis();
ContextNode *matched = graph_->Matched(best.context_state);
if (!matched || best.trailing_blanks <= required_trailing_blanks_ ||
best.probabilities.size() < static_cast<std::size_t>(matched->level)) {
return {};
}
float acoustic_score = 0.0f;
const std::size_t begin = best.probabilities.size() - matched->level;
for (std::size_t i = begin; i < best.probabilities.size(); ++i) {
acoustic_score += best.probabilities[i];
}
acoustic_score /= matched->level;
if (acoustic_score < matched->threshold) return {};
const std::string phrase = matched->phrase;
Reset();
return phrase;
}
private:
struct Hypothesis {
std::vector<int32_t> history;
std::vector<float> probabilities;
ContextNode *context_state = nullptr;
int trailing_blanks = 0;
double log_probability = 0.0;
};
struct Candidate {
std::size_t hypothesis = 0;
int32_t token = 0;
float acoustic_probability = 0.0f;
double selection_score = 0.0;
double score = 0.0;
ContextNode *context_state = nullptr;
};
static double LogAdd(double left, double right) {
const double maximum = std::max(left, right);
return maximum + std::log(std::exp(left - maximum) +
std::exp(right - maximum));
}
static std::string HistoryKey(const std::vector<int32_t> &history) {
std::string result;
for (int32_t token : history) {
if (!result.empty()) result.push_back('-');
result += std::to_string(token);
}
return result;
}
const Hypothesis &BestHypothesis() const {
if (hypotheses_.empty()) throw std::runtime_error("No active hypotheses");
return *std::max_element(
hypotheses_.begin(), hypotheses_.end(),
[](const Hypothesis &left, const Hypothesis &right) {
return left.log_probability < right.log_probability;
});
}
std::array<float, kEncoderDim> DecoderOutput(
const Hypothesis &hypothesis) {
if (hypothesis.history[hypothesis.history.size() - 2] < 0 ||
hypothesis.history.back() < 0) {
return initial_decoder_;
}
const std::pair<int32_t, int32_t> key{
hypothesis.history[hypothesis.history.size() - 2],
hypothesis.history.back()};
const auto cached = cache_.find(key);
if (cached != cache_.end()) return cached->second;
const std::array<int32_t, kContextSize> decoder_input{key.first, key.second};
if (decoder_->SetInputByName("y", decoder_input.data(),
decoder_input.size() * sizeof(int32_t)) != 0 ||
decoder_->RunSync() != 0) {
throw std::runtime_error("Decoder inference failed");
}
std::array<float, kEncoderDim> output{};
if (decoder_->GetOutputByName("decoder_out", output.data(),
output.size() * sizeof(float)) != 0) {
throw std::runtime_error("Failed to read decoder output");
}
cache_.emplace(key, output);
return output;
}
struct PairHash {
std::size_t operator()(const std::pair<int32_t, int32_t> &value) const {
return (static_cast<std::size_t>(static_cast<uint32_t>(value.first))
<< 32) ^
static_cast<uint32_t>(value.second);
}
};
EngineWrapper *decoder_;
EngineWrapper *joiner_;
ContextGraph *graph_;
std::array<float, kEncoderDim> initial_decoder_{};
int required_trailing_blanks_ = 1;
int max_active_paths_ = 1;
std::vector<Hypothesis> hypotheses_;
std::unordered_map<std::pair<int32_t, int32_t>,
std::array<float, kEncoderDim>, PairHash>
cache_;
};
void ResetEncoderStates(EngineWrapper *encoder) {
for (std::size_t i = 1; i < encoder->InputCount(); ++i) {
if (encoder->ZeroInputByName(encoder->InputName(i)) != 0) {
throw std::runtime_error("Failed to reset encoder state: " +
encoder->InputName(i));
}
}
}
void UpdateEncoderStates(EngineWrapper *encoder) {
for (std::size_t i = 1; i < encoder->InputCount(); ++i) {
const std::string &input_name = encoder->InputName(i);
if (encoder->CopyOutputToInputByName("new_" + input_name, input_name) != 0) {
throw std::runtime_error("Failed to update encoder state: " + input_name);
}
}
}
void Run(const Args &args) {
const PcmWav wav = ReadPcmWav(args.audio);
const double audio_seconds =
static_cast<double>(wav.samples.size()) / kSampleRate;
if (audio_seconds <= 0.0) {
throw std::runtime_error("Input WAV contains no samples");
}
const auto feature_begin = Clock::now();
const std::vector<float> features = ComputeFbank(wav);
const double feature_seconds = ElapsedSeconds(feature_begin, Clock::now());
const int feature_frames = static_cast<int>(features.size() / kFeatureDim);
const auto token_table = LoadTokens(args.tokens);
const auto keywords = LoadKeywords(args.keywords, token_table,
args.default_score,
args.default_threshold);
ContextGraph graph(keywords);
const auto initial_decoder = LoadInitialDecoder(args.initial_decoder);
const auto model_load_begin = Clock::now();
AxRuntime runtime;
EngineWrapper encoder;
EngineWrapper decoder;
EngineWrapper joiner;
if (encoder.Init(ModelPath(args, "encoder")) != 0 ||
decoder.Init(ModelPath(args, "decoder")) != 0 ||
joiner.Init(ModelPath(args, "joiner")) != 0) {
throw std::runtime_error("Failed to load Sherpa KWS axmodels");
}
const double model_load_seconds =
ElapsedSeconds(model_load_begin, Clock::now());
KeywordDecoder keyword_decoder(&decoder, &joiner, &graph, initial_decoder,
args.trailing_blanks, args.max_active_paths);
ResetEncoderStates(&encoder);
const int input_frames = args.chunk_size == 8 ? 29 : 45;
const int output_frames = args.chunk_size == 8 ? 4 : 8;
const int chunk_shift = args.chunk_size * 2;
if (encoder.GetInputSizeByName("x") !=
input_frames * kFeatureDim * static_cast<int>(sizeof(float))) {
throw std::runtime_error("Encoder input shape does not match chunk size");
}
std::vector<std::string> detections;
int decode_calls = 0;
const auto inference_begin = Clock::now();
for (int start = 0; start + input_frames < feature_frames;
start += chunk_shift) {
if (keyword_decoder.trailing_blanks() * 0.04f > 1.5f) {
ResetEncoderStates(&encoder);
keyword_decoder.Reset();
}
const float *input =
features.data() + static_cast<std::size_t>(start) * kFeatureDim;
if (encoder.SetInputByName(
"x", input,
input_frames * kFeatureDim * sizeof(float)) != 0 ||
encoder.RunSync() != 0) {
throw std::runtime_error("Encoder inference failed");
}
std::vector<float> encoder_output(
static_cast<std::size_t>(output_frames) * kEncoderDim);
if (encoder.GetOutputByName("encoder_out", encoder_output.data(),
encoder_output.size() * sizeof(float)) != 0) {
throw std::runtime_error("Failed to read encoder output");
}
UpdateEncoderStates(&encoder);
++decode_calls;
bool found = false;
for (int frame = 0; frame < output_frames; ++frame) {
const std::string phrase = keyword_decoder.DecodeFrame(
encoder_output.data() + frame * kEncoderDim);
if (!phrase.empty()) {
detections.push_back(phrase);
found = true;
}
}
if (found) {
ResetEncoderStates(&encoder);
keyword_decoder.Reset();
}
}
const double inference_seconds =
ElapsedSeconds(inference_begin, Clock::now());
const double processing_seconds = feature_seconds + inference_seconds;
const double rtf = processing_seconds / audio_seconds;
std::printf("\nSherpa KWS C++ inference complete\n");
std::printf(
"target: %s\naudio: %s\nchunk_size: %d\nmax_active_paths: %d\n"
"feature_frames: %d\n",
AXERA_TARGET_NAME, args.audio.c_str(), args.chunk_size,
args.max_active_paths,
feature_frames);
std::printf("decode_calls: %d\ndetections:", decode_calls);
if (detections.empty()) {
std::printf(" []\n");
} else {
std::printf("\n");
for (const std::string &phrase : detections) {
std::printf(" WAKEUP %s\n", phrase.c_str());
}
}
std::printf("audio_seconds: %.6f\n", audio_seconds);
std::printf("feature_seconds: %.6f\n", feature_seconds);
std::printf("model_load_seconds: %.6f\n", model_load_seconds);
std::printf("inference_seconds: %.6f\n", inference_seconds);
std::printf("processing_seconds: %.6f\n", processing_seconds);
std::printf("rtf: %.6f\n", rtf);
}
} // namespace
int main(int argc, char **argv) {
try {
const Args args = ParseArgs(argc, argv);
Run(args);
return 0;
} catch (const std::exception &error) {
std::fprintf(stderr, "ERROR: %s\n", error.what());
return 1;
}
}