Add Python API (#31)

sherpa-onnx/python/CMakeLists.txt

@@ -0,0 +1,5 @@
+add_subdirectory(csrc)
+
+if(SHERPA_ONNX_ENABLE_TESTS)
+  add_subdirectory(tests)
+endif()

sherpa-onnx/python/csrc/CMakeLists.txt

@@ -0,0 +1,29 @@
+include_directories(${CMAKE_SOURCE_DIR})
+
+pybind11_add_module(_sherpa_onnx
+  features.cc
+  online-transducer-model-config.cc
+  sherpa-onnx.cc
+  online-stream.cc
+  online-recognizer.cc
+)
+
+if(APPLE)
+  execute_process(
+    COMMAND "${PYTHON_EXECUTABLE}" -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())"
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE PYTHON_SITE_PACKAGE_DIR
+  )
+  message(STATUS "PYTHON_SITE_PACKAGE_DIR: ${PYTHON_SITE_PACKAGE_DIR}")
+  target_link_libraries(_sherpa_onnx PRIVATE "-Wl,-rpath,${PYTHON_SITE_PACKAGE_DIR}")
+endif()
+
+if(NOT WIN32)
+  target_link_libraries(_sherpa_onnx PRIVATE "-Wl,-rpath,${SHERPA_ONNX_RPATH_ORIGIN}/sherpa_onnx/lib")
+endif()
+
+target_link_libraries(_sherpa_onnx PRIVATE sherpa-onnx-core)
+
+install(TARGETS _sherpa_onnx
+  DESTINATION ../
+)

sherpa-onnx/python/csrc/features.cc

@@ -0,0 +1,23 @@
+// sherpa-onnx/python/csrc/features.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/python/csrc/features.h"
+
+#include "sherpa-onnx/csrc/features.h"
+
+namespace sherpa_onnx {
+
+static void PybindFeatureExtractorConfig(py::module *m) {
+  using PyClass = FeatureExtractorConfig;
+  py::class_<PyClass>(*m, "FeatureExtractorConfig")
+      .def(py::init<float, int32_t>(), py::arg("sampling_rate") = 16000,
+           py::arg("feature_dim") = 80)
+      .def_readwrite("sampling_rate", &PyClass::sampling_rate)
+      .def_readwrite("feature_dim", &PyClass::feature_dim)
+      .def("__str__", &PyClass::ToString);
+}
+
+void PybindFeatures(py::module *m) { PybindFeatureExtractorConfig(m); }
+
+}  // namespace sherpa_onnx

sherpa-onnx/python/csrc/features.h

@@ -0,0 +1,16 @@
+// sherpa-onnx/python/csrc/features.h
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#ifndef SHERPA_ONNX_PYTHON_CSRC_FEATURES_H_
+#define SHERPA_ONNX_PYTHON_CSRC_FEATURES_H_
+
+#include "sherpa-onnx/python/csrc/sherpa-onnx.h"
+
+namespace sherpa_onnx {
+
+void PybindFeatures(py::module *m);
+
+}
+
+#endif  // SHERPA_ONNX_PYTHON_CSRC_FEATURES_H_

sherpa-onnx/python/csrc/online-recognizer.cc

@@ -0,0 +1,49 @@
+// sherpa-onnx/python/csrc/online-recongizer.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/python/csrc/online-recognizer.h"
+
+#include <string>
+#include <vector>
+
+#include "sherpa-onnx/csrc/online-recognizer.h"
+
+namespace sherpa_onnx {
+
+static void PybindOnlineRecognizerResult(py::module *m) {
+  using PyClass = OnlineRecognizerResult;
+  py::class_<PyClass>(*m, "OnlineRecognizerResult")
+      .def_property_readonly("text", [](PyClass &self) { return self.text; });
+}
+
+static void PybindOnlineRecognizerConfig(py::module *m) {
+  using PyClass = OnlineRecognizerConfig;
+  py::class_<PyClass>(*m, "OnlineRecognizerConfig")
+      .def(py::init<const FeatureExtractorConfig &,
+                    const OnlineTransducerModelConfig &, const std::string &>(),
+           py::arg("feat_config"), py::arg("model_config"), py::arg("tokens"))
+      .def_readwrite("feat_config", &PyClass::feat_config)
+      .def_readwrite("model_config", &PyClass::model_config)
+      .def_readwrite("tokens", &PyClass::tokens)
+      .def("__str__", &PyClass::ToString);
+}
+
+void PybindOnlineRecognizer(py::module *m) {
+  PybindOnlineRecognizerResult(m);
+  PybindOnlineRecognizerConfig(m);
+
+  using PyClass = OnlineRecognizer;
+  py::class_<PyClass>(*m, "OnlineRecognizer")
+      .def(py::init<const OnlineRecognizerConfig &>(), py::arg("config"))
+      .def("create_stream", &PyClass::CreateStream)
+      .def("is_ready", &PyClass::IsReady)
+      .def("decode_stream", &PyClass::DecodeStream)
+      .def("decode_streams",
+           [](PyClass &self, std::vector<OnlineStream *> ss) {
+             self.DecodeStreams(ss.data(), ss.size());
+           })
+      .def("get_result", &PyClass::GetResult);
+}
+
+}  // namespace sherpa_onnx

sherpa-onnx/python/csrc/online-recognizer.h

@@ -0,0 +1,16 @@
+// sherpa-onnx/python/csrc/online-recongizer.h
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#ifndef SHERPA_ONNX_PYTHON_CSRC_ONLINE_RECOGNIZER_H_
+#define SHERPA_ONNX_PYTHON_CSRC_ONLINE_RECOGNIZER_H_
+
+#include "sherpa-onnx/python/csrc/sherpa-onnx.h"
+
+namespace sherpa_onnx {
+
+void PybindOnlineRecognizer(py::module *m);
+
+}
+
+#endif  // SHERPA_ONNX_PYTHON_CSRC_ONLINE_RECOGNIZER_H_

sherpa-onnx/python/csrc/online-stream.cc

@@ -0,0 +1,21 @@
+// sherpa-onnx/python/csrc/online-stream.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/python/csrc/online-stream.h"
+
+#include "sherpa-onnx/csrc/online-stream.h"
+
+namespace sherpa_onnx {
+
+void PybindOnlineStream(py::module *m) {
+  using PyClass = OnlineStream;
+  py::class_<PyClass>(*m, "OnlineStream")
+      .def("accept_waveform",
+           [](PyClass &self, float sample_rate, py::array_t<float> waveform) {
+             self.AcceptWaveform(sample_rate, waveform.data(), waveform.size());
+           })
+      .def("input_finished", &PyClass::InputFinished);
+}
+
+}  // namespace sherpa_onnx

sherpa-onnx/python/csrc/online-stream.h

@@ -0,0 +1,16 @@
+// sherpa-onnx/python/csrc/online-stream.h
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#ifndef SHERPA_ONNX_PYTHON_CSRC_ONLINE_STREAM_H_
+#define SHERPA_ONNX_PYTHON_CSRC_ONLINE_STREAM_H_
+
+#include "sherpa-onnx/python/csrc/sherpa-onnx.h"
+
+namespace sherpa_onnx {
+
+void PybindOnlineStream(py::module *m);
+
+}
+
+#endif  // SHERPA_ONNX_PYTHON_CSRC_ONLINE_STREAM_H_

sherpa-onnx/python/csrc/online-transducer-model-config.cc

@@ -0,0 +1,29 @@
+// sherpa-onnx/python/csrc/online-transducer-model-config.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/csrc/online-transducer-model-config.h"
+
+#include <string>
+
+#include "sherpa-onnx/python/csrc/online-transducer-model-config.h"
+
+namespace sherpa_onnx {
+
+void PybindOnlineTransducerModelConfig(py::module *m) {
+  using PyClass = OnlineTransducerModelConfig;
+  py::class_<PyClass>(*m, "OnlineTransducerModelConfig")
+      .def(py::init<const std::string &, const std::string &,
+                    const std::string &, int32_t, bool>(),
+           py::arg("encoder_filename"), py::arg("decoder_filename"),
+           py::arg("joiner_filename"), py::arg("num_threads"),
+           py::arg("debug") = false)
+      .def_readwrite("encoder_filename", &PyClass::encoder_filename)
+      .def_readwrite("decoder_filename", &PyClass::decoder_filename)
+      .def_readwrite("joiner_filename", &PyClass::joiner_filename)
+      .def_readwrite("num_threads", &PyClass::num_threads)
+      .def_readwrite("debug", &PyClass::debug)
+      .def("__str__", &PyClass::ToString);
+}
+
+}  // namespace sherpa_onnx

sherpa-onnx/python/csrc/online-transducer-model-config.h

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

sherpa-onnx/python/csrc/sherpa-onnx.cc

@@ -0,0 +1,22 @@
+// sherpa-onnx/python/csrc/sherpa-onnx.cc
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#include "sherpa-onnx/python/csrc/sherpa-onnx.h"
+
+#include "sherpa-onnx/python/csrc/features.h"
+#include "sherpa-onnx/python/csrc/online-recognizer.h"
+#include "sherpa-onnx/python/csrc/online-stream.h"
+#include "sherpa-onnx/python/csrc/online-transducer-model-config.h"
+
+namespace sherpa_onnx {
+
+PYBIND11_MODULE(_sherpa_onnx, m) {
+  m.doc() = "pybind11 binding of sherpa-onnx";
+  PybindFeatures(&m);
+  PybindOnlineTransducerModelConfig(&m);
+  PybindOnlineStream(&m);
+  PybindOnlineRecognizer(&m);
+}
+
+}  // namespace sherpa_onnx

sherpa-onnx/python/csrc/sherpa-onnx.h

@@ -0,0 +1,14 @@
+// sherpa-onnx/python/csrc/sherpa-onnx.h
+//
+// Copyright (c)  2023  Xiaomi Corporation
+
+#ifndef SHERPA_ONNX_PYTHON_CSRC_SHERPA_ONNX_H_
+#define SHERPA_ONNX_PYTHON_CSRC_SHERPA_ONNX_H_
+
+#include "pybind11/numpy.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/stl.h"
+
+namespace py = pybind11;
+
+#endif  // SHERPA_ONNX_PYTHON_CSRC_SHERPA_ONNX_H_

sherpa-onnx/python/sherpa_onnx/__init__.py

@@ -0,0 +1,8 @@
+from _sherpa_onnx import (
+    FeatureExtractorConfig,
+    OnlineRecognizerConfig,
+    OnlineStream,
+    OnlineTransducerModelConfig,
+)
+
+from .online_recognizer import OnlineRecognizer

sherpa-onnx/python/sherpa_onnx/online_recognizer.py

@@ -0,0 +1,96 @@
+from pathlib import Path
+from typing import List
+
+from _sherpa_onnx import (
+    OnlineStream,
+    OnlineTransducerModelConfig,
+    FeatureExtractorConfig,
+    OnlineRecognizerConfig,
+)
+from _sherpa_onnx import OnlineRecognizer as _Recognizer
+
+
+def _assert_file_exists(f: str):
+    assert Path(f).is_file(), f"{f} does not exist"
+
+
+class OnlineRecognizer(object):
+    """A class for streaming speech recognition."""
+
+    def __init__(
+        self,
+        tokens: str,
+        encoder: str,
+        decoder: str,
+        joiner: str,
+        num_threads: int = 4,
+        sample_rate: float = 16000,
+        feature_dim: int = 80,
+    ):
+        """
+        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.
+
+        Args:
+          tokens:
+            Path to ``tokens.txt``. Each line in ``tokens.txt`` contains two
+            columns::
+
+                symbol integer_id
+
+          encoder:
+            Path to ``encoder.onnx``.
+          decoder:
+            Path to ``decoder.onnx``.
+          joiner:
+            Path to ``joiner.onnx``.
+          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.
+        """
+        _assert_file_exists(tokens)
+        _assert_file_exists(encoder)
+        _assert_file_exists(decoder)
+        _assert_file_exists(joiner)
+
+        assert num_threads > 0, num_threads
+
+        model_config = OnlineTransducerModelConfig(
+            encoder_filename=encoder,
+            decoder_filename=decoder,
+            joiner_filename=joiner,
+            num_threads=num_threads,
+        )
+
+        feat_config = FeatureExtractorConfig(
+            sampling_rate=sample_rate,
+            feature_dim=feature_dim,
+        )
+
+        recognizer_config = OnlineRecognizerConfig(
+            feat_config=feat_config,
+            model_config=model_config,
+            tokens=tokens,
+        )
+
+        self.recognizer = _Recognizer(recognizer_config)
+
+    def create_stream(self):
+        return self.recognizer.create_stream()
+
+    def decode_stream(self, s: OnlineStream):
+        self.recognizer.decode_stream(s)
+
+    def decode_streams(self, ss: List[OnlineStream]):
+        self.recognizer.decode_streams(ss)
+
+    def is_ready(self, s: OnlineStream) -> bool:
+        return self.recognizer.is_ready(s)
+
+    def get_result(self, s: OnlineStream) -> str:
+        return self.recognizer.get_result(s).text

sherpa-onnx/python/tests/CMakeLists.txt

@@ -0,0 +1,27 @@
+function(sherpa_onnx_add_py_test source)
+  get_filename_component(name ${source} NAME_WE)
+  set(name "${name}_py")
+
+  add_test(NAME ${name}
+    COMMAND
+      "${PYTHON_EXECUTABLE}"
+      "${CMAKE_CURRENT_SOURCE_DIR}/${source}"
+  )
+
+  get_filename_component(sherpa_onnx_path ${CMAKE_CURRENT_LIST_DIR} DIRECTORY)
+
+  set_property(TEST ${name}
+    PROPERTY ENVIRONMENT "PYTHONPATH=${sherpa_path}:$<TARGET_FILE_DIR:_sherpa_onnx>:$ENV{PYTHONPATH}"
+  )
+endfunction()
+
+# please sort the files in alphabetic order
+set(py_test_files
+  test_feature_extractor_config.py
+  test_online_transducer_model_config.py
+)
+
+foreach(source IN LISTS py_test_files)
+  sherpa_onnx_add_py_test(${source})
+endforeach()
+

sherpa-onnx/python/tests/test_feature_extractor_config.py

@@ -0,0 +1,29 @@
+# sherpa-onnx/python/tests/test_feature_extractor_config.py
+#
+# Copyright (c)  2023  Xiaomi Corporation
+#
+# To run this single test, use
+#
+#  ctest --verbose -R  test_feature_extractor_config_py
+
+import unittest
+
+import sherpa_onnx
+
+
+class TestFeatureExtractorConfig(unittest.TestCase):
+    def test_default_constructor(self):
+        config = sherpa_onnx.FeatureExtractorConfig()
+        assert config.sampling_rate == 16000, config.sampling_rate
+        assert config.feature_dim == 80, config.feature_dim
+        print(config)
+
+    def test_constructor(self):
+        config = sherpa_onnx.FeatureExtractorConfig(sampling_rate=8000, feature_dim=40)
+        assert config.sampling_rate == 8000, config.sampling_rate
+        assert config.feature_dim == 40, config.feature_dim
+        print(config)
+
+
+if __name__ == "__main__":
+    unittest.main()

sherpa-onnx/python/tests/test_online_transducer_model_config.py

@@ -0,0 +1,32 @@
+# sherpa-onnx/python/tests/test_online_transducer_model_config.py
+#
+# Copyright (c)  2023  Xiaomi Corporation
+#
+# To run this single test, use
+#
+#  ctest --verbose -R  test_online_transducer_model_config_py
+
+import unittest
+
+import sherpa_onnx
+
+
+class TestOnlineTransducerModelConfig(unittest.TestCase):
+    def test_constructor(self):
+        config = sherpa_onnx.OnlineTransducerModelConfig(
+            encoder_filename="encoder.onnx",
+            decoder_filename="decoder.onnx",
+            joiner_filename="joiner.onnx",
+            num_threads=8,
+            debug=True,
+        )
+        assert config.encoder_filename == "encoder.onnx", config.encoder_filename
+        assert config.decoder_filename == "decoder.onnx", config.decoder_filename
+        assert config.joiner_filename == "joiner.onnx", config.joiner_filename
+        assert config.num_threads == 8, config.num_threads
+        assert config.debug is True, config.debug
+        print(config)
+
+
+if __name__ == "__main__":
+    unittest.main()