Skip to content

Wav2Vec Aligner

Item dataclass

A data class that represents an item with a sentence, a path to a wav file, and an output path.

Source code in training/preprocess/wav2vec_aligner.py 
15
16
17
18
19
20
21
22
23
@dataclass
class Item:
    r"""A data class that represents an item with a sentence,
    a path to a wav file, and an output path.
    """

    sent: str
    wav_path: str
    out_path: str

Point dataclass

A data class that represents a point with a token index, a time index, and a score.

Source code in training/preprocess/wav2vec_aligner.py 
26
27
28
29
30
31
32
33
34
@dataclass
class Point:
    r"""A data class that represents a point with a token index,
    a time index, and a score.
    """

    token_index: int
    time_index: int
    score: float

Segment dataclass

A data class that represents a segment with a label, a start time, an end time, a duration, and a score.

Source code in training/preprocess/wav2vec_aligner.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
@dataclass
class Segment:
    r"""A data class that represents a segment with a label,
    a start time, an end time, a duration, and a score.
    """

    label: str
    start: int
    end: int
    # TODO: check that the scale of duration is correct...
    duration: float
    score: float

    @property
    def length(self):
        return self.end - self.start

Wav2VecAligner

Bases: Module

Wav2VecAligner model.

The Wav2VecAligner model is designed for aligning audio data with text data. This class handles the training and validation of the Wav2VecAligner model.

Attributes config (AutoConfig): The configuration for the pre-trained model. model (AutoModelForCTC): The pre-trained model. processor (AutoProcessor): The processor for the pre-trained model. labels (List): The labels from the vocabulary of the tokenizer. blank_id (int): The ID of the blank token.

Methods load_audio: Load an audio file from the specified path. encode: Encode the labels. decode: Decode the tokens. align_single_sample: Align a single sample of audio data with the corresponding text. get_trellis: Build a trellis matrix that represents the probabilities of each source token being at a certain time step. backtrack: Walk backwards from the last merge_repeats: Merge repeated tokens into a single segment. merge_words: Merge words in the given path. forward: Perform the forward pass of the model, which involves loading the audio data, aligning the audio with the text, building the trellis, backtracking to find the optimal path, merging repeated tokens, and finally merging words.

Source code in training/preprocess/wav2vec_aligner.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
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
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
class Wav2VecAligner(Module):
    r"""Wav2VecAligner model.

    The Wav2VecAligner model is designed for aligning audio data with text data.
    This class handles the training and validation of the Wav2VecAligner model.

    Attributes
        config (AutoConfig): The configuration for the pre-trained model.
        model (AutoModelForCTC): The pre-trained model.
        processor (AutoProcessor): The processor for the pre-trained model.
        labels (List): The labels from the vocabulary of the tokenizer.
        blank_id (int): The ID of the blank token.

    Methods
        load_audio: Load an audio file from the specified path.
        encode: Encode the labels.
        decode: Decode the tokens.
        align_single_sample: Align a single sample of audio data with the corresponding text.
        get_trellis: Build a trellis matrix that represents the probabilities of each source token being at a certain time step.
        backtrack: Walk backwards from the last
        merge_repeats: Merge repeated tokens into a single segment.
        merge_words: Merge words in the given path.
        forward: Perform the forward pass of the model, which involves loading the audio data, aligning the audio with the text, building the trellis, backtracking to find the optimal path, merging repeated tokens, and finally merging words.

    """

    def __init__(
        self,
        model_name: str = "facebook/wav2vec2-base-960h",
    ):
        r"""Initialize a new instance of the Wav2VecAligner class.

        Args:
            model_name (str): The name of the pre-trained model to use. Defaults to "facebook/wav2vec2-base-960h".
        """
        super().__init__()

        # Load the config
        self.config = AutoConfig.from_pretrained(model_name)

        self.model = AutoModelForCTC.from_pretrained(model_name)
        self.model.eval()

        self.processor = AutoProcessor.from_pretrained(model_name)

        # get labels from vocab
        self.labels = list(self.processor.tokenizer.get_vocab().keys())

        for i in range(len(self.labels)):
            if self.labels[i] == "[PAD]" or self.labels[i] == "<pad>":
                self.blank_id = i

        print("Blank Token id [PAD]/<pad>", self.blank_id)

    def load_audio(self, wav_path: str) -> Tuple[torch.Tensor, int]:
        r"""Load an audio file from the specified path.

        Args:
            wav_path (str): The path to the audio file.

        Returns:
            Tuple[torch.Tensor, int]: A tuple containing the loaded audio data and the sample rate, or a FileNotFoundError if the file does not exist.
        """
        if not os.path.isfile(wav_path):
            raise FileNotFoundError(wav_path, "Not found in wavs directory")

        speech_array, sampling_rate = torchaudio.load(wav_path) # type: ignore
        return speech_array, sampling_rate

    def encode(self, text: str) -> List:
        # encode labels
        with self.processor.as_target_processor():
            return self.processor(text, return_tensors="pt").input_ids

    def decode(self, tokens: List):
        # Decode tokens
        decoded = self.processor.batch_decode(tokens)
        return decoded[0]

    def align_single_sample(
        self, audio_input: torch.Tensor, text: str,
    ) -> Tuple[torch.Tensor, List, str]:
        r"""Align a single sample of audio data with the corresponding text.

        Args:
            audio_input (torch.Tensor): The audio data.
            text (str): The corresponding text.

        Returns:
            Tuple[torch.Tensor, List, str]: A tuple containing the emissions, the tokens, and the transcript.
        """
        transcript = "|".join(text.split(" "))
        transcript = f"|{transcript}|"

        with torch.inference_mode():
            logits = self.model(audio_input).logits

        # Get the emission probability at frame level
        # Compute the probability in log-domain to avoid numerical instability
        # For this purpose, we normalize the emission with `torch.log_softmax()`
        emissions = torch.log_softmax(logits, dim=-1)
        emissions = emissions[0]

        tokens = self.encode(transcript)[0]

        return emissions, tokens, transcript

    def get_trellis(
        self,
        emission: torch.Tensor,
        tokens: List,
    ) -> torch.Tensor:
        r"""Build a trellis matrix of shape (num_frames + 1, num_tokens + 1)
        that represents the probabilities of each source token being at a certain time step.

        Since we are looking for the most likely transitions, we take the more likely path for the value of $k_{(t+1,j+1)}$, that is:

        $k_{t+1, j+1} = \max(k_{t, j} p_{t+1, c_{j+1}}, k_{t, j+1} p_{t+1, \text{repeat}})$

        Args:
            emission (torch.Tensor): The emission tensor.
            tokens (List): The list of tokens.

        Returns:
            torch.Tensor: The trellis matrix.
        """
        num_frames = emission.size(0)
        num_tokens = len(tokens)

        # Trellis has extra diemsions for both time axis and tokens.
        # The extra dim for tokens represents <SoS> (start-of-sentence)
        # The extra dim for time axis is for simplification of the code.

        trellis = torch.zeros((num_frames, num_tokens))
        trellis[1:, 0] = torch.cumsum(emission[1:, self.blank_id], 0)
        trellis[0, 1:] = -float("inf")
        trellis[-num_tokens + 1 :, 0] = float("inf")

        for t in range(num_frames - 1):
            trellis[t + 1, 1:] = torch.maximum(
                # Score for staying at the same token
                trellis[t, 1:] + emission[t, self.blank_id],
                # Score for changing to the next token
                trellis[t, :-1] + emission[t, tokens[1:]],
            )
        return trellis

    def backtrack(
        self,
        trellis: torch.Tensor,
        emission: torch.Tensor,
        tokens: List,
    ) -> List[Point]:
        r"""Walk backwards from the last (sentence_token, time_step) pair to build the optimal sequence alignment path.

        Args:
            trellis (torch.Tensor): The trellis matrix.
            emission (torch.Tensor): The emission tensor.
            tokens (List): The list of tokens.

        Returns:
            List[Point]: The optimal sequence alignment path.
        """
        # Note:
        # j and t are indices for trellis, which has extra dimensions
        # for time and tokens at the beginning.
        # When referring to time frame index `T` in trellis,
        # the corresponding index in emission is `T-1`.
        # Similarly, when referring to token index `J` in trellis,
        # the corresponding index in transcript is `J-1`.
        t, j = trellis.size(0) - 1, trellis.size(1) - 1

        path = [Point(j, t, emission[t, self.blank_id].exp().item())]
        while j > 0:
            # Should not happen but just in case
            assert t > 0

            # 1. Figure out if the current position was stay or change
            # Frame-wise score of stay vs change
            p_stay = emission[t - 1, self.blank_id]
            p_change = emission[t - 1, tokens[j]]

            # Context-aware score for stay vs change
            stayed = trellis[t - 1, j] + p_stay
            changed = trellis[t - 1, j - 1] + p_change

            # Update position
            t -= 1
            if changed > stayed:
                j -= 1

            # Store the path with frame-wise probability.
            prob = (p_change if changed > stayed else p_stay).exp().item()
            path.append(Point(j, t, prob))

        # Now j == 0, which means, it reached the SoS.
        # Fill up the rest for the sake of visualization
        while t > 0:
            prob = emission[t - 1, self.blank_id].exp().item()
            path.append(Point(j, t - 1, prob))
            t -= 1

        return path[::-1]

    def merge_repeats(self, path: List[Point], transcript: str) -> List[Segment]:
        r"""Merge repeated tokens into a single segment.

        Args:
            path (List[Point]): The sequence alignment path.
            transcript (str): The transcript.

        Returns:
            List[Segment]: The list of segments.

        Note: this shouldn't affect repeated characters from the
        original sentences (e.g. `ll` in `hello`)
        """
        i1, i2 = 0, 0
        segments = []
        while i1 < len(path):
            while i2 < len(path) and path[i1].token_index == path[i2].token_index:
                i2 += 1
            score = sum(path[k].score for k in range(i1, i2)) / (i2 - i1)

            x0, x1 = path[i1].time_index, path[i2 - 1].time_index + 1
            duration = x1 - x0

            segments.append(
                Segment(
                    transcript[path[i1].token_index],
                    x0,
                    x1,
                    duration,
                    score,
                ),
            )
            i1 = i2
        return segments

    # Merge words
    def merge_words(
        self, segments: List[Segment], separator: str = "|",
    ) -> List[Segment]:
        r"""Merge words in the given path.

        Args:
            segments (List[Segment]): The list of segments.
            separator (str): The separator character. Defaults to "|".

        Returns:
            List[Segment]: The list of merged words.
        """
        words = []
        i1, i2 = 0, 0
        while i1 < len(segments):
            if i2 >= len(segments) or segments[i2].label == separator:
                if i1 != i2:
                    segs = segments[i1:i2]
                    word = "".join([seg.label for seg in segs])
                    score = sum(seg.score * seg.length for seg in segs) / sum(
                        seg.length for seg in segs
                    )

                    x0, x1 = segments[i1].start, segments[i2 - 1].end
                    duration = x1 - x0

                    words.append(Segment(word, x0, x1, duration, score))
                i1 = i2 + 1
                i2 = i1
            else:
                i2 += 1
        return words

    def forward(self, wav_path: str, text: str) -> List[Segment]:
        r"""Perform the forward pass of the model, which involves loading the audio data, aligning the audio with the text,
        building the trellis, backtracking to find the optimal path, merging repeated tokens, and finally merging words.

        Args:
            wav_path (str): The path to the audio file.
            text (str): The corresponding text.

        Returns:
            List[Segment]: The list of segments representing the alignment of the audio data with the text.
        """
        audio_input, _ = self.load_audio(wav_path)

        emissions, tokens, transcript = self.align_single_sample(audio_input, text)

        trellis = self.get_trellis(emissions, tokens)

        path = self.backtrack(trellis, emissions, tokens)

        merged_path = self.merge_repeats(path, transcript)

        return self.merge_words(merged_path)

    def save_segments(self, wav_path: str, text: str, save_dir: str):
        r"""Perform the forward pass of the model to get the segments and save each segment to a file.
        Used for debugging purposes.

        Args:
            wav_path (str): The path to the audio file.
            text (str): The corresponding text.
            save_dir (str): The directory where the audio files should be saved.

        Returns:
            None
        """
        word_segments = self.forward(wav_path, text)

        waveform, sampling_rate = self.load_audio(wav_path)

        emissions, tokens, _ = self.align_single_sample(waveform, text)

        trellis = self.get_trellis(emissions, tokens)

        ratio = waveform.size(1) / trellis.size(0)

        for i, word in enumerate(word_segments):
            x0 = int(ratio * word.start)
            x1 = int(ratio * word.end)

            print(
                f"{word.label} ({word.score:.2f}): {x0 / sampling_rate:.3f} - {x1 / sampling_rate:.3f} sec",
            )

            segment_waveform = waveform[:, x0:x1]

            # Save the segment waveform to a file
            filename = f"{i}_{word.label}.wav"
            torchaudio.save( # type: ignore
                os.path.join(save_dir, filename), segment_waveform, sampling_rate,
            )

__init__(model_name='facebook/wav2vec2-base-960h')

Initialize a new instance of the Wav2VecAligner class.

Parameters:

Name Type Description Default
model_name str

The name of the pre-trained model to use. Defaults to "facebook/wav2vec2-base-960h".

 'facebook/wav2vec2-base-960h'
Source code in training/preprocess/wav2vec_aligner.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def __init__(
    self,
    model_name: str = "facebook/wav2vec2-base-960h",
):
    r"""Initialize a new instance of the Wav2VecAligner class.

    Args:
        model_name (str): The name of the pre-trained model to use. Defaults to "facebook/wav2vec2-base-960h".
    """
    super().__init__()

    # Load the config
    self.config = AutoConfig.from_pretrained(model_name)

    self.model = AutoModelForCTC.from_pretrained(model_name)
    self.model.eval()

    self.processor = AutoProcessor.from_pretrained(model_name)

    # get labels from vocab
    self.labels = list(self.processor.tokenizer.get_vocab().keys())

    for i in range(len(self.labels)):
        if self.labels[i] == "[PAD]" or self.labels[i] == "<pad>":
            self.blank_id = i

    print("Blank Token id [PAD]/<pad>", self.blank_id)

align_single_sample(audio_input, text)

Align a single sample of audio data with the corresponding text.

Parameters:

Name Type Description Default
audio_input Tensor

The audio data.

 required
text str

The corresponding text.

 required

Returns:

Type Description
Tuple[Tensor, List, str]

Tuple[torch.Tensor, List, str]: A tuple containing the emissions, the tokens, and the transcript.
Source code in training/preprocess/wav2vec_aligner.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def align_single_sample(
    self, audio_input: torch.Tensor, text: str,
) -> Tuple[torch.Tensor, List, str]:
    r"""Align a single sample of audio data with the corresponding text.

    Args:
        audio_input (torch.Tensor): The audio data.
        text (str): The corresponding text.

    Returns:
        Tuple[torch.Tensor, List, str]: A tuple containing the emissions, the tokens, and the transcript.
    """
    transcript = "|".join(text.split(" "))
    transcript = f"|{transcript}|"

    with torch.inference_mode():
        logits = self.model(audio_input).logits

    # Get the emission probability at frame level
    # Compute the probability in log-domain to avoid numerical instability
    # For this purpose, we normalize the emission with `torch.log_softmax()`
    emissions = torch.log_softmax(logits, dim=-1)
    emissions = emissions[0]

    tokens = self.encode(transcript)[0]

    return emissions, tokens, transcript

backtrack(trellis, emission, tokens)

Walk backwards from the last (sentence_token, time_step) pair to build the optimal sequence alignment path.

Parameters:

Name Type Description Default
trellis Tensor

The trellis matrix.

 required
emission Tensor

The emission tensor.

 required
tokens List

The list of tokens.

 required

Returns:

Type Description
List[Point]

List[Point]: The optimal sequence alignment path.
Source code in training/preprocess/wav2vec_aligner.py 
202
203
204
205
206
207
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
251
252
253
254
255
256
257
def backtrack(
    self,
    trellis: torch.Tensor,
    emission: torch.Tensor,
    tokens: List,
) -> List[Point]:
    r"""Walk backwards from the last (sentence_token, time_step) pair to build the optimal sequence alignment path.

    Args:
        trellis (torch.Tensor): The trellis matrix.
        emission (torch.Tensor): The emission tensor.
        tokens (List): The list of tokens.

    Returns:
        List[Point]: The optimal sequence alignment path.
    """
    # Note:
    # j and t are indices for trellis, which has extra dimensions
    # for time and tokens at the beginning.
    # When referring to time frame index `T` in trellis,
    # the corresponding index in emission is `T-1`.
    # Similarly, when referring to token index `J` in trellis,
    # the corresponding index in transcript is `J-1`.
    t, j = trellis.size(0) - 1, trellis.size(1) - 1

    path = [Point(j, t, emission[t, self.blank_id].exp().item())]
    while j > 0:
        # Should not happen but just in case
        assert t > 0

        # 1. Figure out if the current position was stay or change
        # Frame-wise score of stay vs change
        p_stay = emission[t - 1, self.blank_id]
        p_change = emission[t - 1, tokens[j]]

        # Context-aware score for stay vs change
        stayed = trellis[t - 1, j] + p_stay
        changed = trellis[t - 1, j - 1] + p_change

        # Update position
        t -= 1
        if changed > stayed:
            j -= 1

        # Store the path with frame-wise probability.
        prob = (p_change if changed > stayed else p_stay).exp().item()
        path.append(Point(j, t, prob))

    # Now j == 0, which means, it reached the SoS.
    # Fill up the rest for the sake of visualization
    while t > 0:
        prob = emission[t - 1, self.blank_id].exp().item()
        path.append(Point(j, t - 1, prob))
        t -= 1

    return path[::-1]

forward(wav_path, text)

Perform the forward pass of the model, which involves loading the audio data, aligning the audio with the text, building the trellis, backtracking to find the optimal path, merging repeated tokens, and finally merging words.

Parameters:

Name Type Description Default
wav_path str

The path to the audio file.

 required
text str

The corresponding text.

 required

Returns:

Type Description
List[Segment]

List[Segment]: The list of segments representing the alignment of the audio data with the text.
Source code in training/preprocess/wav2vec_aligner.py 
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
def forward(self, wav_path: str, text: str) -> List[Segment]:
    r"""Perform the forward pass of the model, which involves loading the audio data, aligning the audio with the text,
    building the trellis, backtracking to find the optimal path, merging repeated tokens, and finally merging words.

    Args:
        wav_path (str): The path to the audio file.
        text (str): The corresponding text.

    Returns:
        List[Segment]: The list of segments representing the alignment of the audio data with the text.
    """
    audio_input, _ = self.load_audio(wav_path)

    emissions, tokens, transcript = self.align_single_sample(audio_input, text)

    trellis = self.get_trellis(emissions, tokens)

    path = self.backtrack(trellis, emissions, tokens)

    merged_path = self.merge_repeats(path, transcript)

    return self.merge_words(merged_path)

get_trellis(emission, tokens)

Build a trellis matrix of shape (num_frames + 1, num_tokens + 1) that represents the probabilities of each source token being at a certain time step.

Since we are looking for the most likely transitions, we take the more likely path for the value of \(k_{(t+1,j+1)}\), that is:

\(k_{t+1, j+1} = \max(k_{t, j} p_{t+1, c_{j+1}}, k_{t, j+1} p_{t+1, \text{repeat}})\)

Parameters:

Name Type Description Default
emission Tensor

The emission tensor.

 required
tokens List

The list of tokens.

 required

Returns:

Type Description
Tensor

torch.Tensor: The trellis matrix.
Source code in training/preprocess/wav2vec_aligner.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def get_trellis(
    self,
    emission: torch.Tensor,
    tokens: List,
) -> torch.Tensor:
    r"""Build a trellis matrix of shape (num_frames + 1, num_tokens + 1)
    that represents the probabilities of each source token being at a certain time step.

    Since we are looking for the most likely transitions, we take the more likely path for the value of $k_{(t+1,j+1)}$, that is:

    $k_{t+1, j+1} = \max(k_{t, j} p_{t+1, c_{j+1}}, k_{t, j+1} p_{t+1, \text{repeat}})$

    Args:
        emission (torch.Tensor): The emission tensor.
        tokens (List): The list of tokens.

    Returns:
        torch.Tensor: The trellis matrix.
    """
    num_frames = emission.size(0)
    num_tokens = len(tokens)

    # Trellis has extra diemsions for both time axis and tokens.
    # The extra dim for tokens represents <SoS> (start-of-sentence)
    # The extra dim for time axis is for simplification of the code.

    trellis = torch.zeros((num_frames, num_tokens))
    trellis[1:, 0] = torch.cumsum(emission[1:, self.blank_id], 0)
    trellis[0, 1:] = -float("inf")
    trellis[-num_tokens + 1 :, 0] = float("inf")

    for t in range(num_frames - 1):
        trellis[t + 1, 1:] = torch.maximum(
            # Score for staying at the same token
            trellis[t, 1:] + emission[t, self.blank_id],
            # Score for changing to the next token
            trellis[t, :-1] + emission[t, tokens[1:]],
        )
    return trellis

load_audio(wav_path)

Load an audio file from the specified path.

Parameters:

Name Type Description Default
wav_path str

The path to the audio file.

 required

Returns:

Type Description
Tuple[Tensor, int]

Tuple[torch.Tensor, int]: A tuple containing the loaded audio data and the sample rate, or a FileNotFoundError if the file does not exist.
Source code in training/preprocess/wav2vec_aligner.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def load_audio(self, wav_path: str) -> Tuple[torch.Tensor, int]:
    r"""Load an audio file from the specified path.

    Args:
        wav_path (str): The path to the audio file.

    Returns:
        Tuple[torch.Tensor, int]: A tuple containing the loaded audio data and the sample rate, or a FileNotFoundError if the file does not exist.
    """
    if not os.path.isfile(wav_path):
        raise FileNotFoundError(wav_path, "Not found in wavs directory")

    speech_array, sampling_rate = torchaudio.load(wav_path) # type: ignore
    return speech_array, sampling_rate

merge_repeats(path, transcript)

Merge repeated tokens into a single segment.

Parameters:

Name Type Description Default
path List[Point]

The sequence alignment path.

 required
transcript str

The transcript.

 required

Returns:

Type Description
List[Segment]

List[Segment]: The list of segments.

Note: this shouldn't affect repeated characters from the original sentences (e.g. ll in hello)

Source code in training/preprocess/wav2vec_aligner.py 
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
def merge_repeats(self, path: List[Point], transcript: str) -> List[Segment]:
    r"""Merge repeated tokens into a single segment.

    Args:
        path (List[Point]): The sequence alignment path.
        transcript (str): The transcript.

    Returns:
        List[Segment]: The list of segments.

    Note: this shouldn't affect repeated characters from the
    original sentences (e.g. `ll` in `hello`)
    """
    i1, i2 = 0, 0
    segments = []
    while i1 < len(path):
        while i2 < len(path) and path[i1].token_index == path[i2].token_index:
            i2 += 1
        score = sum(path[k].score for k in range(i1, i2)) / (i2 - i1)

        x0, x1 = path[i1].time_index, path[i2 - 1].time_index + 1
        duration = x1 - x0

        segments.append(
            Segment(
                transcript[path[i1].token_index],
                x0,
                x1,
                duration,
                score,
            ),
        )
        i1 = i2
    return segments

merge_words(segments, separator='|')

Merge words in the given path.

Parameters:

Name Type Description Default
segments List[Segment]

The list of segments.

 required
separator str

The separator character. Defaults to "|".

 '|'

Returns:

Type Description
List[Segment]

List[Segment]: The list of merged words.
Source code in training/preprocess/wav2vec_aligner.py 
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
def merge_words(
    self, segments: List[Segment], separator: str = "|",
) -> List[Segment]:
    r"""Merge words in the given path.

    Args:
        segments (List[Segment]): The list of segments.
        separator (str): The separator character. Defaults to "|".

    Returns:
        List[Segment]: The list of merged words.
    """
    words = []
    i1, i2 = 0, 0
    while i1 < len(segments):
        if i2 >= len(segments) or segments[i2].label == separator:
            if i1 != i2:
                segs = segments[i1:i2]
                word = "".join([seg.label for seg in segs])
                score = sum(seg.score * seg.length for seg in segs) / sum(
                    seg.length for seg in segs
                )

                x0, x1 = segments[i1].start, segments[i2 - 1].end
                duration = x1 - x0

                words.append(Segment(word, x0, x1, duration, score))
            i1 = i2 + 1
            i2 = i1
        else:
            i2 += 1
    return words

save_segments(wav_path, text, save_dir)

Perform the forward pass of the model to get the segments and save each segment to a file. Used for debugging purposes.

Parameters:

Name Type Description Default
wav_path str

The path to the audio file.

 required
text str

The corresponding text.

 required
save_dir str

The directory where the audio files should be saved.

 required

Returns:

Type Description

None
Source code in training/preprocess/wav2vec_aligner.py 
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
def save_segments(self, wav_path: str, text: str, save_dir: str):
    r"""Perform the forward pass of the model to get the segments and save each segment to a file.
    Used for debugging purposes.

    Args:
        wav_path (str): The path to the audio file.
        text (str): The corresponding text.
        save_dir (str): The directory where the audio files should be saved.

    Returns:
        None
    """
    word_segments = self.forward(wav_path, text)

    waveform, sampling_rate = self.load_audio(wav_path)

    emissions, tokens, _ = self.align_single_sample(waveform, text)

    trellis = self.get_trellis(emissions, tokens)

    ratio = waveform.size(1) / trellis.size(0)

    for i, word in enumerate(word_segments):
        x0 = int(ratio * word.start)
        x1 = int(ratio * word.end)

        print(
            f"{word.label} ({word.score:.2f}): {x0 / sampling_rate:.3f} - {x1 / sampling_rate:.3f} sec",
        )

        segment_waveform = waveform[:, x0:x1]

        # Save the segment waveform to a file
        filename = f"{i}_{word.label}.wav"
        torchaudio.save( # type: ignore
            os.path.join(save_dir, filename), segment_waveform, sampling_rate,
        )