| #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); |
| } |
| } |
|
|
| 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; |
| } |
| } |
|
|