csukuangfj
commited on
Commit
•
071812d
1
Parent(s):
99b54d7
Add ASR demo with Next-gen Kaldi
Browse files- app.py +153 -0
- decode.py +121 -0
- model.py +74 -0
- offline_asr.py +419 -0
- 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
|