Add C++ support for streaming NeMo CTC models. (#857)

This commit is contained in:
Fangjun Kuang
2024-05-10 16:26:43 +08:00
committed by GitHub
parent 1eb60e8711
commit 46e4e5b7ac
22 changed files with 782 additions and 41 deletions

View File

@@ -23,6 +23,7 @@ set(srcs
online-ctc-fst-decoder-config.cc
online-lm-config.cc
online-model-config.cc
online-nemo-ctc-model-config.cc
online-paraformer-model-config.cc
online-recognizer.cc
online-stream.cc

View File

@@ -9,6 +9,7 @@
#include "sherpa-onnx/csrc/online-model-config.h"
#include "sherpa-onnx/csrc/online-transducer-model-config.h"
#include "sherpa-onnx/python/csrc/online-nemo-ctc-model-config.h"
#include "sherpa-onnx/python/csrc/online-paraformer-model-config.h"
#include "sherpa-onnx/python/csrc/online-transducer-model-config.h"
#include "sherpa-onnx/python/csrc/online-wenet-ctc-model-config.h"
@@ -21,26 +22,30 @@ void PybindOnlineModelConfig(py::module *m) {
PybindOnlineParaformerModelConfig(m);
PybindOnlineWenetCtcModelConfig(m);
PybindOnlineZipformer2CtcModelConfig(m);
PybindOnlineNeMoCtcModelConfig(m);
using PyClass = OnlineModelConfig;
py::class_<PyClass>(*m, "OnlineModelConfig")
.def(py::init<const OnlineTransducerModelConfig &,
const OnlineParaformerModelConfig &,
const OnlineWenetCtcModelConfig &,
const OnlineZipformer2CtcModelConfig &, const std::string &,
const OnlineZipformer2CtcModelConfig &,
const OnlineNeMoCtcModelConfig &, const std::string &,
int32_t, int32_t, bool, const std::string &,
const std::string &>(),
py::arg("transducer") = OnlineTransducerModelConfig(),
py::arg("paraformer") = OnlineParaformerModelConfig(),
py::arg("wenet_ctc") = OnlineWenetCtcModelConfig(),
py::arg("zipformer2_ctc") = OnlineZipformer2CtcModelConfig(),
py::arg("tokens"), py::arg("num_threads"), py::arg("warm_up") = 0,
py::arg("nemo_ctc") = OnlineNeMoCtcModelConfig(), py::arg("tokens"),
py::arg("num_threads"), py::arg("warm_up") = 0,
py::arg("debug") = false, py::arg("provider") = "cpu",
py::arg("model_type") = "")
.def_readwrite("transducer", &PyClass::transducer)
.def_readwrite("paraformer", &PyClass::paraformer)
.def_readwrite("wenet_ctc", &PyClass::wenet_ctc)
.def_readwrite("zipformer2_ctc", &PyClass::zipformer2_ctc)
.def_readwrite("nemo_ctc", &PyClass::nemo_ctc)
.def_readwrite("tokens", &PyClass::tokens)
.def_readwrite("num_threads", &PyClass::num_threads)
.def_readwrite("debug", &PyClass::debug)

View File

@@ -0,0 +1,22 @@
// sherpa-onnx/python/csrc/online-nemo-ctc-model-config.cc
//
// Copyright (c) 2024 Xiaomi Corporation
#include "sherpa-onnx/python/csrc/online-nemo-ctc-model-config.h"
#include <string>
#include <vector>
#include "sherpa-onnx/csrc/online-nemo-ctc-model-config.h"
namespace sherpa_onnx {
void PybindOnlineNeMoCtcModelConfig(py::module *m) {
using PyClass = OnlineNeMoCtcModelConfig;
py::class_<PyClass>(*m, "OnlineNeMoCtcModelConfig")
.def(py::init<const std::string &>(), py::arg("model"))
.def_readwrite("model", &PyClass::model)
.def("__str__", &PyClass::ToString);
}
} // namespace sherpa_onnx

View File

@@ -0,0 +1,16 @@
// sherpa-onnx/python/csrc/online-nemo-ctc-model-config.h
//
// Copyright (c) 2024 Xiaomi Corporation
#ifndef SHERPA_ONNX_PYTHON_CSRC_ONLINE_NEMO_CTC_MODEL_CONFIG_H_
#define SHERPA_ONNX_PYTHON_CSRC_ONLINE_NEMO_CTC_MODEL_CONFIG_H_
#include "sherpa-onnx/python/csrc/sherpa-onnx.h"
namespace sherpa_onnx {
void PybindOnlineNeMoCtcModelConfig(py::module *m);
}
#endif // SHERPA_ONNX_PYTHON_CSRC_ONLINE_NEMO_CTC_MODEL_CONFIG_H_

View File

@@ -42,6 +42,8 @@ static void PybindOnlineRecognizerResult(py::module *m) {
"segment", [](PyClass &self) -> int32_t { return self.segment; })
.def_property_readonly(
"is_final", [](PyClass &self) -> bool { return self.is_final; })
.def("__str__", &PyClass::AsJsonString,
py::call_guard<py::gil_scoped_release>())
.def("as_json_string", &PyClass::AsJsonString,
py::call_guard<py::gil_scoped_release>());
}
@@ -50,29 +52,17 @@ static void PybindOnlineRecognizerConfig(py::module *m) {
using PyClass = OnlineRecognizerConfig;
py::class_<PyClass>(*m, "OnlineRecognizerConfig")
.def(
py::init<const FeatureExtractorConfig &,
const OnlineModelConfig &,
const OnlineLMConfig &,
const EndpointConfig &,
const OnlineCtcFstDecoderConfig &,
bool,
const std::string &,
int32_t,
const std::string &,
float,
float,
float>(),
py::arg("feat_config"),
py::arg("model_config"),
py::init<const FeatureExtractorConfig &, const OnlineModelConfig &,
const OnlineLMConfig &, const EndpointConfig &,
const OnlineCtcFstDecoderConfig &, bool, const std::string &,
int32_t, const std::string &, float, float, float>(),
py::arg("feat_config"), py::arg("model_config"),
py::arg("lm_config") = OnlineLMConfig(),
py::arg("endpoint_config") = EndpointConfig(),
py::arg("ctc_fst_decoder_config") = OnlineCtcFstDecoderConfig(),
py::arg("enable_endpoint"),
py::arg("decoding_method"),
py::arg("max_active_paths") = 4,
py::arg("hotwords_file") = "",
py::arg("hotwords_score") = 0,
py::arg("blank_penalty") = 0.0,
py::arg("enable_endpoint"), py::arg("decoding_method"),
py::arg("max_active_paths") = 4, py::arg("hotwords_file") = "",
py::arg("hotwords_score") = 0, py::arg("blank_penalty") = 0.0,
py::arg("temperature_scale") = 2.0)
.def_readwrite("feat_config", &PyClass::feat_config)
.def_readwrite("model_config", &PyClass::model_config)