update

toolbox/k2_sherpa/nn_models.py

@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 from enum import Enum
 from functools import lru_cache
+import logging
 import os
 import platform
 from pathlib import Path
@@ -10,6 +11,8 @@ import huggingface_hub
 import sherpa
 import sherpa_onnx
 
+main_logger = logging.getLogger("main")
+
 
 class EnumDecodingMethod(Enum):
     greedy_search = "greedy_search"
@@ -104,6 +107,7 @@ def download_model(local_model_dir: str,
     repo_id = kwargs["repo_id"]
 
     if "nn_model_file" in kwargs.keys():
+        main_logger.info("download nn_model_file. filename: {}, subfolder: {}".format(kwargs["nn_model_file"], kwargs["nn_model_file_sub_folder"]))
         _ = huggingface_hub.hf_hub_download(
             repo_id=repo_id,
             filename=kwargs["nn_model_file"],
@@ -112,6 +116,7 @@ def download_model(local_model_dir: str,
         )
 
     if "encoder_model_file" in kwargs.keys():
+        main_logger.info("download encoder_model_file. filename: {}, subfolder: {}".format(kwargs["encoder_model_file"], kwargs["encoder_model_file_sub_folder"]))
         _ = huggingface_hub.hf_hub_download(
             repo_id=repo_id,
             filename=kwargs["encoder_model_file"],
@@ -120,6 +125,7 @@ def download_model(local_model_dir: str,
         )
 
     if "decoder_model_file" in kwargs.keys():
+        main_logger.info("download decoder_model_file. filename: {}, subfolder: {}".format(kwargs["decoder_model_file"], kwargs["decoder_model_file_sub_folder"]))
         _ = huggingface_hub.hf_hub_download(
             repo_id=repo_id,
             filename=kwargs["decoder_model_file"],
@@ -128,6 +134,7 @@ def download_model(local_model_dir: str,
         )
 
     if "joiner_model_file" in kwargs.keys():
+        main_logger.info("download joiner_model_file. filename: {}, subfolder: {}".format(kwargs["joiner_model_file"], kwargs["joiner_model_file_sub_folder"]))
         _ = huggingface_hub.hf_hub_download(
             repo_id=repo_id,
             filename=kwargs["joiner_model_file"],
@@ -136,6 +143,7 @@ def download_model(local_model_dir: str,
         )
 
     if "tokens_file" in kwargs.keys():
+        main_logger.info("download tokens_file. filename: {}, subfolder: {}".format(kwargs["tokens_file"], kwargs["tokens_file_sub_folder"]))
         _ = huggingface_hub.hf_hub_download(
             repo_id=repo_id,
             filename=kwargs["tokens_file"],
@@ -158,6 +166,9 @@ def load_sherpa_offline_recognizer(nn_model_file: str,
     feat_config.fbank_opts.mel_opts.num_bins = num_mel_bins
     feat_config.fbank_opts.frame_opts.dither = frame_dither
 
+    if not os.path.exists(nn_model_file):
+        raise AssertionError("nn_model_file not found. ")
+
     config = sherpa.OfflineRecognizerConfig(
         nn_model=nn_model_file,
         tokens=tokens_file,
@@ -220,7 +231,7 @@ def load_recognizer(local_model_dir: Path,
                     num_active_paths: int = 4,
                     **kwargs
                     ):
-    if not os.path.exists(local_model_dir):
+    if not local_model_dir.exists():
         download_model(
             local_model_dir=local_model_dir.as_posix(),
             **kwargs,