Support TDNN models from the yesno recipe from icefall (#262)
This commit is contained in:
@@ -10,6 +10,7 @@ pybind11_add_module(_sherpa_onnx
|
||||
offline-paraformer-model-config.cc
|
||||
offline-recognizer.cc
|
||||
offline-stream.cc
|
||||
offline-tdnn-model-config.cc
|
||||
offline-transducer-model-config.cc
|
||||
offline-whisper-model-config.cc
|
||||
online-lm-config.cc
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
#include "sherpa-onnx/csrc/offline-model-config.h"
|
||||
#include "sherpa-onnx/python/csrc/offline-nemo-enc-dec-ctc-model-config.h"
|
||||
#include "sherpa-onnx/python/csrc/offline-paraformer-model-config.h"
|
||||
#include "sherpa-onnx/python/csrc/offline-tdnn-model-config.h"
|
||||
#include "sherpa-onnx/python/csrc/offline-transducer-model-config.h"
|
||||
#include "sherpa-onnx/python/csrc/offline-whisper-model-config.h"
|
||||
|
||||
@@ -20,24 +21,28 @@ void PybindOfflineModelConfig(py::module *m) {
|
||||
PybindOfflineParaformerModelConfig(m);
|
||||
PybindOfflineNemoEncDecCtcModelConfig(m);
|
||||
PybindOfflineWhisperModelConfig(m);
|
||||
PybindOfflineTdnnModelConfig(m);
|
||||
|
||||
using PyClass = OfflineModelConfig;
|
||||
py::class_<PyClass>(*m, "OfflineModelConfig")
|
||||
.def(py::init<const OfflineTransducerModelConfig &,
|
||||
const OfflineParaformerModelConfig &,
|
||||
const OfflineNemoEncDecCtcModelConfig &,
|
||||
const OfflineWhisperModelConfig &, const std::string &,
|
||||
const OfflineWhisperModelConfig &,
|
||||
const OfflineTdnnModelConfig &, const std::string &,
|
||||
int32_t, bool, const std::string &, const std::string &>(),
|
||||
py::arg("transducer") = OfflineTransducerModelConfig(),
|
||||
py::arg("paraformer") = OfflineParaformerModelConfig(),
|
||||
py::arg("nemo_ctc") = OfflineNemoEncDecCtcModelConfig(),
|
||||
py::arg("whisper") = OfflineWhisperModelConfig(), py::arg("tokens"),
|
||||
py::arg("whisper") = OfflineWhisperModelConfig(),
|
||||
py::arg("tdnn") = OfflineTdnnModelConfig(), py::arg("tokens"),
|
||||
py::arg("num_threads"), py::arg("debug") = false,
|
||||
py::arg("provider") = "cpu", py::arg("model_type") = "")
|
||||
.def_readwrite("transducer", &PyClass::transducer)
|
||||
.def_readwrite("paraformer", &PyClass::paraformer)
|
||||
.def_readwrite("nemo_ctc", &PyClass::nemo_ctc)
|
||||
.def_readwrite("whisper", &PyClass::whisper)
|
||||
.def_readwrite("tdnn", &PyClass::tdnn)
|
||||
.def_readwrite("tokens", &PyClass::tokens)
|
||||
.def_readwrite("num_threads", &PyClass::num_threads)
|
||||
.def_readwrite("debug", &PyClass::debug)
|
||||
|
||||
22
sherpa-onnx/python/csrc/offline-tdnn-model-config.cc
Normal file
22
sherpa-onnx/python/csrc/offline-tdnn-model-config.cc
Normal file
@@ -0,0 +1,22 @@
|
||||
// sherpa-onnx/python/csrc/offline-tdnn-model-config.cc
|
||||
//
|
||||
// Copyright (c) 2023 Xiaomi Corporation
|
||||
|
||||
#include "sherpa-onnx/csrc/offline-tdnn-model-config.h"
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "sherpa-onnx/python/csrc/offline-tdnn-model-config.h"
|
||||
|
||||
namespace sherpa_onnx {
|
||||
|
||||
void PybindOfflineTdnnModelConfig(py::module *m) {
|
||||
using PyClass = OfflineTdnnModelConfig;
|
||||
py::class_<PyClass>(*m, "OfflineTdnnModelConfig")
|
||||
.def(py::init<const std::string &>(), py::arg("model"))
|
||||
.def_readwrite("model", &PyClass::model)
|
||||
.def("__str__", &PyClass::ToString);
|
||||
}
|
||||
|
||||
} // namespace sherpa_onnx
|
||||
16
sherpa-onnx/python/csrc/offline-tdnn-model-config.h
Normal file
16
sherpa-onnx/python/csrc/offline-tdnn-model-config.h
Normal file
@@ -0,0 +1,16 @@
|
||||
// sherpa-onnx/python/csrc/offline-tdnn-model-config.h
|
||||
//
|
||||
// Copyright (c) 2023 Xiaomi Corporation
|
||||
|
||||
#ifndef SHERPA_ONNX_PYTHON_CSRC_OFFLINE_TDNN_MODEL_CONFIG_H_
|
||||
#define SHERPA_ONNX_PYTHON_CSRC_OFFLINE_TDNN_MODEL_CONFIG_H_
|
||||
|
||||
#include "sherpa-onnx/python/csrc/sherpa-onnx.h"
|
||||
|
||||
namespace sherpa_onnx {
|
||||
|
||||
void PybindOfflineTdnnModelConfig(py::module *m);
|
||||
|
||||
}
|
||||
|
||||
#endif // SHERPA_ONNX_PYTHON_CSRC_OFFLINE_TDNN_MODEL_CONFIG_H_
|
||||
@@ -8,6 +8,7 @@ from _sherpa_onnx import (
|
||||
OfflineModelConfig,
|
||||
OfflineNemoEncDecCtcModelConfig,
|
||||
OfflineParaformerModelConfig,
|
||||
OfflineTdnnModelConfig,
|
||||
OfflineWhisperModelConfig,
|
||||
)
|
||||
from _sherpa_onnx import OfflineRecognizer as _Recognizer
|
||||
@@ -37,7 +38,7 @@ class OfflineRecognizer(object):
|
||||
decoder: str,
|
||||
joiner: str,
|
||||
tokens: str,
|
||||
num_threads: int,
|
||||
num_threads: int = 1,
|
||||
sample_rate: int = 16000,
|
||||
feature_dim: int = 80,
|
||||
decoding_method: str = "greedy_search",
|
||||
@@ -48,7 +49,7 @@ class OfflineRecognizer(object):
|
||||
):
|
||||
"""
|
||||
Please refer to
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/index.html>`_
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/offline-transducer/index.html>`_
|
||||
to download pre-trained models for different languages, e.g., Chinese,
|
||||
English, etc.
|
||||
|
||||
@@ -115,7 +116,7 @@ class OfflineRecognizer(object):
|
||||
cls,
|
||||
paraformer: str,
|
||||
tokens: str,
|
||||
num_threads: int,
|
||||
num_threads: int = 1,
|
||||
sample_rate: int = 16000,
|
||||
feature_dim: int = 80,
|
||||
decoding_method: str = "greedy_search",
|
||||
@@ -124,9 +125,8 @@ class OfflineRecognizer(object):
|
||||
):
|
||||
"""
|
||||
Please refer to
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/index.html>`_
|
||||
to download pre-trained models for different languages, e.g., Chinese,
|
||||
English, etc.
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/offline-paraformer/index.html>`_
|
||||
to download pre-trained models.
|
||||
|
||||
Args:
|
||||
tokens:
|
||||
@@ -179,7 +179,7 @@ class OfflineRecognizer(object):
|
||||
cls,
|
||||
model: str,
|
||||
tokens: str,
|
||||
num_threads: int,
|
||||
num_threads: int = 1,
|
||||
sample_rate: int = 16000,
|
||||
feature_dim: int = 80,
|
||||
decoding_method: str = "greedy_search",
|
||||
@@ -188,7 +188,7 @@ class OfflineRecognizer(object):
|
||||
):
|
||||
"""
|
||||
Please refer to
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/index.html>`_
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/offline-ctc/nemo/index.html>`_
|
||||
to download pre-trained models for different languages, e.g., Chinese,
|
||||
English, etc.
|
||||
|
||||
@@ -244,14 +244,14 @@ class OfflineRecognizer(object):
|
||||
encoder: str,
|
||||
decoder: str,
|
||||
tokens: str,
|
||||
num_threads: int,
|
||||
num_threads: int = 1,
|
||||
decoding_method: str = "greedy_search",
|
||||
debug: bool = False,
|
||||
provider: str = "cpu",
|
||||
):
|
||||
"""
|
||||
Please refer to
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/index.html>`_
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/whisper/index.html>`_
|
||||
to download pre-trained models for different kinds of whisper models,
|
||||
e.g., tiny, tiny.en, base, base.en, etc.
|
||||
|
||||
@@ -301,6 +301,69 @@ class OfflineRecognizer(object):
|
||||
self.config = recognizer_config
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def from_tdnn_ctc(
|
||||
cls,
|
||||
model: str,
|
||||
tokens: str,
|
||||
num_threads: int = 1,
|
||||
sample_rate: int = 8000,
|
||||
feature_dim: int = 23,
|
||||
decoding_method: str = "greedy_search",
|
||||
debug: bool = False,
|
||||
provider: str = "cpu",
|
||||
):
|
||||
"""
|
||||
Please refer to
|
||||
`<https://k2-fsa.github.io/sherpa/onnx/pretrained_models/offline-ctc/yesno/index.html>`_
|
||||
to download pre-trained models.
|
||||
|
||||
Args:
|
||||
model:
|
||||
Path to ``model.onnx``.
|
||||
tokens:
|
||||
Path to ``tokens.txt``. Each line in ``tokens.txt`` contains two
|
||||
columns::
|
||||
|
||||
symbol integer_id
|
||||
|
||||
num_threads:
|
||||
Number of threads for neural network computation.
|
||||
sample_rate:
|
||||
Sample rate of the training data used to train the model.
|
||||
feature_dim:
|
||||
Dimension of the feature used to train the model.
|
||||
decoding_method:
|
||||
Valid values are greedy_search.
|
||||
debug:
|
||||
True to show debug messages.
|
||||
provider:
|
||||
onnxruntime execution providers. Valid values are: cpu, cuda, coreml.
|
||||
"""
|
||||
self = cls.__new__(cls)
|
||||
model_config = OfflineModelConfig(
|
||||
tdnn=OfflineTdnnModelConfig(model=model),
|
||||
tokens=tokens,
|
||||
num_threads=num_threads,
|
||||
debug=debug,
|
||||
provider=provider,
|
||||
model_type="tdnn",
|
||||
)
|
||||
|
||||
feat_config = OfflineFeatureExtractorConfig(
|
||||
sampling_rate=sample_rate,
|
||||
feature_dim=feature_dim,
|
||||
)
|
||||
|
||||
recognizer_config = OfflineRecognizerConfig(
|
||||
feat_config=feat_config,
|
||||
model_config=model_config,
|
||||
decoding_method=decoding_method,
|
||||
)
|
||||
self.recognizer = _Recognizer(recognizer_config)
|
||||
self.config = recognizer_config
|
||||
return self
|
||||
|
||||
def create_stream(self, contexts_list: Optional[List[List[int]]] = None):
|
||||
if contexts_list is None:
|
||||
return self.recognizer.create_stream()
|
||||
|
||||
Reference in New Issue
Block a user