First commit

app.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+#
+# Copyright      2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# References:
+# https://gradio.app/docs/#dropdown
+
+import logging
+import os
+import time
+from datetime import datetime
+
+import gradio as gr
+import torchaudio
+
+from model import get_pretrained_model, language_to_models, sample_rate
+
+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"
+    logging.info(f"Converting '{in_filename}' to '{out_filename}'")
+    _ = os.system(f"ffmpeg -hide_banner -i '{in_filename}' '{out_filename}'")
+    return out_filename
+
+
+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}")
+
+    filename = convert_to_wav(in_filename)
+
+    now = datetime.now()
+    date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
+    logging.info(f"Started at {date_time}")
+
+    start = time.time()
+    wave, wave_sample_rate = torchaudio.load(filename)
+
+    if wave_sample_rate != sample_rate:
+        logging.info(
+            f"Expected sample rate: {sample_rate}. Given: {wave_sample_rate}. "
+            f"Resampling to {sample_rate}."
+        )
+
+        wave = torchaudio.functional.resample(
+            wave,
+            orig_freq=wave_sample_rate,
+            new_freq=sample_rate,
+        )
+    wave = wave[0]  # use only the first channel.
+
+    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()
+
+    duration = wave.shape[0] / sample_rate
+    rtf = (end - start) / duration
+
+    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
+
+
+title = "# Automatic Speech Recognition with Next-gen Kaldi"
+description = """
+This space shows how to do automatic speech recognition with Next-gen Kaldi.
+
+See more information by visiting the following links:
+
+- <https://github.com/k2-fsa/icefall>
+- <https://github.com/k2-fsa/sherpa>
+- <https://github.com/k2-fsa/k2>
+- <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)
+    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.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.Textbox(label="Recognized speech from uploaded file")
+
+        with gr.TabItem("Record from microphone"):
+            microphone = gr.Audio(
+                source="microphone",  # Choose between "microphone", "upload"
+                type="filepath",
+                optional=False,
+                label="Record from microphone",
+            )
+
+            record_button = gr.Button("Submit for recognition")
+            recorded_output = gr.Textbox(label="Recognized speech from recordings")
+
+        upload_button.click(
+            process,
+            inputs=[
+                uploaded_file,
+                language_radio,
+                model_dropdown,
+                decoding_method_radio,
+                num_active_paths_slider,
+            ],
+            outputs=uploaded_output,
+        )
+        record_button.click(
+            process,
+            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()

decode.py

@@ -0,0 +1,121 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# Copied from https://github.com/k2-fsa/sherpa/blob/master/sherpa/bin/conformer_rnnt/decode.py
+#
+# See LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import List
+
+import torch
+from sherpa import RnntConformerModel, greedy_search, modified_beam_search
+from torch.nn.utils.rnn import pad_sequence
+
+LOG_EPS = math.log(1e-10)
+
+
+@torch.no_grad()
+def run_model_and_do_greedy_search(
+    model: RnntConformerModel,
+    features: List[torch.Tensor],
+) -> List[List[int]]:
+    """Run RNN-T model with the given features and use greedy search
+    to decode the output of the model.
+
+    Args:
+      model:
+        The RNN-T model.
+      features:
+        A list of 2-D tensors. Each entry is of shape
+        (num_frames, feature_dim).
+    Returns:
+      Return a list-of-list containing the decoding token IDs.
+    """
+    features_length = torch.tensor(
+        [f.size(0) for f in features],
+        dtype=torch.int64,
+    )
+    features = pad_sequence(
+        features,
+        batch_first=True,
+        padding_value=LOG_EPS,
+    )
+
+    device = model.device
+    features = features.to(device)
+    features_length = features_length.to(device)
+
+    encoder_out, encoder_out_length = model.encoder(
+        features=features,
+        features_length=features_length,
+    )
+
+    hyp_tokens = greedy_search(
+        model=model,
+        encoder_out=encoder_out,
+        encoder_out_length=encoder_out_length.cpu(),
+    )
+    return hyp_tokens
+
+
+@torch.no_grad()
+def run_model_and_do_modified_beam_search(
+    model: RnntConformerModel,
+    features: List[torch.Tensor],
+    num_active_paths: int,
+) -> List[List[int]]:
+    """Run RNN-T model with the given features and use greedy search
+    to decode the output of the model.
+
+    Args:
+      model:
+        The RNN-T model.
+      features:
+        A list of 2-D tensors. Each entry is of shape
+        (num_frames, feature_dim).
+      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-list containing the decoding token IDs.
+    """
+    features_length = torch.tensor(
+        [f.size(0) for f in features],
+        dtype=torch.int64,
+    )
+    features = pad_sequence(
+        features,
+        batch_first=True,
+        padding_value=LOG_EPS,
+    )
+
+    device = model.device
+    features = features.to(device)
+    features_length = features_length.to(device)
+
+    encoder_out, encoder_out_length = model.encoder(
+        features=features,
+        features_length=features_length,
+    )
+
+    hyp_tokens = modified_beam_search(
+        model=model,
+        encoder_out=encoder_out,
+        encoder_out_length=encoder_out_length.cpu(),
+        num_active_paths=num_active_paths,
+    )
+    return hyp_tokens

model.py

@@ -0,0 +1,211 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# See LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from huggingface_hub import hf_hub_download
+from functools import lru_cache
+
+
+from offline_asr import OfflineAsr
+
+sample_rate = 16000
+
+
+@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(
+        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=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,
+        sample_rate=sample_rate,
+        device="cpu",
+    )
+
+
+@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",
+    )
+    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_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,
+        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

@@ -0,0 +1,427 @@
+#!/usr/bin/env python3
+# Copyright      2022  Xiaomi Corp.        (authors: Fangjun Kuang)
+#
+# Copied from https://github.com/k2-fsa/sherpa/blob/master/sherpa/bin/conformer_rnnt/offline_asr.py
+#
+# See LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+A standalone script for offline ASR recognition.
+
+It loads a torchscript model, decodes the given wav files, and exits.
+
+Usage:
+    ./offline_asr.py --help
+
+For BPE based models (e.g., LibriSpeech):
+
+    ./offline_asr.py \
+        --nn-model-filename /path/to/cpu_jit.pt \
+        --bpe-model-filename /path/to/bpe.model \
+        --decoding-method greedy_search \
+        ./foo.wav \
+        ./bar.wav \
+        ./foobar.wav
+
+For character based models (e.g., aishell):
+
+    ./offline.py \
+        --nn-model-filename /path/to/cpu_jit.pt \
+        --token-filename /path/to/lang_char/tokens.txt \
+        --decoding-method greedy_search \
+        ./foo.wav \
+        ./bar.wav \
+        ./foobar.wav
+
+Note: We provide pre-trained models for testing.
+
+(1) Pre-trained model with the LibriSpeech dataset
+
+    sudo apt-get install git-lfs
+    git lfs install
+    git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+
+    nn_model_filename=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/exp/cpu_jit-torch-1.6.0.pt
+    bpe_model=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/data/lang_bpe_500/bpe.model
+
+    wav1=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/test_wavs/1089-134686-0001.wav
+    wav2=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/test_wavs/1221-135766-0001.wav
+    wav3=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/test_wavs/1221-135766-0002.wav
+
+    sherpa/bin/conformer_rnnt/offline_asr.py \
+      --nn-model-filename $nn_model_filename \
+      --bpe-model $bpe_model \
+      $wav1 \
+      $wav2 \
+      $wav3
+
+(2) Pre-trained model with the aishell dataset
+
+    sudo apt-get install git-lfs
+    git lfs install
+    git clone https://huggingface.co/csukuangfj/icefall-aishell-pruned-transducer-stateless3-2022-06-20
+
+    nn_model_filename=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/exp/cpu_jit-epoch-29-avg-5-torch-1.6.0.pt
+    token_filename=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/data/lang_char/tokens.txt
+
+    wav1=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/test_wavs/BAC009S0764W0121.wav
+    wav2=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/test_wavs/BAC009S0764W0122.wav
+    wav3=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/test_wavs/BAC009S0764W0123.wav
+
+    sherpa/bin/conformer_rnnt/offline_asr.py \
+      --nn-model-filename $nn_model_filename \
+      --token-filename $token_filename \
+      $wav1 \
+      $wav2 \
+      $wav3
+"""
+import argparse
+import functools
+import logging
+from typing import List, Optional, Union
+
+import k2
+import kaldifeat
+import sentencepiece as spm
+import torch
+import torchaudio
+from sherpa import RnntConformerModel
+
+from decode import run_model_and_do_greedy_search, run_model_and_do_modified_beam_search
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--nn-model-filename",
+        type=str,
+        help="""The torchscript model. You can use
+          icefall/egs/librispeech/ASR/pruned_transducer_statelessX/export.py \
+             --jit=1
+        to generate this model.
+        """,
+    )
+
+    parser.add_argument(
+        "--bpe-model-filename",
+        type=str,
+        help="""The BPE model
+        You can find it in the directory egs/librispeech/ASR/data/lang_bpe_xxx
+        from icefall,
+        where xxx is the number of BPE tokens you used to train the model.
+        Note: Use it only when your model is using BPE. You don't need to
+        provide it if you provide `--token-filename`
+        """,
+    )
+
+    parser.add_argument(
+        "--token-filename",
+        type=str,
+        help="""Filename for tokens.txt
+        You can find it in the directory
+        egs/aishell/ASR/data/lang_char/tokens.txt from icefall.
+        Note: You don't need to provide it if you provide `--bpe-model`
+        """,
+    )
+
+    parser.add_argument(
+        "--decoding-method",
+        type=str,
+        default="greedy_search",
+        help="""Decoding method to use. Currently, only greedy_search and
+        modified_beam_search are implemented.
+        """,
+    )
+
+    parser.add_argument(
+        "--num-active-paths",
+        type=int,
+        default=4,
+        help="""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".
+        """,
+    )
+
+    parser.add_argument(
+        "--sample-rate",
+        type=int,
+        default=16000,
+        help="The expected sample rate of the input sound files",
+    )
+
+    parser.add_argument(
+        "sound_files",
+        type=str,
+        nargs="+",
+        help="The input sound file(s) to transcribe. "
+        "Supported formats are those supported by torchaudio.load(). "
+        "For example, wav and flac are supported. "
+        "The sample rate has to equal to `--sample-rate`.",
+    )
+
+    return parser.parse_args()
+
+
+def read_sound_files(
+    filenames: List[str],
+    expected_sample_rate: int,
+) -> List[torch.Tensor]:
+    """Read a list of sound files into a list 1-D float32 torch tensors.
+    Args:
+      filenames:
+        A list of sound filenames.
+      expected_sample_rate:
+        The expected sample rate of the sound files.
+    Returns:
+      Return a list of 1-D float32 torch tensors.
+    """
+    ans = []
+    for f in filenames:
+        wave, sample_rate = torchaudio.load(f)
+        assert sample_rate == expected_sample_rate, (
+            f"expected sample rate: {expected_sample_rate}. " f"Given: {sample_rate}"
+        )
+        # We use only the first channel
+        ans.append(wave[0])
+    return ans
+
+
+class OfflineAsr(object):
+    def __init__(
+        self,
+        nn_model_filename: str,
+        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",
+    ):
+        """
+        Args:
+          nn_model_filename:
+            Path to the torch script model.
+          bpe_model_filename:
+            Path to the BPE model. If it is None, you have to provide
+            `token_filename`.
+          token_filename:
+            Path to tokens.txt. If it is None, you have to provide
+            `bpe_model_filename`.
+          sample_rate:
+            Expected sample rate of the feature extractor.
+          device:
+            The device to use for computation.
+        """
+        self.model = RnntConformerModel(
+            filename=nn_model_filename,
+            device=device,
+            optimize_for_inference=False,
+        )
+
+        if bpe_model_filename:
+            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(
+            sample_rate=sample_rate,
+            device=device,
+        )
+
+        self.device = device
+
+    def _build_feature_extractor(
+        self,
+        sample_rate: int = 16000,
+        device: Union[str, torch.device] = "cpu",
+    ) -> kaldifeat.OfflineFeature:
+        """Build a fbank feature extractor for extracting features.
+
+        Args:
+          sample_rate:
+            Expected sample rate of the feature extractor.
+          device:
+            The device to use for computation.
+        Returns:
+          Return a fbank feature extractor.
+        """
+        opts = kaldifeat.FbankOptions()
+        opts.device = device
+        opts.frame_opts.dither = 0
+        opts.frame_opts.snip_edges = False
+        opts.frame_opts.samp_freq = sample_rate
+        opts.mel_opts.num_bins = 80
+
+        fbank = kaldifeat.Fbank(opts)
+
+        return fbank
+
+    def decode_waves(
+        self,
+        waves: List[torch.Tensor],
+        decoding_method: str,
+        num_active_paths: int,
+    ) -> List[List[str]]:
+        """
+        Args:
+          waves:
+            A list of 1-D torch.float32 tensors containing audio samples.
+            wavs[i] contains audio samples for the i-th utterance.
+
+            Note:
+              Whether it should be in the range [-32768, 32767] or be normalized
+              to [-1, 1] depends on which range you used for your training data.
+              For instance, if your training data used [-32768, 32767],
+              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 = 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
+
+
+@torch.no_grad()
+def main():
+    args = get_args()
+    logging.info(vars(args))
+
+    nn_model_filename = args.nn_model_filename
+    bpe_model_filename = args.bpe_model_filename
+    token_filename = args.token_filename
+    decoding_method = args.decoding_method
+    num_active_paths = args.num_active_paths
+    sample_rate = args.sample_rate
+    sound_files = args.sound_files
+
+    assert decoding_method in ("greedy_search", "modified_beam_search"), decoding_method
+
+    if decoding_method == "modified_beam_search":
+        assert num_active_paths >= 1, num_active_paths
+
+    if bpe_model_filename:
+        assert token_filename is None
+
+    if token_filename:
+        assert bpe_model_filename is None
+
+    device = torch.device("cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+
+    logging.info(f"device: {device}")
+
+    offline_asr = OfflineAsr(
+        nn_model_filename=nn_model_filename,
+        bpe_model_filename=bpe_model_filename,
+        token_filename=token_filename,
+        decoding_method=decoding_method,
+        num_active_paths=num_active_paths,
+        sample_rate=sample_rate,
+        device=device,
+    )
+
+    waves = read_sound_files(
+        filenames=sound_files,
+        expected_sample_rate=sample_rate,
+    )
+
+    logging.info("Decoding started.")
+
+    hyps = offline_asr.decode_waves(waves)
+
+    s = "\n"
+    for filename, hyp in zip(sound_files, hyps):
+        s += f"{filename}:\n{hyp}\n\n"
+    logging.info(s)
+
+    logging.info("Decoding done.")
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+# See https://github.com/pytorch/pytorch/issues/38342
+# and https://github.com/pytorch/pytorch/issues/33354
+#
+# If we don't do this, the delay increases whenever there is
+# a new request that changes the actual batch size.
+# If you use `py-spy dump --pid <server-pid> --native`, you will
+# see a lot of time is spent in re-compiling the torch script model.
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_set_profiling_mode(False)
+torch._C._set_graph_executor_optimize(False)
+"""
+// Use the following in C++
+torch::jit::getExecutorMode() = false;
+torch::jit::getProfilingMode() = false;
+torch::jit::setGraphExecutorOptimize(false);
+"""
+
+if __name__ == "__main__":
+    torch.manual_seed(20220609)
+
+    formatter = (
+        "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"  # noqa
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+
+    main()

requirements.txt

@@ -0,0 +1,13 @@
+https://download.pytorch.org/whl/cpu/torch-1.10.0%2Bcpu-cp38-cp38-linux_x86_64.whl
+https://k2-fsa.org/nightly/whl/k2-1.17.dev20220711+cpu.torch1.10.0-cp38-cp38-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torchaudio-0.10.0%2Bcpu-cp38-cp38-linux_x86_64.whl
+
+
+https://huggingface.co/csukuangfj/wheels/resolve/main/kaldifeat-1.17-cp38-cp38-linux_x86_64.whl
+https://huggingface.co/csukuangfj/wheels/resolve/main/k2_sherpa-0.6-cp38-cp38-linux_x86_64.whl
+
+
+sentencepiece>=0.1.96
+numpy
+
+huggingface_hub