Begin to support CTC models (#119)
Please see https://k2-fsa.github.io/sherpa/onnx/pretrained_models/offline-ctc/nemo/index.html for a list of pre-trained CTC models from NeMo.
This commit is contained in:
128
sherpa-onnx/csrc/offline-recognizer-ctc-impl.h
Normal file
128
sherpa-onnx/csrc/offline-recognizer-ctc-impl.h
Normal file
@@ -0,0 +1,128 @@
|
||||
// sherpa-onnx/csrc/offline-recognizer-ctc-impl.h
|
||||
//
|
||||
// Copyright (c) 2022-2023 Xiaomi Corporation
|
||||
|
||||
#ifndef SHERPA_ONNX_CSRC_OFFLINE_RECOGNIZER_CTC_IMPL_H_
|
||||
#define SHERPA_ONNX_CSRC_OFFLINE_RECOGNIZER_CTC_IMPL_H_
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "sherpa-onnx/csrc/offline-ctc-decoder.h"
|
||||
#include "sherpa-onnx/csrc/offline-ctc-greedy-search-decoder.h"
|
||||
#include "sherpa-onnx/csrc/offline-ctc-model.h"
|
||||
#include "sherpa-onnx/csrc/offline-recognizer-impl.h"
|
||||
#include "sherpa-onnx/csrc/pad-sequence.h"
|
||||
#include "sherpa-onnx/csrc/symbol-table.h"
|
||||
|
||||
namespace sherpa_onnx {
|
||||
|
||||
static OfflineRecognitionResult Convert(const OfflineCtcDecoderResult &src,
|
||||
const SymbolTable &sym_table) {
|
||||
OfflineRecognitionResult r;
|
||||
r.tokens.reserve(src.tokens.size());
|
||||
|
||||
std::string text;
|
||||
|
||||
for (int32_t i = 0; i != src.tokens.size(); ++i) {
|
||||
auto sym = sym_table[src.tokens[i]];
|
||||
text.append(sym);
|
||||
r.tokens.push_back(std::move(sym));
|
||||
}
|
||||
r.text = std::move(text);
|
||||
|
||||
return r;
|
||||
}
|
||||
|
||||
class OfflineRecognizerCtcImpl : public OfflineRecognizerImpl {
|
||||
public:
|
||||
explicit OfflineRecognizerCtcImpl(const OfflineRecognizerConfig &config)
|
||||
: config_(config),
|
||||
symbol_table_(config_.model_config.tokens),
|
||||
model_(OfflineCtcModel::Create(config_.model_config)) {
|
||||
config_.feat_config.nemo_normalize_type =
|
||||
model_->FeatureNormalizationMethod();
|
||||
|
||||
if (config.decoding_method == "greedy_search") {
|
||||
if (!symbol_table_.contains("<blk>")) {
|
||||
SHERPA_ONNX_LOGE(
|
||||
"We expect that tokens.txt contains "
|
||||
"the symbol <blk> and its ID.");
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
int32_t blank_id = symbol_table_["<blk>"];
|
||||
decoder_ = std::make_unique<OfflineCtcGreedySearchDecoder>(blank_id);
|
||||
} else {
|
||||
SHERPA_ONNX_LOGE("Only greedy_search is supported at present. Given %s",
|
||||
config.decoding_method.c_str());
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<OfflineStream> CreateStream() const override {
|
||||
return std::make_unique<OfflineStream>(config_.feat_config);
|
||||
}
|
||||
|
||||
void DecodeStreams(OfflineStream **ss, int32_t n) const override {
|
||||
auto memory_info =
|
||||
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
|
||||
|
||||
int32_t feat_dim = config_.feat_config.feature_dim;
|
||||
|
||||
std::vector<Ort::Value> features;
|
||||
features.reserve(n);
|
||||
|
||||
std::vector<std::vector<float>> features_vec(n);
|
||||
std::vector<int64_t> features_length_vec(n);
|
||||
|
||||
for (int32_t i = 0; i != n; ++i) {
|
||||
std::vector<float> f = ss[i]->GetFrames();
|
||||
|
||||
int32_t num_frames = f.size() / feat_dim;
|
||||
features_vec[i] = std::move(f);
|
||||
|
||||
features_length_vec[i] = num_frames;
|
||||
|
||||
std::array<int64_t, 2> shape = {num_frames, feat_dim};
|
||||
|
||||
Ort::Value x = Ort::Value::CreateTensor(
|
||||
memory_info, features_vec[i].data(), features_vec[i].size(),
|
||||
shape.data(), shape.size());
|
||||
features.push_back(std::move(x));
|
||||
} // for (int32_t i = 0; i != n; ++i)
|
||||
|
||||
std::vector<const Ort::Value *> features_pointer(n);
|
||||
for (int32_t i = 0; i != n; ++i) {
|
||||
features_pointer[i] = &features[i];
|
||||
}
|
||||
|
||||
std::array<int64_t, 1> features_length_shape = {n};
|
||||
Ort::Value x_length = Ort::Value::CreateTensor(
|
||||
memory_info, features_length_vec.data(), n,
|
||||
features_length_shape.data(), features_length_shape.size());
|
||||
|
||||
Ort::Value x = PadSequence(model_->Allocator(), features_pointer,
|
||||
-23.025850929940457f);
|
||||
auto t = model_->Forward(std::move(x), std::move(x_length));
|
||||
|
||||
auto results = decoder_->Decode(std::move(t.first), std::move(t.second));
|
||||
|
||||
for (int32_t i = 0; i != n; ++i) {
|
||||
auto r = Convert(results[i], symbol_table_);
|
||||
ss[i]->SetResult(r);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
OfflineRecognizerConfig config_;
|
||||
SymbolTable symbol_table_;
|
||||
std::unique_ptr<OfflineCtcModel> model_;
|
||||
std::unique_ptr<OfflineCtcDecoder> decoder_;
|
||||
};
|
||||
|
||||
} // namespace sherpa_onnx
|
||||
|
||||
#endif // SHERPA_ONNX_CSRC_OFFLINE_RECOGNIZER_CTC_IMPL_H_
|
||||
Reference in New Issue
Block a user