online_asr_tutorial.py 7.66 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
"""
Online ASR with Emformer RNN-T
==============================

**Author**: `Jeff Hwang <jeffhwang@fb.com>`__, `Moto Hira <moto@fb.com>`__

This tutorial shows how to use Emformer RNN-T and streaming API
to perform online speech recognition.

"""

######################################################################
#
14
# .. note::
15
#
16
#    This tutorial requires FFmpeg libraries (>=4.1, <4.4) and SentencePiece.
17
#
18
#    There are multiple ways to install FFmpeg libraries.
19
#    If you are using Anaconda Python distribution,
20
#    ``conda install 'ffmpeg<4.4'`` will install
21
#    the required FFmpeg libraries.
22
#
23
#    You can install SentencePiece by running ``pip install sentencepiece``.
24

25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
######################################################################
# 1. Overview
# -----------
#
# Performing online speech recognition is composed of the following steps
#
# 1. Build the inference pipeline
#    Emformer RNN-T is composed of three components: feature extractor,
#    decoder and token processor.
# 2. Format the waveform into chunks of expected sizes.
# 3. Pass data through the pipeline.

######################################################################
# 2. Preparation
# --------------
#

42
43
44
import torch
import torchaudio

45
46
47
48
49
50
51
print(torch.__version__)
print(torchaudio.__version__)

######################################################################
#
import IPython

52
try:
53
    from torchaudio.io import StreamReader
54
55
56
57
58
59
except ModuleNotFoundError:
    try:
        import google.colab

        print(
            """
60
61
            To enable running this notebook in Google Colab, install the requisite
            third party libraries by running the following code block:
62
63
64
65
66
67
68
69
70

            !add-apt-repository -y ppa:savoury1/ffmpeg4
            !apt-get -qq install -y ffmpeg
            """
        )
    except ModuleNotFoundError:
        pass
    raise

71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

######################################################################
# 3. Construct the pipeline
# -------------------------
#
# Pre-trained model weights and related pipeline components are
# bundled as :py:func:`torchaudio.pipelines.RNNTBundle`.
#
# We use :py:func:`torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH`,
# which is a Emformer RNN-T model trained on LibriSpeech dataset.
#

bundle = torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH

feature_extractor = bundle.get_streaming_feature_extractor()
decoder = bundle.get_decoder()
token_processor = bundle.get_token_processor()

######################################################################
# Streaming inference works on input data with overlap.
moto's avatar
moto committed
91
92
93
94
95
# Emformer RNN-T model treats the newest portion of the input data
# as the "right context" — a preview of future context.
# In each inference call, the model expects the main segment
# to start from this right context from the previous inference call.
# The following figure illustrates this.
96
97
98
99
100
101
102
103
#
# .. image:: https://download.pytorch.org/torchaudio/tutorial-assets/emformer_rnnt_context.png
#
# The size of main segment and right context, along with
# the expected sample rate can be retrieved from bundle.
#

sample_rate = bundle.sample_rate
104
105
segment_length = bundle.segment_length * bundle.hop_length
context_length = bundle.right_context_length * bundle.hop_length
106
107

print(f"Sample rate: {sample_rate}")
108
109
print(f"Main segment: {segment_length} frames ({segment_length / sample_rate} seconds)")
print(f"Right context: {context_length} frames ({context_length / sample_rate} seconds)")
110
111
112
113
114

######################################################################
# 4. Configure the audio stream
# -----------------------------
#
115
# Next, we configure the input audio stream using :py:func:`~torchaudio.io.StreamReader`.
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#
# For the detail of this API, please refer to the
# `Media Stream API tutorial <./streaming_api_tutorial.html>`__.
#

######################################################################
# The following audio file was originally published by LibriVox project,
# and it is in the public domain.
#
# https://librivox.org/great-pirate-stories-by-joseph-lewis-french/
#
# It was re-uploaded for the sake of the tutorial.
#
src = "https://download.pytorch.org/torchaudio/tutorial-assets/greatpiratestories_00_various.mp3"

131
streamer = StreamReader(src)
132
streamer.add_basic_audio_stream(frames_per_chunk=segment_length, sample_rate=bundle.sample_rate)
133
134
135
136
137

print(streamer.get_src_stream_info(0))
print(streamer.get_out_stream_info(0))

######################################################################
moto's avatar
moto committed
138
139
140
141
142
143
144
145
146
# As previously explained, Emformer RNN-T model expects input data with
# overlaps; however, `Streamer` iterates the source media without overlap,
# so we make a helper structure that caches a part of input data from
# `Streamer` as right context and then appends it to the next input data from
# `Streamer`.
#
# The following figure illustrates this.
#
# .. image:: https://download.pytorch.org/torchaudio/tutorial-assets/emformer_rnnt_streamer_context.png
147
148
149
150
#


class ContextCacher:
moto's avatar
moto committed
151
    """Cache the end of input data and prepend the next input data with it.
152
153

    Args:
154
155
156
        segment_length (int): The size of main segment.
            If the incoming segment is shorter, then the segment is padded.
        context_length (int): The size of the context, cached and appended.
157
158
    """

159
160
161
162
    def __init__(self, segment_length: int, context_length: int):
        self.segment_length = segment_length
        self.context_length = context_length
        self.context = torch.zeros([context_length])
163
164

    def __call__(self, chunk: torch.Tensor):
165
166
167
168
        if chunk.size(0) < self.segment_length:
            chunk = torch.nn.functional.pad(chunk, (0, self.segment_length - chunk.size(0)))
        chunk_with_context = torch.cat((self.context, chunk))
        self.context = chunk[-self.context_length :]
169
170
171
172
173
174
175
176
177
178
179
180
181
182
        return chunk_with_context


######################################################################
# 5. Run stream inference
# -----------------------
#
# Finally, we run the recognition.
#
# First, we initialize the stream iterator, context cacher, and
# state and hypothesis that are used by decoder to carry over the
# decoding state between inference calls.
#

183
cacher = ContextCacher(segment_length, context_length)
184
185
186
187
188
189
190
191
192
193
194

state, hypothesis = None, None

######################################################################
# Next we, run the inference.
#
# For the sake of better display, we create a helper function which
# processes the source stream up to the given times and call it
# repeatedly.
#

moto's avatar
moto committed
195
196
stream_iterator = streamer.stream()

197
198
199
200
201

@torch.inference_mode()
def run_inference(num_iter=200):
    global state, hypothesis
    chunks = []
moto's avatar
moto committed
202
    for i, (chunk,) in enumerate(stream_iterator, start=1):
203
        segment = cacher(chunk[:, 0])
204
205
206
        features, length = feature_extractor(segment)
        hypos, state = decoder.infer(features, length, 10, state=state, hypothesis=hypothesis)
        hypothesis = hypos[0]
207
        transcript = token_processor(hypothesis[0], lstrip=False)
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
        print(transcript, end="", flush=True)

        chunks.append(chunk)
        if i == num_iter:
            break

    return IPython.display.Audio(torch.cat(chunks).T.numpy(), rate=bundle.sample_rate)


######################################################################
#

run_inference()

######################################################################
#

run_inference()

######################################################################
#

run_inference()

######################################################################
#

run_inference()

######################################################################
#

run_inference()

######################################################################
#

run_inference()

######################################################################
#

run_inference()