small fixes

app.py

@@ -26,16 +26,9 @@ 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:
@@ -46,12 +39,10 @@ def convert_to_wav(in_filename: str) -> str:
     return out_filename
 
 
-demo = gr.Blocks()
-
-
-def process(in_filename: str, language: str) -> str:
+def process(in_filename: str, language: str, repo_id: str) -> str:
     print("in_filename", in_filename)
     print("language", language)
+    print("repo_id", repo_id)
     filename = convert_to_wav(in_filename)
 
     now = datetime.now()
@@ -74,7 +65,7 @@ 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])[0]
 
     date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
     end = time.time()
@@ -103,14 +94,32 @@ 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,
     )
+    model_dropdown = gr.Dropdown(choices=[], label="Select a model")
+    language_radio.change(
+        update_model_dropdown,
+        inputs=language_radio,
+        outputs=model_dropdown,
+    )
 
     with gr.Tabs():
         with gr.TabItem("Upload from disk"):
@@ -140,12 +149,12 @@ with demo:
 
         upload_button.click(
             process,
-            inputs=[uploaded_file, language],
+            inputs=[uploaded_file, language_radio, model_dropdown],
             outputs=uploaded_output,
         )
         record_button.click(
             process,
-            inputs=[microphone, language],
+            inputs=[microphone, language_radio, model_dropdown],
             outputs=recorded_output,
         )
 

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:
+        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",
     ):
@@ -246,6 +246,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(