minor fixes.

app.py

@@ -19,6 +19,7 @@
 # References:
 # https://gradio.app/docs/#dropdown
 
+import logging
 import os
 import time
 from datetime import datetime
@@ -26,43 +27,43 @@ from datetime import datetime
 import gradio as gr
 import torchaudio
 
-from model import (
-    get_gigaspeech_pre_trained_model,
-    sample_rate,
-    get_wenetspeech_pre_trained_model,
-)
+from model import get_pretrained_model, language_to_models, sample_rate
 
-models = {
-    "Chinese": get_wenetspeech_pre_trained_model(),
-    "English": get_gigaspeech_pre_trained_model(),
-}
+languages = sorted(language_to_models.keys())
 
 
 def convert_to_wav(in_filename: str) -> str:
     """Convert the input audio file to a wave file"""
     out_filename = in_filename + ".wav"
-    print(f"Converting '{in_filename}' to '{out_filename}'")
+    logging.info(f"Converting '{in_filename}' to '{out_filename}'")
     _ = os.system(f"ffmpeg -hide_banner -i '{in_filename}' '{out_filename}'")
     return out_filename
 
 
-demo = gr.Blocks()
-
+def process(
+    in_filename: str,
+    language: str,
+    repo_id: str,
+    decoding_method: str,
+    num_active_paths: int,
+) -> str:
+    logging.info(f"in_filename: {in_filename}")
+    logging.info(f"language: {language}")
+    logging.info(f"repo_id: {repo_id}")
+    logging.info(f"decoding_method: {decoding_method}")
+    logging.info(f"num_active_paths: {num_active_paths}")
 
-def process(in_filename: str, language: str) -> str:
-    print("in_filename", in_filename)
-    print("language", language)
     filename = convert_to_wav(in_filename)
 
     now = datetime.now()
     date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
-    print(f"Started at {date_time}")
+    logging.info(f"Started at {date_time}")
 
     start = time.time()
     wave, wave_sample_rate = torchaudio.load(filename)
 
     if wave_sample_rate != sample_rate:
-        print(
+        logging.info(
             f"Expected sample rate: {sample_rate}. Given: {wave_sample_rate}. "
             f"Resampling to {sample_rate}."
         )
@@ -74,7 +75,11 @@ def process(in_filename: str, language: str) -> str:
         )
     wave = wave[0]  # use only the first channel.
 
-    hyp = models[language].decode_waves([wave])[0]
+    hyp = get_pretrained_model(repo_id).decode_waves(
+        [wave],
+        decoding_method=decoding_method,
+        num_active_paths=num_active_paths,
+    )[0]
 
     date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
     end = time.time()
@@ -82,11 +87,10 @@ def process(in_filename: str, language: str) -> str:
     duration = wave.shape[0] / sample_rate
     rtf = (end - start) / duration
 
-    print(f"Finished at {date_time} s. Elapsed: {end - start: .3f} s")
-    print(f"Duration {duration: .3f} s")
-    print(f"RTF {rtf: .3f}")
-    print("hyp")
-    print(hyp)
+    logging.info(f"Finished at {date_time} s. Elapsed: {end - start: .3f} s")
+    logging.info(f"Duration {duration: .3f} s")
+    logging.info(f"RTF {rtf: .3f}")
+    logging.info(f"hyp:\n{hyp}")
 
     return hyp
 
@@ -103,51 +107,100 @@ See more information by visiting the following links:
 - <https://github.com/lhotse-speech/lhotse>
 """
 
+
+def update_model_dropdown(language: str):
+    if language in language_to_models:
+        choices = language_to_models[language]
+        return gr.Dropdown.update(choices=choices, value=choices[0])
+
+    raise ValueError(f"Unsupported language: {language}")
+
+
+demo = gr.Blocks()
+
 with demo:
     gr.Markdown(title)
-    gr.Markdown(description)
-    language_choices = list(models.keys())
-    language = gr.inputs.Radio(
+    language_choices = list(language_to_models.keys())
+
+    language_radio = gr.Radio(
         label="Language",
         choices=language_choices,
+        value=language_choices[0],
+    )
+    model_dropdown = gr.Dropdown(
+        choices=language_to_models[language_choices[0]],
+        label="Select a model",
+        value=language_to_models[language_choices[0]][0],
+    )
+
+    language_radio.change(
+        update_model_dropdown,
+        inputs=language_radio,
+        outputs=model_dropdown,
+    )
+
+    decoding_method_radio = gr.Radio(
+        label="Decoding method",
+        choices=["greedy_search", "modified_beam_search"],
+        value="greedy_search",
+    )
+
+    num_active_paths_slider = gr.Slider(
+        minimum=1,
+        value=4,
+        step=1,
+        label="Number of active paths for modified_beam_search",
     )
 
     with gr.Tabs():
         with gr.TabItem("Upload from disk"):
-            uploaded_file = gr.inputs.Audio(
+            uploaded_file = gr.Audio(
                 source="upload",  # Choose between "microphone", "upload"
                 type="filepath",
                 optional=False,
                 label="Upload from disk",
             )
             upload_button = gr.Button("Submit for recognition")
-            uploaded_output = gr.outputs.Textbox(
-                label="Recognized speech from uploaded file"
-            )
+            uploaded_output = gr.Textbox(label="Recognized speech from uploaded file")
 
         with gr.TabItem("Record from microphone"):
-            microphone = gr.inputs.Audio(
+            microphone = gr.Audio(
                 source="microphone",  # Choose between "microphone", "upload"
                 type="filepath",
                 optional=False,
                 label="Record from microphone",
             )
-            recorded_output = gr.outputs.Textbox(
-                label="Recognized speech from recordings"
-            )
 
             record_button = gr.Button("Submit for recognition")
+            recorded_output = gr.Textbox(label="Recognized speech from recordings")
 
         upload_button.click(
             process,
-            inputs=[uploaded_file, language],
+            inputs=[
+                uploaded_file,
+                language_radio,
+                model_dropdown,
+                decoding_method_radio,
+                num_active_paths_slider,
+            ],
             outputs=uploaded_output,
         )
         record_button.click(
             process,
-            inputs=[microphone, language],
+            inputs=[
+                microphone,
+                language_radio,
+                model_dropdown,
+                decoding_method_radio,
+                num_active_paths_slider,
+            ],
             outputs=recorded_output,
         )
+    gr.Markdown(description)
 
 if __name__ == "__main__":
+    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+
+    logging.basicConfig(format=formatter, level=logging.INFO)
+
     demo.launch()

model.py

@@ -23,52 +23,189 @@ from offline_asr import OfflineAsr
 sample_rate = 16000
 
 
-@lru_cache(maxsize=1)
-def get_gigaspeech_pre_trained_model():
+@lru_cache(maxsize=30)
+def get_pretrained_model(repo_id: str) -> OfflineAsr:
+    if repo_id in chinese_models:
+        return chinese_models[repo_id](repo_id)
+    elif repo_id in english_models:
+        return english_models[repo_id](repo_id)
+    elif repo_id in chinese_english_mixed_models:
+        return chinese_english_mixed_models[repo_id](repo_id)
+    else:
+        raise ValueError(f"Unsupported repo_id: {repo_id}")
+
+
+def _get_nn_model_filename(
+    repo_id: str,
+    filename: str,
+    subfolder: str = "exp",
+) -> str:
     nn_model_filename = hf_hub_download(
-        # It is converted from https://huggingface.co/wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2
-        repo_id="csukuangfj/icefall-asr-gigaspeech-pruned-transducer-stateless2",
-        filename="cpu_jit-epoch-29-avg-11-torch-1.10.0.pt",
-        subfolder="exp",
+        repo_id=repo_id,
+        filename=filename,
+        subfolder=subfolder,
     )
+    return nn_model_filename
+
 
+def _get_bpe_model_filename(
+    repo_id: str,
+    filename: str = "bpe.model",
+    subfolder: str = "data/lang_bpe_500",
+) -> str:
     bpe_model_filename = hf_hub_download(
-        repo_id="wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2",
-        filename="bpe.model",
-        subfolder="data/lang_bpe_500",
+        repo_id=repo_id,
+        filename=filename,
+        subfolder=subfolder,
+    )
+    return bpe_model_filename
+
+
+def _get_token_filename(
+    repo_id: str,
+    filename: str = "tokens.txt",
+    subfolder: str = "data/lang_char",
+) -> str:
+    token_filename = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        subfolder=subfolder,
+    )
+    return token_filename
+
+
+@lru_cache(maxsize=10)
+def _get_aishell2_pretrained_model(repo_id: str) -> OfflineAsr:
+    assert repo_id in [
+        # context-size 1
+        "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-A-2022-07-12",  # noqa
+        # context-size 2
+        "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-B-2022-07-12",  # noqa
+    ]
+
+    nn_model_filename = _get_nn_model_filename(
+        repo_id=repo_id,
+        filename="cpu_jit.pt",
+    )
+    token_filename = _get_token_filename(repo_id=repo_id)
+
+    return OfflineAsr(
+        nn_model_filename=nn_model_filename,
+        bpe_model_filename=None,
+        token_filename=token_filename,
+        sample_rate=sample_rate,
+        device="cpu",
+    )
+
+
+@lru_cache(maxsize=10)
+def _get_gigaspeech_pre_trained_model(repo_id: str) -> OfflineAsr:
+    assert repo_id in [
+        "wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2",
+    ]
+
+    nn_model_filename = _get_nn_model_filename(
+        # It is converted from https://huggingface.co/wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2  # noqa
+        repo_id="csukuangfj/icefall-asr-gigaspeech-pruned-transducer-stateless2",  # noqa
+        filename="cpu_jit-epoch-29-avg-11-torch-1.10.0.pt",
     )
+    bpe_model_filename = _get_bpe_model_filename(repo_id=repo_id)
 
     return OfflineAsr(
         nn_model_filename=nn_model_filename,
         bpe_model_filename=bpe_model_filename,
         token_filename=None,
-        decoding_method="greedy_search",
-        num_active_paths=4,
         sample_rate=sample_rate,
         device="cpu",
     )
 
 
-@lru_cache(maxsize=1)
-def get_wenetspeech_pre_trained_model():
-    nn_model_filename = hf_hub_download(
-        repo_id="luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
+@lru_cache(maxsize=10)
+def _get_librispeech_pre_trained_model(repo_id: str) -> OfflineAsr:
+    assert repo_id in [
+        "csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13",  # noqa
+    ]
+
+    nn_model_filename = _get_nn_model_filename(
+        repo_id=repo_id,
+        filename="cpu_jit.pt",
+    )
+    bpe_model_filename = _get_bpe_model_filename(repo_id=repo_id)
+
+    return OfflineAsr(
+        nn_model_filename=nn_model_filename,
+        bpe_model_filename=bpe_model_filename,
+        token_filename=None,
+        sample_rate=sample_rate,
+        device="cpu",
+    )
+
+
+@lru_cache(maxsize=10)
+def _get_wenetspeech_pre_trained_model(repo_id: str):
+    assert repo_id in [
+        "luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
+    ]
+
+    nn_model_filename = _get_nn_model_filename(
+        repo_id=repo_id,
         filename="cpu_jit_epoch_10_avg_2_torch_1.7.1.pt",
-        subfolder="exp",
     )
+    token_filename = _get_token_filename(repo_id=repo_id)
 
-    token_filename = hf_hub_download(
-        repo_id="luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
-        filename="tokens.txt",
-        subfolder="data/lang_char",
+    return OfflineAsr(
+        nn_model_filename=nn_model_filename,
+        bpe_model_filename=None,
+        token_filename=token_filename,
+        sample_rate=sample_rate,
+        device="cpu",
+    )
+
+
+@lru_cache(maxsize=10)
+def _get_tal_csasr_pre_trained_model(repo_id: str):
+    assert repo_id in [
+        "luomingshuang/icefall_asr_tal-csasr_pruned_transducer_stateless5",
+    ]
+
+    nn_model_filename = _get_nn_model_filename(
+        repo_id=repo_id,
+        filename="cpu_jit.pt",
     )
+    token_filename = _get_token_filename(repo_id=repo_id)
 
     return OfflineAsr(
         nn_model_filename=nn_model_filename,
         bpe_model_filename=None,
         token_filename=token_filename,
-        decoding_method="greedy_search",
-        num_active_paths=4,
         sample_rate=sample_rate,
         device="cpu",
     )
+
+
+chinese_models = {
+    "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-A-2022-07-12": _get_aishell2_pretrained_model,  # noqa
+    "yuekai/icefall-asr-aishell2-pruned-transducer-stateless5-B-2022-07-12": _get_aishell2_pretrained_model,  # noqa
+    "luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2": _get_wenetspeech_pre_trained_model,  # noqa
+}
+
+english_models = {
+    "wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2": _get_gigaspeech_pre_trained_model,  # noqa
+    "csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13": _get_librispeech_pre_trained_model,  # noqa
+}
+
+chinese_english_mixed_models = {
+    "luomingshuang/icefall_asr_tal-csasr_pruned_transducer_stateless5": _get_tal_csasr_pre_trained_model,  # noqa
+}
+
+all_models = {
+    **chinese_models,
+    **english_models,
+    **chinese_english_mixed_models,
+}
+
+language_to_models = {
+    "Chinese": sorted(chinese_models.keys()),
+    "English": sorted(english_models.keys()),
+    "Chinese+English": sorted(chinese_english_mixed_models.keys()),
+}

offline_asr.py

@@ -206,10 +206,10 @@ class OfflineAsr(object):
     def __init__(
         self,
         nn_model_filename: str,
-        bpe_model_filename: Optional[str],
-        token_filename: Optional[str],
-        decoding_method: str,
-        num_active_paths: int,
+        bpe_model_filename: Optional[str] = None,
+        token_filename: Optional[str] = None,
+        decoding_method: str = "greedy_search",
+        num_active_paths: int = 4,
         sample_rate: int = 16000,
         device: Union[str, torch.device] = "cpu",
     ):
@@ -223,14 +223,6 @@ class OfflineAsr(object):
           token_filename:
             Path to tokens.txt. If it is None, you have to provide
             `bpe_model_filename`.
-          decoding_method:
-            The decoding method to use. Currently, only greedy_search and
-            modified_beam_search are implemented.
-          num_active_paths:
-            Used only when decoding_method is modified_beam_search.
-            It specifies number of active paths for each utterance. Due to
-            merging paths with identical token sequences, the actual number
-            may be less than "num_active_paths".
           sample_rate:
             Expected sample rate of the feature extractor.
           device:
@@ -246,6 +238,7 @@ class OfflineAsr(object):
             self.sp = spm.SentencePieceProcessor()
             self.sp.load(bpe_model_filename)
         else:
+            assert token_filename is not None, token_filename
             self.token_table = k2.SymbolTable.from_file(token_filename)
 
         self.feature_extractor = self._build_feature_extractor(
@@ -253,24 +246,6 @@ class OfflineAsr(object):
             device=device,
         )
 
-        assert decoding_method in (
-            "greedy_search",
-            "modified_beam_search",
-        ), decoding_method
-        if decoding_method == "greedy_search":
-            nn_and_decoding_func = run_model_and_do_greedy_search
-        elif decoding_method == "modified_beam_search":
-            nn_and_decoding_func = functools.partial(
-                run_model_and_do_modified_beam_search,
-                num_active_paths=num_active_paths,
-            )
-        else:
-            raise ValueError(
-                f"Unsupported decoding_method: {decoding_method} "
-                "Please use greedy_search or modified_beam_search"
-            )
-
-        self.nn_and_decoding_func = nn_and_decoding_func
         self.device = device
 
     def _build_feature_extractor(
@@ -299,7 +274,12 @@ class OfflineAsr(object):
 
         return fbank
 
-    def decode_waves(self, waves: List[torch.Tensor]) -> List[List[str]]:
+    def decode_waves(
+        self,
+        waves: List[torch.Tensor],
+        decoding_method: str,
+        num_active_paths: int,
+    ) -> List[List[str]]:
         """
         Args:
           waves:
@@ -313,20 +293,48 @@ class OfflineAsr(object):
               then the given waves have to contain samples in this range.
 
               All models trained in icefall use the normalized range [-1, 1].
+          decoding_method:
+            The decoding method to use. Currently, only greedy_search and
+            modified_beam_search are implemented.
+          num_active_paths:
+            Used only when decoding_method is modified_beam_search.
+            It specifies number of active paths for each utterance. Due to
+            merging paths with identical token sequences, the actual number
+            may be less than "num_active_paths".
         Returns:
           Return a list of decoded results. `ans[i]` contains the decoded
           results for `wavs[i]`.
         """
+        assert decoding_method in (
+            "greedy_search",
+            "modified_beam_search",
+        ), decoding_method
+
+        if decoding_method == "greedy_search":
+            nn_and_decoding_func = run_model_and_do_greedy_search
+        elif decoding_method == "modified_beam_search":
+            nn_and_decoding_func = functools.partial(
+                run_model_and_do_modified_beam_search,
+                num_active_paths=num_active_paths,
+            )
+        else:
+            raise ValueError(
+                f"Unsupported decoding_method: {decoding_method} "
+                "Please use greedy_search or modified_beam_search"
+            )
+
         waves = [w.to(self.device) for w in waves]
         features = self.feature_extractor(waves)
 
-        tokens = self.nn_and_decoding_func(self.model, features)
+        tokens = nn_and_decoding_func(self.model, features)
 
         if hasattr(self, "sp"):
             results = self.sp.decode(tokens)
         else:
             results = [[self.token_table[i] for i in hyp] for hyp in tokens]
+            blank = chr(0x2581)
             results = ["".join(r) for r in results]
+            results = [r.replace(blank, " ") for r in results]
 
         return results