update

examples/wenet/downaload_model.py

@@ -15,15 +15,26 @@ from project_settings import project_path
 
 def get_args():
     parser = argparse.ArgumentParser()
+    # parser.add_argument(
+    #     "--repo_id",
+    #     default="csukuangfj/wenet-chinese-model",
+    #     # default="csukuangfj/wenet-english-model",
+    #     type=str
+    # )
+    # parser.add_argument("--model_filename", default="final.zip", type=str)
+    # parser.add_argument("--model_sub_folder", default=".", type=str)
+    # parser.add_argument("--tokens_filename", default="units.txt", type=str)
+    # parser.add_argument("--tokens_sub_folder", default=".", type=str)
+
     parser.add_argument(
         "--repo_id",
-        default="csukuangfj/wenet-chinese-model",
-        # default="csukuangfj/wenet-english-model",
+        default="luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
         type=str
     )
-
-    parser.add_argument("--model_filename", default="final.zip", type=str)
-    parser.add_argument("--tokens_filename", default="units.txt", type=str)
+    parser.add_argument("--model_filename", default="cpu_jit_epoch_10_avg_2_torch_1.7.1.pt", type=str)
+    parser.add_argument("--model_sub_folder", default="exp", type=str)
+    parser.add_argument("--tokens_filename", default="tokens.txt", type=str)
+    parser.add_argument("--tokens_sub_folder", default="data/lang_char", type=str)
 
     parser.add_argument(
         "--pretrained_model_dir",
@@ -41,13 +52,13 @@ def main():
     pretrained_model_dir.mkdir(exist_ok=True)
 
     model_dir = pretrained_model_dir / "huggingface" / args.repo_id
-    model_dir.mkdir(exist_ok=True)
+    model_dir.mkdir(parents=True, exist_ok=True)
 
     print("download model")
     model_filename = huggingface_hub.hf_hub_download(
         repo_id=args.repo_id,
         filename=args.model_filename,
-        subfolder=".",
+        subfolder=args.model_sub_folder,
         local_dir=model_dir.as_posix(),
     )
     print(model_filename)
@@ -56,7 +67,7 @@ def main():
     tokens_filename = huggingface_hub.hf_hub_download(
         repo_id=args.repo_id,
         filename=args.tokens_filename,
-        subfolder=".",
+        subfolder=args.tokens_sub_folder,
         local_dir=model_dir.as_posix(),
     )
     print(tokens_filename)

main.py

@@ -101,6 +101,7 @@ def process(
 
     # load recognizer
     local_model_dir = pretrained_model_dir / "huggingface" / repo_id
+    local_model_dir.mkdir(parents=True, exist_ok=True)
 
     recognizer = nn_models.load_recognizer(
         local_model_dir=local_model_dir,