update

main.py

@@ -105,8 +105,9 @@ def process(
     recognizer = nn_models.load_recognizer(
         repo_id=m_dict["repo_id"],
         nn_model_file=m_dict["nn_model_file"],
+        nn_model_file_sub_folder=m_dict["nn_model_file_sub_folder"],
         tokens_file=m_dict["tokens_file"],
-        sub_folder=m_dict["sub_folder"],
+        tokens_file_sub_folder=m_dict["tokens_file_sub_folder"],
         local_model_dir=local_model_dir,
         loader=m_dict["loader"],
         decoding_method=decoding_method,

toolbox/k2_sherpa/nn_models.py

@@ -20,30 +20,34 @@ model_map = {
         {
             "repo_id": "csukuangfj/wenet-chinese-model",
             "nn_model_file": "final.zip",
+            "nn_model_file_sub_folder": ".",
             "tokens_file": "units.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": ".",
             "loader": "load_sherpa_offline_recognizer",
             "normalize_samples": False,
         },
         {
             "repo_id": "csukuangfj/sherpa-onnx-paraformer-zh-2024-03-09",
             "nn_model_file": "model.int8.onnx",
+            "nn_model_file_sub_folder": ".",
             "tokens_file": "tokens.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": ".",
             "loader": "load_sherpa_offline_recognizer_from_paraformer",
         },
         {
             "repo_id": "csukuangfj/sherpa-onnx-paraformer-zh-small-2024-03-09",
             "nn_model_file": "model.int8.onnx",
+            "nn_model_file_sub_folder": ".",
             "tokens_file": "tokens.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": ".",
             "loader": "load_sherpa_offline_recognizer_from_paraformer",
         },
         {
             "repo_id": "luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
             "nn_model_file": "cpu_jit_epoch_10_avg_2_torch_1.7.1.pt",
+            "nn_model_file_sub_folder": "exp",
             "tokens_file": "tokens.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": "data/lang_char",
             "loader": "load_sherpa_offline_recognizer",
             "normalize_samples": True,
         }
@@ -52,8 +56,9 @@ model_map = {
         {
             "repo_id": "csukuangfj/sherpa-onnx-paraformer-en-2024-03-09",
             "nn_model_file": "model.int8.onnx",
+            "nn_model_file_sub_folder": ".",
             "tokens_file": "tokens.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": ".",
             "loader": "load_sherpa_offline_recognizer_from_paraformer",
         },
     ],
@@ -61,8 +66,9 @@ model_map = {
         {
             "repo_id": "csukuangfj/sherpa-onnx-paraformer-zh-2023-03-28",
             "nn_model_file": "model.int8.onnx",
+            "nn_model_file_sub_folder": ".",
             "tokens_file": "tokens.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": ".",
             "loader": "load_sherpa_offline_recognizer_from_paraformer",
         },
     ],
@@ -70,8 +76,9 @@ model_map = {
         {
             "repo_id": "csukuangfj/sherpa-onnx-paraformer-trilingual-zh-cantonese-en",
             "nn_model_file": "model.int8.onnx",
+            "nn_model_file_sub_folder": ".",
             "tokens_file": "tokens.txt",
-            "sub_folder": ".",
+            "tokens_file_sub_folder": ".",
             "loader": "load_sherpa_offline_recognizer_from_paraformer",
         },
     ]
@@ -80,22 +87,23 @@ model_map = {
 
 def download_model(repo_id: str,
                    nn_model_file: str,
+                   nn_model_file_sub_folder: str,
                    tokens_file: str,
-                   sub_folder: str,
+                   tokens_file_sub_folder: str,
                    local_model_dir: str,
                    ):
 
     nn_model_file = huggingface_hub.hf_hub_download(
         repo_id=repo_id,
         filename=nn_model_file,
-        subfolder=sub_folder,
+        subfolder=nn_model_file_sub_folder,
         local_dir=local_model_dir,
     )
 
     tokens_file = huggingface_hub.hf_hub_download(
         repo_id=repo_id,
         filename=tokens_file,
-        subfolder=sub_folder,
+        subfolder=tokens_file_sub_folder,
         local_dir=local_model_dir,
     )
     return nn_model_file, tokens_file
@@ -150,8 +158,9 @@ def load_sherpa_offline_recognizer_from_paraformer(nn_model_file: str,
 
 def load_recognizer(repo_id: str,
                     nn_model_file: str,
+                    nn_model_file_sub_folder: str,
                     tokens_file: str,
-                    sub_folder: str,
+                    tokens_file_sub_folder: str,
                     local_model_dir: Path,
                     loader: str,
                     decoding_method: str = "greedy_search",
@@ -161,8 +170,9 @@ def load_recognizer(repo_id: str,
         download_model(
             repo_id=repo_id,
             nn_model_file=nn_model_file,
+            nn_model_file_sub_folder=nn_model_file_sub_folder,
             tokens_file=tokens_file,
-            sub_folder=sub_folder,
+            tokens_file_sub_folder=tokens_file_sub_folder,
             local_model_dir=local_model_dir.as_posix(),
         )