csukuangfj commited on
Commit
071812d
1 Parent(s): 99b54d7

Add ASR demo with Next-gen Kaldi

Browse files
Files changed (5) hide show
  1. app.py +153 -0
  2. decode.py +121 -0
  3. model.py +74 -0
  4. offline_asr.py +419 -0
  5. requirements.txt +13 -0
app.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ #
3
+ # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
4
+ #
5
+ # See LICENSE for clarification regarding multiple authors
6
+ #
7
+ # Licensed under the Apache License, Version 2.0 (the "License");
8
+ # you may not use this file except in compliance with the License.
9
+ # You may obtain a copy of the License at
10
+ #
11
+ # http://www.apache.org/licenses/LICENSE-2.0
12
+ #
13
+ # Unless required by applicable law or agreed to in writing, software
14
+ # distributed under the License is distributed on an "AS IS" BASIS,
15
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+ # See the License for the specific language governing permissions and
17
+ # limitations under the License.
18
+
19
+ # References:
20
+ # https://gradio.app/docs/#dropdown
21
+
22
+ import os
23
+ import time
24
+ from datetime import datetime
25
+
26
+ import gradio as gr
27
+ import torchaudio
28
+
29
+ from model import (
30
+ get_gigaspeech_pre_trained_model,
31
+ sample_rate,
32
+ get_wenetspeech_pre_trained_model,
33
+ )
34
+
35
+ models = {
36
+ "Chinese": get_wenetspeech_pre_trained_model(),
37
+ "English": get_gigaspeech_pre_trained_model(),
38
+ }
39
+
40
+
41
+ def convert_to_wav(in_filename: str) -> str:
42
+ """Convert the input audio file to a wave file"""
43
+ out_filename = in_filename + ".wav"
44
+ print(f"Converting '{in_filename}' to '{out_filename}'")
45
+ _ = os.system(f"ffmpeg -hide_banner -i '{in_filename}' '{out_filename}'")
46
+ return out_filename
47
+
48
+
49
+ demo = gr.Blocks()
50
+
51
+
52
+ def process(in_filename: str, language: str) -> str:
53
+ print("in_filename", in_filename)
54
+ print("language", language)
55
+ filename = convert_to_wav(in_filename)
56
+
57
+ now = datetime.now()
58
+ date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
59
+ print(f"Started at {date_time}")
60
+
61
+ start = time.time()
62
+ wave, wave_sample_rate = torchaudio.load(filename)
63
+
64
+ if wave_sample_rate != sample_rate:
65
+ print(
66
+ f"Expected sample rate: {sample_rate}. Given: {wave_sample_rate}. "
67
+ f"Resampling to {sample_rate}."
68
+ )
69
+
70
+ wave = torchaudio.functional.resample(
71
+ wave,
72
+ orig_freq=wave_sample_rate,
73
+ new_freq=sample_rate,
74
+ )
75
+ wave = wave[0] # use only the first channel.
76
+
77
+ hyp = models[language].decode_waves([wave])[0]
78
+
79
+ date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
80
+ end = time.time()
81
+
82
+ duration = wave.shape[0] / sample_rate
83
+ rtf = (end - start) / duration
84
+
85
+ print(f"Finished at {date_time} s. Elapsed: {end - start: .3f} s")
86
+ print(f"Duration {duration: .3f} s")
87
+ print(f"RTF {rtf: .3f}")
88
+ print("hyp")
89
+ print(hyp)
90
+
91
+ return hyp
92
+
93
+
94
+ title = "# Automatic Speech Recognition with Next-gen Kaldi"
95
+ description = """
96
+ This space shows how to do automatic speech recognition with Next-gen Kaldi.
97
+
98
+ See more information by visiting the following links:
99
+
100
+ - <https://github.com/k2-fsa/icefall>
101
+ - <https://github.com/k2-fsa/sherpa>
102
+ - <https://github.com/k2-fsa/k2>
103
+ - <https://github.com/lhotse-speech/lhotse>
104
+ """
105
+
106
+ with demo:
107
+ gr.Markdown(title)
108
+ gr.Markdown(description)
109
+ language_choices = list(models.keys())
110
+ language = gr.inputs.Radio(
111
+ label="Language",
112
+ choices=language_choices,
113
+ )
114
+
115
+ with gr.Tabs():
116
+ with gr.TabItem("Upload from disk"):
117
+ uploaded_file = gr.inputs.Audio(
118
+ source="upload", # Choose between "microphone", "upload"
119
+ type="filepath",
120
+ optional=False,
121
+ label="Upload from disk",
122
+ )
123
+ upload_button = gr.Button("Submit for recognition")
124
+ uploaded_output = gr.outputs.Textbox(
125
+ label="Recognized speech from uploaded file"
126
+ )
127
+
128
+ with gr.TabItem("Record from microphone"):
129
+ microphone = gr.inputs.Audio(
130
+ source="microphone", # Choose between "microphone", "upload"
131
+ type="filepath",
132
+ optional=False,
133
+ label="Record from microphone",
134
+ )
135
+ recorded_output = gr.outputs.Textbox(
136
+ label="Recognized speech from recordings"
137
+ )
138
+
139
+ record_button = gr.Button("Submit for recordings")
140
+
141
+ upload_button.click(
142
+ process,
143
+ inputs=[uploaded_file, language],
144
+ outputs=uploaded_output,
145
+ )
146
+ record_button.click(
147
+ process,
148
+ inputs=[microphone, language],
149
+ outputs=recorded_output,
150
+ )
151
+
152
+ if __name__ == "__main__":
153
+ demo.launch()
decode.py ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
2
+ #
3
+ # Copied from https://github.com/k2-fsa/sherpa/blob/master/sherpa/bin/conformer_rnnt/decode.py
4
+ #
5
+ # See LICENSE for clarification regarding multiple authors
6
+ #
7
+ # Licensed under the Apache License, Version 2.0 (the "License");
8
+ # you may not use this file except in compliance with the License.
9
+ # You may obtain a copy of the License at
10
+ #
11
+ # http://www.apache.org/licenses/LICENSE-2.0
12
+ #
13
+ # Unless required by applicable law or agreed to in writing, software
14
+ # distributed under the License is distributed on an "AS IS" BASIS,
15
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+ # See the License for the specific language governing permissions and
17
+ # limitations under the License.
18
+
19
+ import math
20
+ from typing import List
21
+
22
+ import torch
23
+ from sherpa import RnntConformerModel, greedy_search, modified_beam_search
24
+ from torch.nn.utils.rnn import pad_sequence
25
+
26
+ LOG_EPS = math.log(1e-10)
27
+
28
+
29
+ @torch.no_grad()
30
+ def run_model_and_do_greedy_search(
31
+ model: RnntConformerModel,
32
+ features: List[torch.Tensor],
33
+ ) -> List[List[int]]:
34
+ """Run RNN-T model with the given features and use greedy search
35
+ to decode the output of the model.
36
+
37
+ Args:
38
+ model:
39
+ The RNN-T model.
40
+ features:
41
+ A list of 2-D tensors. Each entry is of shape
42
+ (num_frames, feature_dim).
43
+ Returns:
44
+ Return a list-of-list containing the decoding token IDs.
45
+ """
46
+ features_length = torch.tensor(
47
+ [f.size(0) for f in features],
48
+ dtype=torch.int64,
49
+ )
50
+ features = pad_sequence(
51
+ features,
52
+ batch_first=True,
53
+ padding_value=LOG_EPS,
54
+ )
55
+
56
+ device = model.device
57
+ features = features.to(device)
58
+ features_length = features_length.to(device)
59
+
60
+ encoder_out, encoder_out_length = model.encoder(
61
+ features=features,
62
+ features_length=features_length,
63
+ )
64
+
65
+ hyp_tokens = greedy_search(
66
+ model=model,
67
+ encoder_out=encoder_out,
68
+ encoder_out_length=encoder_out_length.cpu(),
69
+ )
70
+ return hyp_tokens
71
+
72
+
73
+ @torch.no_grad()
74
+ def run_model_and_do_modified_beam_search(
75
+ model: RnntConformerModel,
76
+ features: List[torch.Tensor],
77
+ num_active_paths: int,
78
+ ) -> List[List[int]]:
79
+ """Run RNN-T model with the given features and use greedy search
80
+ to decode the output of the model.
81
+
82
+ Args:
83
+ model:
84
+ The RNN-T model.
85
+ features:
86
+ A list of 2-D tensors. Each entry is of shape
87
+ (num_frames, feature_dim).
88
+ num_active_paths:
89
+ Used only when decoding_method is modified_beam_search.
90
+ It specifies number of active paths for each utterance. Due to
91
+ merging paths with identical token sequences, the actual number
92
+ may be less than "num_active_paths".
93
+ Returns:
94
+ Return a list-of-list containing the decoding token IDs.
95
+ """
96
+ features_length = torch.tensor(
97
+ [f.size(0) for f in features],
98
+ dtype=torch.int64,
99
+ )
100
+ features = pad_sequence(
101
+ features,
102
+ batch_first=True,
103
+ padding_value=LOG_EPS,
104
+ )
105
+
106
+ device = model.device
107
+ features = features.to(device)
108
+ features_length = features_length.to(device)
109
+
110
+ encoder_out, encoder_out_length = model.encoder(
111
+ features=features,
112
+ features_length=features_length,
113
+ )
114
+
115
+ hyp_tokens = modified_beam_search(
116
+ model=model,
117
+ encoder_out=encoder_out,
118
+ encoder_out_length=encoder_out_length.cpu(),
119
+ num_active_paths=num_active_paths,
120
+ )
121
+ return hyp_tokens
model.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
2
+ #
3
+ # See LICENSE for clarification regarding multiple authors
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+
17
+ from huggingface_hub import hf_hub_download
18
+ from functools import lru_cache
19
+
20
+
21
+ from offline_asr import OfflineAsr
22
+
23
+ sample_rate = 16000
24
+
25
+
26
+ @lru_cache(maxsize=1)
27
+ def get_gigaspeech_pre_trained_model():
28
+ nn_model_filename = hf_hub_download(
29
+ # It is converted from https://huggingface.co/wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2
30
+ repo_id="csukuangfj/icefall-asr-gigaspeech-pruned-transducer-stateless2",
31
+ filename="cpu_jit-epoch-29-avg-11-torch-1.10.0.pt",
32
+ subfolder="exp",
33
+ )
34
+
35
+ bpe_model_filename = hf_hub_download(
36
+ repo_id="wgb14/icefall-asr-gigaspeech-pruned-transducer-stateless2",
37
+ filename="bpe.model",
38
+ subfolder="data/lang_bpe_500",
39
+ )
40
+
41
+ return OfflineAsr(
42
+ nn_model_filename=nn_model_filename,
43
+ bpe_model_filename=bpe_model_filename,
44
+ token_filename=None,
45
+ decoding_method="greedy_search",
46
+ num_active_paths=4,
47
+ sample_rate=sample_rate,
48
+ device="cpu",
49
+ )
50
+
51
+
52
+ @lru_cache(maxsize=1)
53
+ def get_wenetspeech_pre_trained_model():
54
+ nn_model_filename = hf_hub_download(
55
+ repo_id="luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
56
+ filename="cpu_jit_epoch_10_avg_2_torch_1.7.1.pt",
57
+ subfolder="exp",
58
+ )
59
+
60
+ token_filename = hf_hub_download(
61
+ repo_id="luomingshuang/icefall_asr_wenetspeech_pruned_transducer_stateless2",
62
+ filename="tokens.txt",
63
+ subfolder="data/lang_char",
64
+ )
65
+
66
+ return OfflineAsr(
67
+ nn_model_filename=nn_model_filename,
68
+ bpe_model_filename=None,
69
+ token_filename=token_filename,
70
+ decoding_method="greedy_search",
71
+ num_active_paths=4,
72
+ sample_rate=sample_rate,
73
+ device="cpu",
74
+ )
offline_asr.py ADDED
@@ -0,0 +1,419 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
3
+ #
4
+ # Copied from https://github.com/k2-fsa/sherpa/blob/master/sherpa/bin/conformer_rnnt/offline_asr.py
5
+ #
6
+ # See LICENSE for clarification regarding multiple authors
7
+ #
8
+ # Licensed under the Apache License, Version 2.0 (the "License");
9
+ # you may not use this file except in compliance with the License.
10
+ # You may obtain a copy of the License at
11
+ #
12
+ # http://www.apache.org/licenses/LICENSE-2.0
13
+ #
14
+ # Unless required by applicable law or agreed to in writing, software
15
+ # distributed under the License is distributed on an "AS IS" BASIS,
16
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
+ # See the License for the specific language governing permissions and
18
+ # limitations under the License.
19
+ """
20
+ A standalone script for offline ASR recognition.
21
+
22
+ It loads a torchscript model, decodes the given wav files, and exits.
23
+
24
+ Usage:
25
+ ./offline_asr.py --help
26
+
27
+ For BPE based models (e.g., LibriSpeech):
28
+
29
+ ./offline_asr.py \
30
+ --nn-model-filename /path/to/cpu_jit.pt \
31
+ --bpe-model-filename /path/to/bpe.model \
32
+ --decoding-method greedy_search \
33
+ ./foo.wav \
34
+ ./bar.wav \
35
+ ./foobar.wav
36
+
37
+ For character based models (e.g., aishell):
38
+
39
+ ./offline.py \
40
+ --nn-model-filename /path/to/cpu_jit.pt \
41
+ --token-filename /path/to/lang_char/tokens.txt \
42
+ --decoding-method greedy_search \
43
+ ./foo.wav \
44
+ ./bar.wav \
45
+ ./foobar.wav
46
+
47
+ Note: We provide pre-trained models for testing.
48
+
49
+ (1) Pre-trained model with the LibriSpeech dataset
50
+
51
+ sudo apt-get install git-lfs
52
+ git lfs install
53
+ git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
54
+
55
+ nn_model_filename=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/exp/cpu_jit-torch-1.6.0.pt
56
+ bpe_model=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/data/lang_bpe_500/bpe.model
57
+
58
+ wav1=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/test_wavs/1089-134686-0001.wav
59
+ wav2=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/test_wavs/1221-135766-0001.wav
60
+ wav3=./icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/test_wavs/1221-135766-0002.wav
61
+
62
+ sherpa/bin/conformer_rnnt/offline_asr.py \
63
+ --nn-model-filename $nn_model_filename \
64
+ --bpe-model $bpe_model \
65
+ $wav1 \
66
+ $wav2 \
67
+ $wav3
68
+
69
+ (2) Pre-trained model with the aishell dataset
70
+
71
+ sudo apt-get install git-lfs
72
+ git lfs install
73
+ git clone https://huggingface.co/csukuangfj/icefall-aishell-pruned-transducer-stateless3-2022-06-20
74
+
75
+ nn_model_filename=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/exp/cpu_jit-epoch-29-avg-5-torch-1.6.0.pt
76
+ token_filename=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/data/lang_char/tokens.txt
77
+
78
+ wav1=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/test_wavs/BAC009S0764W0121.wav
79
+ wav2=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/test_wavs/BAC009S0764W0122.wav
80
+ wav3=./icefall-aishell-pruned-transducer-stateless3-2022-06-20/test_wavs/BAC009S0764W0123.wav
81
+
82
+ sherpa/bin/conformer_rnnt/offline_asr.py \
83
+ --nn-model-filename $nn_model_filename \
84
+ --token-filename $token_filename \
85
+ $wav1 \
86
+ $wav2 \
87
+ $wav3
88
+ """
89
+ import argparse
90
+ import functools
91
+ import logging
92
+ from typing import List, Optional, Union
93
+
94
+ import k2
95
+ import kaldifeat
96
+ import sentencepiece as spm
97
+ import torch
98
+ import torchaudio
99
+ from sherpa import RnntConformerModel
100
+
101
+ from decode import run_model_and_do_greedy_search, run_model_and_do_modified_beam_search
102
+
103
+
104
+ def get_args():
105
+ parser = argparse.ArgumentParser(
106
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
107
+ )
108
+
109
+ parser.add_argument(
110
+ "--nn-model-filename",
111
+ type=str,
112
+ help="""The torchscript model. You can use
113
+ icefall/egs/librispeech/ASR/pruned_transducer_statelessX/export.py \
114
+ --jit=1
115
+ to generate this model.
116
+ """,
117
+ )
118
+
119
+ parser.add_argument(
120
+ "--bpe-model-filename",
121
+ type=str,
122
+ help="""The BPE model
123
+ You can find it in the directory egs/librispeech/ASR/data/lang_bpe_xxx
124
+ from icefall,
125
+ where xxx is the number of BPE tokens you used to train the model.
126
+ Note: Use it only when your model is using BPE. You don't need to
127
+ provide it if you provide `--token-filename`
128
+ """,
129
+ )
130
+
131
+ parser.add_argument(
132
+ "--token-filename",
133
+ type=str,
134
+ help="""Filename for tokens.txt
135
+ You can find it in the directory
136
+ egs/aishell/ASR/data/lang_char/tokens.txt from icefall.
137
+ Note: You don't need to provide it if you provide `--bpe-model`
138
+ """,
139
+ )
140
+
141
+ parser.add_argument(
142
+ "--decoding-method",
143
+ type=str,
144
+ default="greedy_search",
145
+ help="""Decoding method to use. Currently, only greedy_search and
146
+ modified_beam_search are implemented.
147
+ """,
148
+ )
149
+
150
+ parser.add_argument(
151
+ "--num-active-paths",
152
+ type=int,
153
+ default=4,
154
+ help="""Used only when decoding_method is modified_beam_search.
155
+ It specifies number of active paths for each utterance. Due to
156
+ merging paths with identical token sequences, the actual number
157
+ may be less than "num_active_paths".
158
+ """,
159
+ )
160
+
161
+ parser.add_argument(
162
+ "--sample-rate",
163
+ type=int,
164
+ default=16000,
165
+ help="The expected sample rate of the input sound files",
166
+ )
167
+
168
+ parser.add_argument(
169
+ "sound_files",
170
+ type=str,
171
+ nargs="+",
172
+ help="The input sound file(s) to transcribe. "
173
+ "Supported formats are those supported by torchaudio.load(). "
174
+ "For example, wav and flac are supported. "
175
+ "The sample rate has to equal to `--sample-rate`.",
176
+ )
177
+
178
+ return parser.parse_args()
179
+
180
+
181
+ def read_sound_files(
182
+ filenames: List[str],
183
+ expected_sample_rate: int,
184
+ ) -> List[torch.Tensor]:
185
+ """Read a list of sound files into a list 1-D float32 torch tensors.
186
+ Args:
187
+ filenames:
188
+ A list of sound filenames.
189
+ expected_sample_rate:
190
+ The expected sample rate of the sound files.
191
+ Returns:
192
+ Return a list of 1-D float32 torch tensors.
193
+ """
194
+ ans = []
195
+ for f in filenames:
196
+ wave, sample_rate = torchaudio.load(f)
197
+ assert sample_rate == expected_sample_rate, (
198
+ f"expected sample rate: {expected_sample_rate}. " f"Given: {sample_rate}"
199
+ )
200
+ # We use only the first channel
201
+ ans.append(wave[0])
202
+ return ans
203
+
204
+
205
+ class OfflineAsr(object):
206
+ def __init__(
207
+ self,
208
+ nn_model_filename: str,
209
+ bpe_model_filename: Optional[str],
210
+ token_filename: Optional[str],
211
+ decoding_method: str,
212
+ num_active_paths: int,
213
+ sample_rate: int = 16000,
214
+ device: Union[str, torch.device] = "cpu",
215
+ ):
216
+ """
217
+ Args:
218
+ nn_model_filename:
219
+ Path to the torch script model.
220
+ bpe_model_filename:
221
+ Path to the BPE model. If it is None, you have to provide
222
+ `token_filename`.
223
+ token_filename:
224
+ Path to tokens.txt. If it is None, you have to provide
225
+ `bpe_model_filename`.
226
+ decoding_method:
227
+ The decoding method to use. Currently, only greedy_search and
228
+ modified_beam_search are implemented.
229
+ num_active_paths:
230
+ Used only when decoding_method is modified_beam_search.
231
+ It specifies number of active paths for each utterance. Due to
232
+ merging paths with identical token sequences, the actual number
233
+ may be less than "num_active_paths".
234
+ sample_rate:
235
+ Expected sample rate of the feature extractor.
236
+ device:
237
+ The device to use for computation.
238
+ """
239
+ self.model = RnntConformerModel(
240
+ filename=nn_model_filename,
241
+ device=device,
242
+ optimize_for_inference=False,
243
+ )
244
+
245
+ if bpe_model_filename:
246
+ self.sp = spm.SentencePieceProcessor()
247
+ self.sp.load(bpe_model_filename)
248
+ else:
249
+ self.token_table = k2.SymbolTable.from_file(token_filename)
250
+
251
+ self.feature_extractor = self._build_feature_extractor(
252
+ sample_rate=sample_rate,
253
+ device=device,
254
+ )
255
+
256
+ assert decoding_method in (
257
+ "greedy_search",
258
+ "modified_beam_search",
259
+ ), decoding_method
260
+ if decoding_method == "greedy_search":
261
+ nn_and_decoding_func = run_model_and_do_greedy_search
262
+ elif decoding_method == "modified_beam_search":
263
+ nn_and_decoding_func = functools.partial(
264
+ run_model_and_do_modified_beam_search,
265
+ num_active_paths=num_active_paths,
266
+ )
267
+ else:
268
+ raise ValueError(
269
+ f"Unsupported decoding_method: {decoding_method} "
270
+ "Please use greedy_search or modified_beam_search"
271
+ )
272
+
273
+ self.nn_and_decoding_func = nn_and_decoding_func
274
+ self.device = device
275
+
276
+ def _build_feature_extractor(
277
+ self,
278
+ sample_rate: int = 16000,
279
+ device: Union[str, torch.device] = "cpu",
280
+ ) -> kaldifeat.OfflineFeature:
281
+ """Build a fbank feature extractor for extracting features.
282
+
283
+ Args:
284
+ sample_rate:
285
+ Expected sample rate of the feature extractor.
286
+ device:
287
+ The device to use for computation.
288
+ Returns:
289
+ Return a fbank feature extractor.
290
+ """
291
+ opts = kaldifeat.FbankOptions()
292
+ opts.device = device
293
+ opts.frame_opts.dither = 0
294
+ opts.frame_opts.snip_edges = False
295
+ opts.frame_opts.samp_freq = sample_rate
296
+ opts.mel_opts.num_bins = 80
297
+
298
+ fbank = kaldifeat.Fbank(opts)
299
+
300
+ return fbank
301
+
302
+ def decode_waves(self, waves: List[torch.Tensor]) -> List[List[str]]:
303
+ """
304
+ Args:
305
+ waves:
306
+ A list of 1-D torch.float32 tensors containing audio samples.
307
+ wavs[i] contains audio samples for the i-th utterance.
308
+
309
+ Note:
310
+ Whether it should be in the range [-32768, 32767] or be normalized
311
+ to [-1, 1] depends on which range you used for your training data.
312
+ For instance, if your training data used [-32768, 32767],
313
+ then the given waves have to contain samples in this range.
314
+
315
+ All models trained in icefall use the normalized range [-1, 1].
316
+ Returns:
317
+ Return a list of decoded results. `ans[i]` contains the decoded
318
+ results for `wavs[i]`.
319
+ """
320
+ waves = [w.to(self.device) for w in waves]
321
+ features = self.feature_extractor(waves)
322
+
323
+ tokens = self.nn_and_decoding_func(self.model, features)
324
+
325
+ if hasattr(self, "sp"):
326
+ results = self.sp.decode(tokens)
327
+ else:
328
+ results = [[self.token_table[i] for i in hyp] for hyp in tokens]
329
+ results = ["".join(r) for r in results]
330
+
331
+ return results
332
+
333
+
334
+ @torch.no_grad()
335
+ def main():
336
+ args = get_args()
337
+ logging.info(vars(args))
338
+
339
+ nn_model_filename = args.nn_model_filename
340
+ bpe_model_filename = args.bpe_model_filename
341
+ token_filename = args.token_filename
342
+ decoding_method = args.decoding_method
343
+ num_active_paths = args.num_active_paths
344
+ sample_rate = args.sample_rate
345
+ sound_files = args.sound_files
346
+
347
+ assert decoding_method in ("greedy_search", "modified_beam_search"), decoding_method
348
+
349
+ if decoding_method == "modified_beam_search":
350
+ assert num_active_paths >= 1, num_active_paths
351
+
352
+ if bpe_model_filename:
353
+ assert token_filename is None
354
+
355
+ if token_filename:
356
+ assert bpe_model_filename is None
357
+
358
+ device = torch.device("cpu")
359
+ if torch.cuda.is_available():
360
+ device = torch.device("cuda", 0)
361
+
362
+ logging.info(f"device: {device}")
363
+
364
+ offline_asr = OfflineAsr(
365
+ nn_model_filename=nn_model_filename,
366
+ bpe_model_filename=bpe_model_filename,
367
+ token_filename=token_filename,
368
+ decoding_method=decoding_method,
369
+ num_active_paths=num_active_paths,
370
+ sample_rate=sample_rate,
371
+ device=device,
372
+ )
373
+
374
+ waves = read_sound_files(
375
+ filenames=sound_files,
376
+ expected_sample_rate=sample_rate,
377
+ )
378
+
379
+ logging.info("Decoding started.")
380
+
381
+ hyps = offline_asr.decode_waves(waves)
382
+
383
+ s = "\n"
384
+ for filename, hyp in zip(sound_files, hyps):
385
+ s += f"{filename}:\n{hyp}\n\n"
386
+ logging.info(s)
387
+
388
+ logging.info("Decoding done.")
389
+
390
+
391
+ torch.set_num_threads(1)
392
+ torch.set_num_interop_threads(1)
393
+
394
+ # See https://github.com/pytorch/pytorch/issues/38342
395
+ # and https://github.com/pytorch/pytorch/issues/33354
396
+ #
397
+ # If we don't do this, the delay increases whenever there is
398
+ # a new request that changes the actual batch size.
399
+ # If you use `py-spy dump --pid <server-pid> --native`, you will
400
+ # see a lot of time is spent in re-compiling the torch script model.
401
+ torch._C._jit_set_profiling_executor(False)
402
+ torch._C._jit_set_profiling_mode(False)
403
+ torch._C._set_graph_executor_optimize(False)
404
+ """
405
+ // Use the following in C++
406
+ torch::jit::getExecutorMode() = false;
407
+ torch::jit::getProfilingMode() = false;
408
+ torch::jit::setGraphExecutorOptimize(false);
409
+ """
410
+
411
+ if __name__ == "__main__":
412
+ torch.manual_seed(20220609)
413
+
414
+ formatter = (
415
+ "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" # noqa
416
+ )
417
+ logging.basicConfig(format=formatter, level=logging.INFO)
418
+
419
+ main()
requirements.txt ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ https://download.pytorch.org/whl/cpu/torch-1.10.0%2Bcpu-cp38-cp38-linux_x86_64.whl
2
+ https://k2-fsa.org/nightly/whl/k2-1.17.dev20220711+cpu.torch1.10.0-cp38-cp38-linux_x86_64.whl
3
+ https://download.pytorch.org/whl/cpu/torchaudio-0.10.0%2Bcpu-cp38-cp38-linux_x86_64.whl
4
+
5
+
6
+ https://huggingface.co/csukuangfj/wheels/resolve/main/kaldifeat-1.17-cp38-cp38-linux_x86_64.whl
7
+ https://huggingface.co/csukuangfj/wheels/resolve/main/k2_sherpa-0.6-cp38-cp38-linux_x86_64.whl
8
+
9
+
10
+ sentencepiece>=0.1.96
11
+ numpy
12
+
13
+ huggingface_hub