update

Dockerfile

@@ -9,7 +9,7 @@ COPY . /code/
 
 RUN pip install --upgrade pip
 
-RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 # libk2_torch_api.so
 RUN export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.8/site-packages/k2/lib/

main.py

@@ -4,13 +4,17 @@ import argparse
 from collections import defaultdict
 from datetime import datetime
 import functools
-import io
 import logging
+import os
 from pathlib import Path
 import platform
 import time
 import tempfile
 
+os.system(
+    "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/python3.8/site-packages/k2/lib/"
+)
+
 from project_settings import project_path, log_directory
 import log
 
@@ -109,8 +113,8 @@ def process(
         nn_model_file=nn_model_file.as_posix(),
         tokens_file=tokens_file.as_posix(),
         sub_folder=m_dict["sub_folder"],
-        local_model_dir=local_model_dir,
-        recognizer_type=m_dict["recognizer_type"],
+        local_model_dir=local_model_dir.as_posix(),
+        loader=m_dict["loader"],
         decoding_method=decoding_method,
         num_active_paths=num_active_paths,
     )

toolbox/k2_sherpa/nn_models.py

@@ -6,6 +6,7 @@ import os
 
 import huggingface_hub
 import sherpa
+import sherpa_onnx
 
 
 class EnumDecodingMethod(Enum):
@@ -13,13 +14,6 @@ class EnumDecodingMethod(Enum):
     modified_beam_search = "modified_beam_search"
 
 
-class EnumRecognizerType(Enum):
-    sherpa_offline_recognizer = "sherpa.OfflineRecognizer"
-    sherpa_online_recognizer = "sherpa.OnlineRecognizer"
-    sherpa_onnx_offline_recognizer = "sherpa_onnx.OfflineRecognizer"
-    sherpa_onnx_online_recognizer = "sherpa_onnx.OnlineRecognizer"
-
-
 model_map = {
     "Chinese": [
         {
@@ -27,7 +21,14 @@ model_map = {
             "nn_model_file": "final.zip",
             "tokens_file": "units.txt",
             "sub_folder": ".",
-            "recognizer_type": EnumRecognizerType.sherpa_offline_recognizer.value,
+            "loader": "load_sherpa_offline_recognizer",
+        },
+        {
+            "repo_id": "csukuangfj/sherpa-onnx-paraformer-zh-2023-03-28",
+            "nn_model_file": "model.int8.onnx",
+            "tokens_file": "tokens.txt",
+            "sub_folder": ".",
+            "loader": "load_sherpa_offline_recognizer_from_paraformer",
         }
     ]
 }
@@ -83,12 +84,31 @@ def load_sherpa_offline_recognizer(nn_model_file: str,
     return recognizer
 
 
+def load_sherpa_offline_recognizer_from_paraformer(nn_model_file: str,
+                                                   tokens_file: str,
+                                                   sample_rate: int = 16000,
+                                                   decoding_method: str = "greedy_search",
+                                                   feature_dim: int = 80,
+                                                   num_threads: int = 2,
+                                                   ):
+    recognizer = sherpa_onnx.OfflineRecognizer.from_paraformer(
+        paraformer=nn_model_file,
+        tokens=tokens_file,
+        num_threads=num_threads,
+        sample_rate=sample_rate,
+        feature_dim=feature_dim,
+        decoding_method=decoding_method,
+        debug=False,
+    )
+    return recognizer
+
+
 def load_recognizer(repo_id: str,
                     nn_model_file: str,
                     tokens_file: str,
                     sub_folder: str,
                     local_model_dir: str,
-                    recognizer_type: str,
+                    loader: str,
                     decoding_method: str = "greedy_search",
                     num_active_paths: int = 4,
                     ):
@@ -101,15 +121,21 @@ def load_recognizer(repo_id: str,
             local_model_dir=local_model_dir,
         )
 
-    if recognizer_type == EnumRecognizerType.sherpa_offline_recognizer.value:
+    if loader == "load_sherpa_offline_recognizer":
         recognizer = load_sherpa_offline_recognizer(
             nn_model_file=nn_model_file,
             tokens_file=tokens_file,
             decoding_method=decoding_method,
             num_active_paths=num_active_paths,
         )
+    elif loader == "load_sherpa_offline_recognizer_from_paraformer":
+        recognizer = load_sherpa_offline_recognizer_from_paraformer(
+            nn_model_file=nn_model_file,
+            tokens_file=tokens_file,
+            decoding_method=decoding_method,
+        )
     else:
-        raise NotImplementedError("recognizer_type not support: {}".format(recognizer_type))
+        raise NotImplementedError("loader not support: {}".format(loader))
     return recognizer