test_modeling_wav2vec2.py 73.2 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# coding=utf-8
# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Testing suite for the PyTorch Wav2Vec2 model. """

import math
18
import multiprocessing
19
20
21
import os
import pickle
import tempfile
22
import traceback
Patrick von Platen's avatar
Patrick von Platen committed
23
24
import unittest

25
import numpy as np
26
from datasets import load_dataset
27

28
from transformers import Wav2Vec2Config, is_torch_available
29
from transformers.testing_utils import (
30
    CaptureLogger,
31
    is_pt_flax_cross_test,
32
33
34
    is_pyctcdecode_available,
    is_torchaudio_available,
    require_pyctcdecode,
35
36
    require_soundfile,
    require_torch,
37
    require_torchaudio,
38
    run_test_in_subprocess,
39
40
41
    slow,
    torch_device,
)
42
from transformers.utils import is_torch_fx_available
Patrick von Platen's avatar
Patrick von Platen committed
43

Yih-Dar's avatar
Yih-Dar committed
44
45
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
46
47
48
49
50
51
    ModelTesterMixin,
    _config_zero_init,
    floats_tensor,
    ids_tensor,
    random_attention_mask,
)
52
from ...test_pipeline_mixin import PipelineTesterMixin
Patrick von Platen's avatar
Patrick von Platen committed
53
54
55
56
57


if is_torch_available():
    import torch

Anton Lozhkov's avatar
Anton Lozhkov committed
58
59
    from transformers import (
        Wav2Vec2FeatureExtractor,
60
        Wav2Vec2ForAudioFrameClassification,
Anton Lozhkov's avatar
Anton Lozhkov committed
61
62
63
        Wav2Vec2ForCTC,
        Wav2Vec2ForMaskedLM,
        Wav2Vec2ForPreTraining,
64
        Wav2Vec2ForSequenceClassification,
65
        Wav2Vec2ForXVector,
Anton Lozhkov's avatar
Anton Lozhkov committed
66
67
68
        Wav2Vec2Model,
        Wav2Vec2Processor,
    )
69
70
71
72
73
    from transformers.models.wav2vec2.modeling_wav2vec2 import (
        Wav2Vec2GumbelVectorQuantizer,
        _compute_mask_indices,
        _sample_negative_indices,
    )
Patrick von Platen's avatar
Patrick von Platen committed
74
75


76
77
78
79
80
if is_torchaudio_available():
    import torchaudio


if is_pyctcdecode_available():
81
    import pyctcdecode.decoder
82

83
    from transformers import Wav2Vec2ProcessorWithLM
84
    from transformers.models.wav2vec2_with_lm import processing_wav2vec2_with_lm
85
86


87
88
89
90
if is_torch_fx_available():
    from transformers.utils.fx import symbolic_trace


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
def _test_wav2vec2_with_lm_invalid_pool(in_queue, out_queue, timeout):
    error = None
    try:
        _ = in_queue.get(timeout=timeout)

        ds = load_dataset("common_voice", "es", split="test", streaming=True)
        sample = next(iter(ds))

        resampled_audio = torchaudio.functional.resample(
            torch.tensor(sample["audio"]["array"]), 48_000, 16_000
        ).numpy()

        model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm").to(
            torch_device
        )
        processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")

        input_values = processor(resampled_audio, return_tensors="pt").input_values

        with torch.no_grad():
            logits = model(input_values.to(torch_device)).logits

        # use a spawn pool, which should trigger a warning if different than fork
        with CaptureLogger(pyctcdecode.decoder.logger) as cl, multiprocessing.get_context("spawn").Pool(1) as pool:
            transcription = processor.batch_decode(logits.cpu().numpy(), pool).text

        unittest.TestCase().assertIn("Falling back to sequential decoding.", cl.out)
        unittest.TestCase().assertEqual(transcription[0], "bien y qué regalo vas a abrir primero")

        # force batch_decode to internally create a spawn pool, which should trigger a warning if different than fork
        multiprocessing.set_start_method("spawn", force=True)
        with CaptureLogger(processing_wav2vec2_with_lm.logger) as cl:
            transcription = processor.batch_decode(logits.cpu().numpy()).text

        unittest.TestCase().assertIn("Falling back to sequential decoding.", cl.out)
        unittest.TestCase().assertEqual(transcription[0], "bien y qué regalo vas a abrir primero")
    except Exception:
        error = f"{traceback.format_exc()}"

    results = {"error": error}
    out_queue.put(results, timeout=timeout)
    out_queue.join()


Patrick von Platen's avatar
Patrick von Platen committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
class Wav2Vec2ModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=1024,  # speech is longer
        is_training=False,
        hidden_size=16,
        feat_extract_norm="group",
        feat_extract_dropout=0.0,
        feat_extract_activation="gelu",
        conv_dim=(32, 32, 32),
        conv_stride=(4, 4, 4),
        conv_kernel=(8, 8, 8),
        conv_bias=False,
        num_conv_pos_embeddings=16,
        num_conv_pos_embedding_groups=2,
        num_hidden_layers=4,
        num_attention_heads=2,
        hidden_dropout_prob=0.1,  # this is most likely not correctly set yet
        intermediate_size=20,
        layer_norm_eps=1e-5,
        hidden_act="gelu",
        initializer_range=0.02,
159
160
        mask_time_prob=0.5,
        mask_time_length=2,
Patrick von Platen's avatar
Patrick von Platen committed
161
162
        vocab_size=32,
        do_stable_layer_norm=False,
163
164
        num_adapter_layers=1,
        adapter_stride=2,
165
166
167
168
        tdnn_dim=(32, 32),
        tdnn_kernel=(5, 3),
        tdnn_dilation=(1, 2),
        xvector_output_dim=32,
Patrick von Platen's avatar
Patrick von Platen committed
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
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.hidden_size = hidden_size
        self.feat_extract_norm = feat_extract_norm
        self.feat_extract_dropout = feat_extract_dropout
        self.feat_extract_activation = feat_extract_activation
        self.conv_dim = conv_dim
        self.conv_stride = conv_stride
        self.conv_kernel = conv_kernel
        self.conv_bias = conv_bias
        self.num_conv_pos_embeddings = num_conv_pos_embeddings
        self.num_conv_pos_embedding_groups = num_conv_pos_embedding_groups
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_dropout_prob = hidden_dropout_prob
        self.intermediate_size = intermediate_size
        self.layer_norm_eps = layer_norm_eps
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.vocab_size = vocab_size
        self.do_stable_layer_norm = do_stable_layer_norm
194
195
        self.num_adapter_layers = num_adapter_layers
        self.adapter_stride = adapter_stride
196
197
        self.mask_time_prob = mask_time_prob
        self.mask_time_length = mask_time_length
Patrick von Platen's avatar
Patrick von Platen committed
198
        self.scope = scope
199
200
201
202
        self.tdnn_dim = tdnn_dim
        self.tdnn_kernel = tdnn_kernel
        self.tdnn_dilation = tdnn_dilation
        self.xvector_output_dim = xvector_output_dim
Patrick von Platen's avatar
Patrick von Platen committed
203
204
205
206
207
208
209

        output_seq_length = self.seq_length
        for kernel, stride in zip(self.conv_kernel, self.conv_stride):
            output_seq_length = (output_seq_length - (kernel - 1)) / stride
        self.output_seq_length = int(math.ceil(output_seq_length))
        self.encoder_seq_length = self.output_seq_length

210
211
        self.adapter_output_seq_length = (self.output_seq_length - 1) // adapter_stride + 1

Patrick von Platen's avatar
Patrick von Platen committed
212
    def prepare_config_and_inputs(self):
213
        input_values = floats_tensor([self.batch_size, self.seq_length], scale=1.0)
214
        attention_mask = random_attention_mask([self.batch_size, self.seq_length])
Patrick von Platen's avatar
Patrick von Platen committed
215

216
217
218
219
220
221
        config = self.get_config()

        return config, input_values, attention_mask

    def get_config(self):
        return Wav2Vec2Config(
Patrick von Platen's avatar
Patrick von Platen committed
222
223
224
225
226
227
228
229
            hidden_size=self.hidden_size,
            feat_extract_norm=self.feat_extract_norm,
            feat_extract_dropout=self.feat_extract_dropout,
            feat_extract_activation=self.feat_extract_activation,
            conv_dim=self.conv_dim,
            conv_stride=self.conv_stride,
            conv_kernel=self.conv_kernel,
            conv_bias=self.conv_bias,
230
231
            mask_time_prob=self.mask_time_prob,
            mask_time_length=self.mask_time_length,
Patrick von Platen's avatar
Patrick von Platen committed
232
233
234
235
236
237
238
            num_conv_pos_embeddings=self.num_conv_pos_embeddings,
            num_conv_pos_embedding_groups=self.num_conv_pos_embedding_groups,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            hidden_dropout_prob=self.hidden_dropout_prob,
            intermediate_size=self.intermediate_size,
            layer_norm_eps=self.layer_norm_eps,
239
            do_stable_layer_norm=self.do_stable_layer_norm,
Patrick von Platen's avatar
Patrick von Platen committed
240
241
242
            hidden_act=self.hidden_act,
            initializer_range=self.initializer_range,
            vocab_size=self.vocab_size,
243
244
            num_adapter_layers=self.num_adapter_layers,
            adapter_stride=self.adapter_stride,
245
246
247
248
            tdnn_dim=self.tdnn_dim,
            tdnn_kernel=self.tdnn_kernel,
            tdnn_dilation=self.tdnn_dilation,
            xvector_output_dim=self.xvector_output_dim,
Patrick von Platen's avatar
Patrick von Platen committed
249
250
        )

251
    def create_and_check_model(self, config, input_values, attention_mask):
Patrick von Platen's avatar
Patrick von Platen committed
252
253
254
        model = Wav2Vec2Model(config=config)
        model.to(torch_device)
        model.eval()
255
        result = model(input_values, attention_mask=attention_mask)
Patrick von Platen's avatar
Patrick von Platen committed
256
257
258
259
        self.parent.assertEqual(
            result.last_hidden_state.shape, (self.batch_size, self.output_seq_length, self.hidden_size)
        )

260
261
262
263
264
265
266
267
268
269
    def create_and_check_model_with_adapter(self, config, input_values, attention_mask):
        config.add_adapter = True
        model = Wav2Vec2Model(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_values, attention_mask=attention_mask)
        self.parent.assertEqual(
            result.last_hidden_state.shape, (self.batch_size, self.adapter_output_seq_length, self.hidden_size)
        )

270
271
272
273
274
275
276
277
278
279
280
    def create_and_check_model_with_adapter_for_ctc(self, config, input_values, attention_mask):
        config.add_adapter = True
        config.output_hidden_size = 2 * config.hidden_size
        model = Wav2Vec2ForCTC(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_values, attention_mask=attention_mask)
        self.parent.assertEqual(
            result.logits.shape, (self.batch_size, self.adapter_output_seq_length, self.vocab_size)
        )

281
282
283
284
285
286
287
288
289
290
291
292
    def create_and_check_model_with_adapter_proj_dim(self, config, input_values, attention_mask):
        config.add_adapter = True
        config.output_hidden_size = 8
        model = Wav2Vec2Model(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_values, attention_mask=attention_mask)
        self.parent.assertEqual(
            result.last_hidden_state.shape,
            (self.batch_size, self.adapter_output_seq_length, config.output_hidden_size),
        )

293
    def create_and_check_batch_inference(self, config, input_values, *args):
294
        # test does not pass for models making use of `group_norm`
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
        # check: https://github.com/pytorch/fairseq/issues/3227
        model = Wav2Vec2Model(config=config)
        model.to(torch_device)
        model.eval()

        input_values = input_values[:3]
        attention_mask = torch.ones(input_values.shape, device=torch_device, dtype=torch.bool)

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0
            attention_mask[i, input_lengths[i] :] = 0.0

        batch_outputs = model(input_values, attention_mask=attention_mask).last_hidden_state

        for i in range(input_values.shape[0]):
            input_slice = input_values[i : i + 1, : input_lengths[i]]
            output = model(input_slice).last_hidden_state

            batch_output = batch_outputs[i : i + 1, : output.shape[1]]
            self.parent.assertTrue(torch.allclose(output, batch_output, atol=1e-3))

319
320
321
322
323
324
325
326
    def check_ctc_loss(self, config, input_values, *args):
        model = Wav2Vec2ForCTC(config=config)
        model.to(torch_device)

        # make sure that dropout is disabled
        model.eval()

        input_values = input_values[:3]
327
        attention_mask = torch.ones(input_values.shape, device=torch_device, dtype=torch.long)
328
329
330
331
332
333
334
335

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], min(max_length_labels) - 1), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0
336
            attention_mask[i, input_lengths[i] :] = 0
337
338

        model.config.ctc_loss_reduction = "sum"
339
        sum_loss = model(input_values, attention_mask=attention_mask, labels=labels).loss.item()
340
341

        model.config.ctc_loss_reduction = "mean"
342
        mean_loss = model(input_values, attention_mask=attention_mask, labels=labels).loss.item()
343

344
345
        self.parent.assertTrue(isinstance(sum_loss, float))
        self.parent.assertTrue(isinstance(mean_loss, float))
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
    def check_seq_classifier_loss(self, config, input_values, *args):
        model = Wav2Vec2ForSequenceClassification(config=config)
        model.to(torch_device)

        # make sure that dropout is disabled
        model.eval()

        input_values = input_values[:3]
        attention_mask = torch.ones(input_values.shape, device=torch_device, dtype=torch.long)

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        labels = ids_tensor((input_values.shape[0], 1), len(model.config.id2label))

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0
            attention_mask[i, input_lengths[i] :] = 0

        masked_loss = model(input_values, attention_mask=attention_mask, labels=labels).loss.item()
        unmasked_loss = model(input_values, labels=labels).loss.item()

        self.parent.assertTrue(isinstance(masked_loss, float))
        self.parent.assertTrue(isinstance(unmasked_loss, float))
        self.parent.assertTrue(masked_loss != unmasked_loss)

    def check_ctc_training(self, config, input_values, *args):
373
374
375
376
377
378
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
379
        model.freeze_feature_encoder()
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                # it's important that we make sure that target lenghts are at least
                # one shorter than logit lenghts to prevent -inf
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()

401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
    def check_seq_classifier_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForSequenceClassification(config=config)
        model.to(torch_device)
        model.train()

        # freeze everything but the classification head
        model.freeze_base_model()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        labels = ids_tensor((input_values.shape[0], 1), len(model.config.id2label))

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()

424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
    def check_xvector_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForXVector(config=config)
        model.to(torch_device)
        model.train()

        # freeze everything but the classification head
        model.freeze_base_model()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        labels = ids_tensor((input_values.shape[0], 1), len(model.config.id2label))

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()

447
448
449
450
451
452
453
454
455
456
457
    def check_labels_out_of_vocab(self, config, input_values, *args):
        model = Wav2Vec2ForCTC(config)
        model.to(torch_device)
        model.train()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size + 100)

458
        with self.parent.assertRaises(ValueError):
459
460
            model(input_values, labels=labels)

Patrick von Platen's avatar
Patrick von Platen committed
461
    def prepare_config_and_inputs_for_common(self):
462
463
        config, input_values, attention_mask = self.prepare_config_and_inputs()
        inputs_dict = {"input_values": input_values, "attention_mask": attention_mask}
Patrick von Platen's avatar
Patrick von Platen committed
464
465
466
467
        return config, inputs_dict


@require_torch
468
class Wav2Vec2ModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
469
    all_model_classes = (
470
471
472
        (Wav2Vec2ForCTC, Wav2Vec2Model, Wav2Vec2ForMaskedLM, Wav2Vec2ForSequenceClassification, Wav2Vec2ForPreTraining)
        if is_torch_available()
        else ()
Patrick von Platen's avatar
Patrick von Platen committed
473
    )
474
475
476
477
478
479
480
481
482
483
    pipeline_model_mapping = (
        {
            "audio-classification": Wav2Vec2ForSequenceClassification,
            "automatic-speech-recognition": Wav2Vec2ForCTC,
            "feature-extraction": Wav2Vec2Model,
            "fill-mask": Wav2Vec2ForMaskedLM,
        }
        if is_torch_available()
        else {}
    )
484
    fx_compatible = True
Patrick von Platen's avatar
Patrick von Platen committed
485
486
487
488
489
490
491
492
493
494
495
496
497
498
    test_pruning = False
    test_headmasking = False

    def setUp(self):
        self.model_tester = Wav2Vec2ModelTester(self)
        self.config_tester = ConfigTester(self, config_class=Wav2Vec2Config, hidden_size=37)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

499
500
501
502
    def test_model_with_adapter(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_with_adapter(*config_and_inputs)

503
504
505
506
    def test_model_with_adapter_for_ctc(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_with_adapter_for_ctc(*config_and_inputs)

507
508
509
510
    def test_model_with_adapter_proj_dim(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_with_adapter_proj_dim(*config_and_inputs)

511
512
513
514
    def test_ctc_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_loss(*config_and_inputs)

515
    def test_seq_classifier_loss_inference(self):
516
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
517
518
519
520
521
522
523
524
525
        self.model_tester.check_seq_classifier_loss(*config_and_inputs)

    def test_ctc_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_training(*config_and_inputs)

    def test_seq_classifier_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_training(*config_and_inputs)
526

527
528
529
530
    def test_xvector_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_xvector_training(*config_and_inputs)

531
532
533
534
    def test_labels_out_of_vocab(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_labels_out_of_vocab(*config_and_inputs)

Patrick von Platen's avatar
Patrick von Platen committed
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
    # Wav2Vec2 has no inputs_embeds
    def test_inputs_embeds(self):
        pass

    # `input_ids` is renamed to `input_values`
    def test_forward_signature(self):
        pass

    # Wav2Vec2 cannot resize token embeddings
    # since it has no tokens embeddings
    def test_resize_tokens_embeddings(self):
        pass

    # Wav2Vec2 has no inputs_embeds
    # and thus the `get_input_embeddings` fn
    # is not implemented
    def test_model_common_attributes(self):
        pass

554
555
556
557
558
559
560
561
562
563
    @is_pt_flax_cross_test
    # non-robust architecture does not exist in Flax
    def test_equivalence_flax_to_pt(self):
        pass

    @is_pt_flax_cross_test
    # non-robust architecture does not exist in Flax
    def test_equivalence_pt_to_flax(self):
        pass

564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
    def test_retain_grad_hidden_states_attentions(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        # no need to test all models as different heads yield the same functionality
        model_class = self.all_model_classes[0]
        model = model_class(config)
        model.to(torch_device)

        # set layer drop to 0
        model.config.layerdrop = 0.0

        input_values = inputs_dict["input_values"]

        input_lengths = torch.tensor(
            [input_values.shape[1] for _ in range(input_values.shape[0])], dtype=torch.long, device=torch_device
        )
        output_lengths = model._get_feat_extract_output_lengths(input_lengths)

        labels = ids_tensor((input_values.shape[0], output_lengths[0] - 2), self.model_tester.vocab_size)
        inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["attention_mask"])
        inputs_dict["labels"] = labels

        outputs = model(**inputs_dict)

        output = outputs[0]

        # Encoder-/Decoder-only models
        hidden_states = outputs.hidden_states[0]
        attentions = outputs.attentions[0]

        hidden_states.retain_grad()
        attentions.retain_grad()

        output.flatten()[0].backward(retain_graph=True)

        self.assertIsNotNone(hidden_states.grad)
        self.assertIsNotNone(attentions.grad)

Patrick von Platen's avatar
Patrick von Platen committed
604
605
606
607
608
609
610
    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
Anton Lozhkov's avatar
Anton Lozhkov committed
611
612
613
614
615
                uniform_init_parms = [
                    "conv.weight",
                    "masked_spec_embed",
                    "codevectors",
                    "quantizer.weight_proj.weight",
616
617
618
619
620
621
                    "project_hid.weight",
                    "project_hid.bias",
                    "project_q.weight",
                    "project_q.bias",
                    "feature_projection.projection.weight",
                    "feature_projection.projection.bias",
622
                    "objective.weight",
Anton Lozhkov's avatar
Anton Lozhkov committed
623
                ]
Patrick von Platen's avatar
Patrick von Platen committed
624
                if param.requires_grad:
Anton Lozhkov's avatar
Anton Lozhkov committed
625
                    if any([x in name for x in uniform_init_parms]):
Patrick von Platen's avatar
Patrick von Platen committed
626
627
                        self.assertTrue(
                            -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
628
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
Patrick von Platen's avatar
Patrick von Platen committed
629
630
631
632
633
                        )
                    else:
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item(),
                            [0.0, 1.0],
634
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
Patrick von Platen's avatar
Patrick von Platen committed
635
636
                        )

637
638
639
640
    # overwrite from test_modeling_common
    def _mock_init_weights(self, module):
        if hasattr(module, "weight") and module.weight is not None:
            module.weight.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
641
        if hasattr(module, "weight_g") and module.weight_g is not None:
642
            module.weight_g.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
643
644
        if hasattr(module, "weight_v") and module.weight_v is not None:
            module.weight_v.data.fill_(3)
645
646
        if hasattr(module, "bias") and module.bias is not None:
            module.bias.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
647
648
        if hasattr(module, "codevectors") and module.codevectors is not None:
            module.codevectors.data.fill_(3)
649
650
        if hasattr(module, "masked_spec_embed") and module.masked_spec_embed is not None:
            module.masked_spec_embed.data.fill_(3)
651

652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
    def test_mask_feature_prob_ctc(self):
        model = Wav2Vec2ForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", mask_feature_prob=0.2, mask_feature_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

    def test_mask_time_prob_ctc(self):
        model = Wav2Vec2ForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", mask_time_prob=0.2, mask_time_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

698
699
700
701
    @unittest.skip(reason="Feed forward chunking is not implemented")
    def test_feed_forward_chunking(self):
        pass

Patrick von Platen's avatar
Patrick von Platen committed
702
703
704
705
706
    @slow
    def test_model_from_pretrained(self):
        model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")
        self.assertIsNotNone(model)

707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
    # Wav2Vec2 cannot be torchscripted because of group norm.
    def _create_and_check_torch_fx_tracing(self, config, inputs_dict, output_loss=False):
        if not is_torch_fx_available() or not self.fx_compatible:
            return

        configs_no_init = _config_zero_init(config)  # To be sure we have no Nan
        configs_no_init.return_dict = False

        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            model.to(torch_device)
            model.eval()
            inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=output_loss)

            try:
                input_names = [
                    "attention_mask",
                    "bbox",
                    "input_features",
                    "input_ids",
                    "input_values",
                    "pixel_values",
                    "token_type_ids",
                    "visual_feats",
                    "visual_pos",
                ]

                labels = inputs.get("labels", None)
                start_positions = inputs.get("start_positions", None)
                end_positions = inputs.get("end_positions", None)
                if labels is not None:
                    input_names.append("labels")
                if start_positions is not None:
                    input_names.append("start_positions")
                if end_positions is not None:
                    input_names.append("end_positions")

                filtered_inputs = {k: v for (k, v) in inputs.items() if k in input_names}
                input_names = list(filtered_inputs.keys())

                model_output = model(**filtered_inputs)

                if (
                    isinstance(model, Wav2Vec2ForSequenceClassification)
                    and not hasattr(model.config, "problem_type")
                    or model.config.problem_type is None
                ):
                    model.config.problem_type = "single_label_classification"

                traced_model = symbolic_trace(model, input_names)
                traced_output = traced_model(**filtered_inputs)

            except Exception as e:
                self.fail(f"Couldn't trace module: {e}")

            def flatten_output(output):
                flatten = []
                for x in output:
                    if isinstance(x, (tuple, list)):
                        flatten += flatten_output(x)
                    elif not isinstance(x, torch.Tensor):
                        continue
                    else:
                        flatten.append(x)
                return flatten

            model_output = flatten_output(model_output)
            traced_output = flatten_output(traced_output)
            num_outputs = len(model_output)

            for i in range(num_outputs):
                self.assertTrue(
                    torch.allclose(model_output[i], traced_output[i]),
                    f"traced {i}th output doesn't match model {i}th output for {model_class}",
                )

            # Test that the model can be serialized and restored properly
            with tempfile.TemporaryDirectory() as tmp_dir_name:
                pkl_file_name = os.path.join(tmp_dir_name, "model.pkl")
                try:
                    with open(pkl_file_name, "wb") as f:
                        pickle.dump(traced_model, f)
                    with open(pkl_file_name, "rb") as f:
                        loaded = pickle.load(f)
                except Exception as e:
                    self.fail(f"Couldn't serialize / deserialize the traced model: {e}")

                loaded_output = loaded(**filtered_inputs)
                loaded_output = flatten_output(loaded_output)

                for i in range(num_outputs):
                    self.assertTrue(
                        torch.allclose(model_output[i], loaded_output[i]),
                        f"serialized model {i}th output doesn't match model {i}th output for {model_class}",
                    )

            # Avoid memory leak. Without this, each call increase RAM usage by ~20MB.
            # (Even with this call, there are still memory leak by ~0.04MB)
            self.clear_torch_jit_class_registry()

Patrick von Platen's avatar
Patrick von Platen committed
807
808
809

@require_torch
class Wav2Vec2RobustModelTest(ModelTesterMixin, unittest.TestCase):
Anton Lozhkov's avatar
Anton Lozhkov committed
810
    all_model_classes = (
811
812
813
814
815
816
817
818
819
        (
            Wav2Vec2ForCTC,
            Wav2Vec2Model,
            Wav2Vec2ForMaskedLM,
            Wav2Vec2ForSequenceClassification,
            Wav2Vec2ForPreTraining,
            Wav2Vec2ForAudioFrameClassification,
            Wav2Vec2ForXVector,
        )
820
821
        if is_torch_available()
        else ()
Anton Lozhkov's avatar
Anton Lozhkov committed
822
    )
Patrick von Platen's avatar
Patrick von Platen committed
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
    test_pruning = False
    test_headmasking = False

    def setUp(self):
        self.model_tester = Wav2Vec2ModelTester(
            self, conv_stride=(3, 3, 3), feat_extract_norm="layer", do_stable_layer_norm=True
        )
        self.config_tester = ConfigTester(self, config_class=Wav2Vec2Config, hidden_size=37)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)
838
839
840
841
842
843
844
845

    def test_model_with_adapter(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_with_adapter(*config_and_inputs)

    def test_model_with_adapter_proj_dim(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_with_adapter_proj_dim(*config_and_inputs)
Patrick von Platen's avatar
Patrick von Platen committed
846

847
848
849
850
    def test_batched_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_batch_inference(*config_and_inputs)

851
852
853
854
    def test_ctc_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_loss(*config_and_inputs)

855
856
857
858
859
    def test_seq_classifier_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_loss(*config_and_inputs)

    def test_ctc_train(self):
860
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
861
862
863
864
865
        self.model_tester.check_ctc_training(*config_and_inputs)

    def test_seq_classifier_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_training(*config_and_inputs)
866

867
868
869
870
    def test_xvector_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_xvector_training(*config_and_inputs)

871
872
873
874
    def test_labels_out_of_vocab(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_labels_out_of_vocab(*config_and_inputs)

Patrick von Platen's avatar
Patrick von Platen committed
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
    # Wav2Vec2 has no inputs_embeds
    def test_inputs_embeds(self):
        pass

    # `input_ids` is renamed to `input_values`
    def test_forward_signature(self):
        pass

    # Wav2Vec2 cannot resize token embeddings
    # since it has no tokens embeddings
    def test_resize_tokens_embeddings(self):
        pass

    # Wav2Vec2 has no inputs_embeds
    # and thus the `get_input_embeddings` fn
    # is not implemented
    def test_model_common_attributes(self):
        pass

894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
    def test_retain_grad_hidden_states_attentions(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        # no need to test all models as different heads yield the same functionality
        model_class = self.all_model_classes[0]
        model = model_class(config)
        model.to(torch_device)

        # set layer drop to 0
        model.config.layerdrop = 0.0

        input_values = inputs_dict["input_values"]

        input_lengths = torch.tensor(
            [input_values.shape[1] for _ in range(input_values.shape[0])], dtype=torch.long, device=torch_device
        )
        output_lengths = model._get_feat_extract_output_lengths(input_lengths)

        labels = ids_tensor((input_values.shape[0], output_lengths[0] - 2), self.model_tester.vocab_size)
        inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["attention_mask"])
        inputs_dict["labels"] = labels

        outputs = model(**inputs_dict)

        output = outputs[0]

        # Encoder-/Decoder-only models
        hidden_states = outputs.hidden_states[0]
        attentions = outputs.attentions[0]

        hidden_states.retain_grad()
        attentions.retain_grad()

        output.flatten()[0].backward(retain_graph=True)

        self.assertIsNotNone(hidden_states.grad)
        self.assertIsNotNone(attentions.grad)

Patrick von Platen's avatar
Patrick von Platen committed
934
935
936
937
938
939
940
    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
Anton Lozhkov's avatar
Anton Lozhkov committed
941
942
943
944
945
                uniform_init_parms = [
                    "conv.weight",
                    "masked_spec_embed",
                    "codevectors",
                    "quantizer.weight_proj.weight",
946
947
948
949
950
951
                    "project_hid.weight",
                    "project_hid.bias",
                    "project_q.weight",
                    "project_q.bias",
                    "feature_projection.projection.weight",
                    "feature_projection.projection.bias",
952
                    "objective.weight",
Anton Lozhkov's avatar
Anton Lozhkov committed
953
                ]
Patrick von Platen's avatar
Patrick von Platen committed
954
                if param.requires_grad:
Anton Lozhkov's avatar
Anton Lozhkov committed
955
                    if any([x in name for x in uniform_init_parms]):
Patrick von Platen's avatar
Patrick von Platen committed
956
957
                        self.assertTrue(
                            -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
958
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
Patrick von Platen's avatar
Patrick von Platen committed
959
960
961
962
963
                        )
                    else:
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item(),
                            [0.0, 1.0],
964
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
Patrick von Platen's avatar
Patrick von Platen committed
965
966
                        )

967
968
969
970
    # overwrite from test_modeling_common
    def _mock_init_weights(self, module):
        if hasattr(module, "weight") and module.weight is not None:
            module.weight.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
971
        if hasattr(module, "weight_g") and module.weight_g is not None:
972
            module.weight_g.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
973
974
        if hasattr(module, "weight_v") and module.weight_v is not None:
            module.weight_v.data.fill_(3)
975
976
        if hasattr(module, "bias") and module.bias is not None:
            module.bias.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
977
978
        if hasattr(module, "codevectors") and module.codevectors is not None:
            module.codevectors.data.fill_(3)
979
980
        if hasattr(module, "masked_spec_embed") and module.masked_spec_embed is not None:
            module.masked_spec_embed.data.fill_(3)
Anton Lozhkov's avatar
Anton Lozhkov committed
981
982
983
984
985

    def test_model_for_pretraining(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        model = Wav2Vec2ForPreTraining(config).to(torch_device)

986
987
988
989
        batch_size = inputs_dict["input_values"].shape[0]
        feature_seq_length = int(model._get_feat_extract_output_lengths(inputs_dict["input_values"].shape[1]))

        features_shape = (batch_size, feature_seq_length)
Anton Lozhkov's avatar
Anton Lozhkov committed
990
991
992
993
994
995

        mask_time_indices = _compute_mask_indices(
            features_shape,
            model.config.mask_time_prob,
            model.config.mask_time_length,
            min_masks=2,
996
997
998
999
1000
        )
        sampled_negative_indices = _sample_negative_indices(features_shape, 10, mask_time_indices)

        mask_time_indices = torch.from_numpy(mask_time_indices).to(torch_device)
        sampled_negative_indices = torch.from_numpy(sampled_negative_indices).to(torch_device)
Anton Lozhkov's avatar
Anton Lozhkov committed
1001
1002
1003
1004
1005

        loss = model(
            inputs_dict["input_values"],
            attention_mask=inputs_dict["attention_mask"],
            mask_time_indices=mask_time_indices,
1006
            sampled_negative_indices=sampled_negative_indices,
Anton Lozhkov's avatar
Anton Lozhkov committed
1007
1008
        ).loss

1009
        # more losses
Anton Lozhkov's avatar
Anton Lozhkov committed
1010
        mask_time_indices[:, : mask_time_indices.shape[-1] // 2] = True
1011
1012
1013

        sampled_negative_indices = _sample_negative_indices(features_shape, 10, mask_time_indices.cpu().numpy())
        sampled_negative_indices = torch.from_numpy(sampled_negative_indices).to(torch_device)
Anton Lozhkov's avatar
Anton Lozhkov committed
1014
1015
1016
1017
        loss_more_masked = model(
            inputs_dict["input_values"],
            attention_mask=inputs_dict["attention_mask"],
            mask_time_indices=mask_time_indices,
1018
            sampled_negative_indices=sampled_negative_indices,
Anton Lozhkov's avatar
Anton Lozhkov committed
1019
1020
1021
1022
        ).loss

        # loss_more_masked has to be bigger or equal loss since more masked inputs have to be predicted
        self.assertTrue(loss.detach().item() <= loss_more_masked.detach().item())
1023

1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
    def test_mask_feature_prob_ctc(self):
        model = Wav2Vec2ForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", mask_feature_prob=0.2, mask_feature_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

    def test_mask_time_prob_ctc(self):
        model = Wav2Vec2ForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", mask_time_prob=0.2, mask_time_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
    def test_mask_time_feature_prob_ctc_single_batch(self):
        model = Wav2Vec2ForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2",
            mask_time_prob=0.2,
            mask_feature_prob=0.2,
            mask_time_length=2,
            mask_feature_length=2,
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-wav2vec2", return_attention_mask=True
        )

        batch_duration_in_seconds = [6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (1, 1498, 32))

1097
1098
1099
1100
    @unittest.skip(reason="Feed forward chunking is not implemented")
    def test_feed_forward_chunking(self):
        pass

Patrick von Platen's avatar
Patrick von Platen committed
1101
1102
1103
1104
1105
1106
    @slow
    def test_model_from_pretrained(self):
        model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h")
        self.assertIsNotNone(model)


1107
1108
1109
1110
1111
1112
1113
1114
@require_torch
class Wav2Vec2UtilsTest(unittest.TestCase):
    def test_compute_mask_indices(self):
        batch_size = 4
        sequence_length = 60
        mask_prob = 0.5
        mask_length = 1

1115
1116
        mask = _compute_mask_indices((batch_size, sequence_length), mask_prob, mask_length)
        mask = torch.from_numpy(mask).to(torch_device)
1117
1118
1119

        self.assertListEqual(mask.sum(axis=-1).tolist(), [mask_prob * sequence_length for _ in range(batch_size)])

1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
    def test_compute_mask_indices_low_prob(self):
        # with these settings num_masked_spans=0.5, which means probabilistic rounding
        # ensures that in 5 out of 10 method calls, num_masked_spans=0, and in
        # the other 5 out of 10, cases num_masked_spans=1
        n_trials = 100
        batch_size = 4
        sequence_length = 100
        mask_prob = 0.05
        mask_length = 10

        count_dimensions_masked = 0
        count_dimensions_not_masked = 0

        for _ in range(n_trials):
            mask = _compute_mask_indices((batch_size, sequence_length), mask_prob, mask_length)
            mask = torch.from_numpy(mask).to(torch_device)

            num_masks = torch.sum(mask).item()

            if num_masks > 0:
                count_dimensions_masked += 1
            else:
                count_dimensions_not_masked += 1

        # as we test for at least 10 masked dimension and at least
        # 10 non-masked dimension, this test could fail with probability:
        # P(100 coin flips, at most 9 heads) = 1.66e-18
        self.assertGreater(count_dimensions_masked, int(n_trials * 0.1))
        self.assertGreater(count_dimensions_not_masked, int(n_trials * 0.1))

1150
1151
    def test_compute_mask_indices_overlap(self):
        batch_size = 4
Anton Lozhkov's avatar
Anton Lozhkov committed
1152
        sequence_length = 80
1153
1154
1155
        mask_prob = 0.5
        mask_length = 4

1156
1157
        mask = _compute_mask_indices((batch_size, sequence_length), mask_prob, mask_length)
        mask = torch.from_numpy(mask).to(torch_device)
1158

Anton Lozhkov's avatar
Anton Lozhkov committed
1159
        # because of overlap mask don't have to add up exactly to `mask_prob * sequence_length`, but have to be smaller or equal
1160
        for batch_sum in mask.sum(axis=-1):
Anton Lozhkov's avatar
Anton Lozhkov committed
1161
1162
            self.assertTrue(int(batch_sum) <= mask_prob * sequence_length)

1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
    def test_compute_mask_indices_attn_mask_overlap(self):
        batch_size = 4
        sequence_length = 80
        mask_prob = 0.5
        mask_length = 4

        attention_mask = torch.ones((batch_size, sequence_length), dtype=torch.long, device=torch_device)
        attention_mask[:2, sequence_length // 2 :] = 0

        mask = _compute_mask_indices(
1173
            (batch_size, sequence_length), mask_prob, mask_length, attention_mask=attention_mask
1174
        )
1175
        mask = torch.from_numpy(mask).to(torch_device)
1176
1177
1178
1179
1180
1181

        for batch_sum in mask.sum(axis=-1):
            self.assertTrue(int(batch_sum) <= mask_prob * sequence_length)

        self.assertTrue(mask[:2, sequence_length // 2 :].sum() == 0)

1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
    def test_compute_mask_indices_short_audio(self):
        batch_size = 4
        sequence_length = 100
        mask_prob = 0.05
        mask_length = 10

        attention_mask = torch.ones((batch_size, sequence_length), dtype=torch.long, device=torch_device)
        # force one example to be heavily padded
        attention_mask[0, 5:] = 0

        mask = _compute_mask_indices(
            (batch_size, sequence_length), mask_prob, mask_length, attention_mask=attention_mask, min_masks=2
        )

        # make sure that non-padded examples cannot be padded
        self.assertFalse(mask[0][attention_mask[0].to(torch.bool).cpu()].any())

Anton Lozhkov's avatar
Anton Lozhkov committed
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
    def test_compute_perplexity(self):
        probs = torch.arange(100, device=torch_device).reshape(2, 5, 10) / 100

        ppl = Wav2Vec2GumbelVectorQuantizer._compute_perplexity(probs)
        self.assertTrue(abs(ppl.item() - 141.4291) < 1e-3)

        # mask half of the input
        mask = torch.ones((2,), device=torch_device, dtype=torch.bool)
        mask[0] = 0

        ppl = Wav2Vec2GumbelVectorQuantizer._compute_perplexity(probs, mask)
        self.assertTrue(abs(ppl.item() - 58.6757) < 1e-3)

    def test_sample_negatives(self):
        batch_size = 2
        sequence_length = 10
        hidden_size = 4
        num_negatives = 3
1217
1218
1219
        sequence = torch.div(
            torch.arange(sequence_length * hidden_size, device=torch_device), hidden_size, rounding_mode="floor"
        )
1220
        features = sequence.view(sequence_length, hidden_size)  # each value in vector consits of same value
Anton Lozhkov's avatar
Anton Lozhkov committed
1221
1222
        features = features[None, :].expand(batch_size, sequence_length, hidden_size).contiguous()

1223
1224
1225
1226
1227
        # sample negative indices
        sampled_negative_indices = _sample_negative_indices((batch_size, sequence_length), num_negatives, None)
        sampled_negative_indices = torch.from_numpy(sampled_negative_indices).to(torch_device)
        negatives = features.view(-1, hidden_size)[sampled_negative_indices.long().view(-1)]
        negatives = negatives.view(batch_size, sequence_length, -1, hidden_size).permute(2, 0, 1, 3)
Anton Lozhkov's avatar
Anton Lozhkov committed
1228
1229
1230
1231
1232
1233
1234
        self.assertTrue(negatives.shape == (num_negatives, batch_size, sequence_length, hidden_size))

        # make sure no negatively sampled vector is actually a positive one
        for negative in negatives:
            self.assertTrue(((negative - features) == 0).sum() == 0.0)

        # make sure that full vectors are sampled and not values of vectors => this means that `unique()` yields a single value for `hidden_size` dim
1235
        self.assertEqual(negatives.unique(dim=-1).shape, (num_negatives, batch_size, sequence_length, 1))
1236

1237
    def test_sample_negatives_with_mask(self):
1238
1239
1240
1241
1242
1243
        batch_size = 2
        sequence_length = 10
        hidden_size = 4
        num_negatives = 3

        # second half of last input tensor is padded
1244
1245
        mask = torch.ones((batch_size, sequence_length), dtype=torch.long, device=torch_device)
        mask[-1, sequence_length // 2 :] = 0
1246

1247
1248
1249
        sequence = torch.div(
            torch.arange(sequence_length * hidden_size, device=torch_device), hidden_size, rounding_mode="floor"
        )
1250
        features = sequence.view(sequence_length, hidden_size)  # each value in vector consits of same value
1251
1252
1253
        features = features[None, :].expand(batch_size, sequence_length, hidden_size).contiguous()

        # replace masked feature vectors with -100 to test that those are not sampled
1254
        features = torch.where(mask[:, :, None].expand(features.shape).bool(), features, -100)
1255

1256
1257
1258
1259
1260
1261
1262
        # sample negative indices
        sampled_negative_indices = _sample_negative_indices(
            (batch_size, sequence_length), num_negatives, mask.cpu().numpy()
        )
        sampled_negative_indices = torch.from_numpy(sampled_negative_indices).to(torch_device)
        negatives = features.view(-1, hidden_size)[sampled_negative_indices.long().view(-1)]
        negatives = negatives.view(batch_size, sequence_length, -1, hidden_size).permute(2, 0, 1, 3)
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272

        self.assertTrue((negatives >= 0).all().item())

        self.assertTrue(negatives.shape == (num_negatives, batch_size, sequence_length, hidden_size))

        # make sure no negatively sampled vector is actually a positive one
        for negative in negatives:
            self.assertTrue(((negative - features) == 0).sum() == 0.0)

        # make sure that full vectors are sampled and not values of vectors => this means that `unique()` yields a single value for `hidden_size` dim
1273
        self.assertEqual(negatives.unique(dim=-1).shape, (num_negatives, batch_size, sequence_length, 1))
1274

1275

Patrick von Platen's avatar
Patrick von Platen committed
1276
1277
@require_torch
@require_soundfile
Anton Lozhkov's avatar
Anton Lozhkov committed
1278
@slow
Patrick von Platen's avatar
Patrick von Platen committed
1279
1280
class Wav2Vec2ModelIntegrationTest(unittest.TestCase):
    def _load_datasamples(self, num_samples):
Patrick von Platen's avatar
Patrick von Platen committed
1281
        ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
1282
1283
1284
1285
        # automatic decoding with librispeech
        speech_samples = ds.sort("id").filter(
            lambda x: x["id"] in [f"1272-141231-000{i}" for i in range(num_samples)]
        )[:num_samples]["audio"]
1286

1287
        return [x["array"] for x in speech_samples]
Patrick von Platen's avatar
Patrick von Platen committed
1288

1289
1290
1291
1292
1293
    def _load_superb(self, task, num_samples):
        ds = load_dataset("anton-l/superb_dummy", task, split="test")

        return ds[:num_samples]

1294
    def test_inference_ctc_normal(self):
1295
        model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
Patrick von Platen's avatar
Patrick von Platen committed
1296
        model.to(torch_device)
1297
        processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h", do_lower_case=True)
Patrick von Platen's avatar
Patrick von Platen committed
1298
1299
        input_speech = self._load_datasamples(1)

1300
        input_values = processor(input_speech, return_tensors="pt").input_values.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
1301
1302
1303
1304
1305

        with torch.no_grad():
            logits = model(input_values).logits

        predicted_ids = torch.argmax(logits, dim=-1)
1306
        predicted_trans = processor.batch_decode(predicted_ids)
Patrick von Platen's avatar
Patrick von Platen committed
1307
1308
1309
1310

        EXPECTED_TRANSCRIPTIONS = ["a man said to the universe sir i exist"]
        self.assertListEqual(predicted_trans, EXPECTED_TRANSCRIPTIONS)

1311
    def test_inference_ctc_normal_batched(self):
1312
        model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
Patrick von Platen's avatar
Patrick von Platen committed
1313
        model.to(torch_device)
1314
        processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h", do_lower_case=True)
Patrick von Platen's avatar
Patrick von Platen committed
1315
1316
1317

        input_speech = self._load_datasamples(2)

1318
        inputs = processor(input_speech, return_tensors="pt", padding=True)
1319
1320

        input_values = inputs.input_values.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
1321
1322
1323
1324
1325

        with torch.no_grad():
            logits = model(input_values).logits

        predicted_ids = torch.argmax(logits, dim=-1)
1326
        predicted_trans = processor.batch_decode(predicted_ids)
Patrick von Platen's avatar
Patrick von Platen committed
1327
1328
1329
1330
1331
1332
1333

        EXPECTED_TRANSCRIPTIONS = [
            "a man said to the universe sir i exist",
            "sweat covered brion's body trickling into the tight lowing cloth that was the only garment he wore",
        ]
        self.assertListEqual(predicted_trans, EXPECTED_TRANSCRIPTIONS)

1334
    def test_inference_ctc_robust_batched(self):
1335
        model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h-lv60-self").to(torch_device)
1336
        processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h-lv60-self", do_lower_case=True)
Patrick von Platen's avatar
Patrick von Platen committed
1337
1338
1339

        input_speech = self._load_datasamples(4)

1340
        inputs = processor(input_speech, return_tensors="pt", padding=True)
1341
1342
1343

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
1344
1345

        with torch.no_grad():
1346
            logits = model(input_values, attention_mask=attention_mask).logits
Patrick von Platen's avatar
Patrick von Platen committed
1347
1348

        predicted_ids = torch.argmax(logits, dim=-1)
1349
        predicted_trans = processor.batch_decode(predicted_ids)
Patrick von Platen's avatar
Patrick von Platen committed
1350
1351
1352
1353

        EXPECTED_TRANSCRIPTIONS = [
            "a man said to the universe sir i exist",
            "sweat covered brion's body trickling into the tight loin cloth that was the only garment he wore",
Sylvain Gugger's avatar
Sylvain Gugger committed
1354
1355
            "the cut on his chest still dripping blood the ache of his overstrained eyes even the soaring arena around"
            " him with the thousands of spectators were trivialities not worth thinking about",
Patrick von Platen's avatar
Patrick von Platen committed
1356
1357
1358
            "his instant panic was followed by a small sharp blow high on his chest",
        ]
        self.assertListEqual(predicted_trans, EXPECTED_TRANSCRIPTIONS)
Anton Lozhkov's avatar
Anton Lozhkov committed
1359

1360
    @unittest.skipIf(torch_device != "cpu", "cannot make deterministic on GPU")
Anton Lozhkov's avatar
Anton Lozhkov committed
1361
    def test_inference_integration(self):
1362
        model = Wav2Vec2ForPreTraining.from_pretrained("facebook/wav2vec2-base")
Anton Lozhkov's avatar
Anton Lozhkov committed
1363
        model.to(torch_device)
1364
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base")
Anton Lozhkov's avatar
Anton Lozhkov committed
1365
1366
1367
1368
        input_speech = self._load_datasamples(2)

        inputs_dict = feature_extractor(input_speech, return_tensors="pt", padding=True)

1369
1370
1371
1372
        batch_size = inputs_dict["input_values"].shape[0]
        feature_seq_length = int(model._get_feat_extract_output_lengths(inputs_dict["input_values"].shape[1]))

        features_shape = (batch_size, feature_seq_length)
Anton Lozhkov's avatar
Anton Lozhkov committed
1373

1374
        np.random.seed(4)
Anton Lozhkov's avatar
Anton Lozhkov committed
1375
1376
1377
1378
1379
        mask_time_indices = _compute_mask_indices(
            features_shape,
            model.config.mask_time_prob,
            model.config.mask_time_length,
            min_masks=2,
1380
1381
        )
        mask_time_indices = torch.from_numpy(mask_time_indices).to(torch_device)
Anton Lozhkov's avatar
Anton Lozhkov committed
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394

        with torch.no_grad():
            outputs = model(
                inputs_dict.input_values.to(torch_device),
                mask_time_indices=mask_time_indices,
            )

        # compute cosine similarity
        cosine_sim = torch.cosine_similarity(outputs.projected_states, outputs.projected_quantized_states, dim=-1)

        # retrieve cosine sim of masked features
        cosine_sim_masked = cosine_sim[mask_time_indices]

1395
1396
        # cosine similarity of model is all > 0.5 as model is
        # pre-trained on contrastive loss
Anton Lozhkov's avatar
Anton Lozhkov committed
1397
        # fmt: off
1398
1399
1400
1401
1402
1403
1404
        expected_cosine_sim_masked = torch.tensor([
            0.8523, 0.5860, 0.6905, 0.5557, 0.7456, 0.5249, 0.6639, 0.7654, 0.7565,
            0.8167, 0.8222, 0.7960, 0.8034, 0.8166, 0.8310, 0.8263, 0.8274, 0.8258,
            0.8179, 0.8412, 0.8536, 0.5098, 0.4728, 0.6461, 0.4498, 0.6002, 0.5774,
            0.6457, 0.7123, 0.5668, 0.6866, 0.4960, 0.6293, 0.7423, 0.7419, 0.7526,
            0.7768, 0.4898, 0.5393, 0.8183
        ], device=torch_device)
Anton Lozhkov's avatar
Anton Lozhkov committed
1405
1406
1407
1408
1409
        # fmt: on

        self.assertTrue(torch.allclose(cosine_sim_masked, expected_cosine_sim_masked, atol=1e-3))

    def test_inference_pretrained(self):
1410
        model = Wav2Vec2ForPreTraining.from_pretrained("facebook/wav2vec2-base")
Anton Lozhkov's avatar
Anton Lozhkov committed
1411
1412
        model.to(torch_device)
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
1413
            "facebook/wav2vec2-base", return_attention_mask=True
Anton Lozhkov's avatar
Anton Lozhkov committed
1414
1415
1416
1417
1418
        )
        input_speech = self._load_datasamples(2)

        inputs_dict = feature_extractor(input_speech, return_tensors="pt", padding=True)

1419
1420
1421
1422
        batch_size = inputs_dict["input_values"].shape[0]
        feature_seq_length = int(model._get_feat_extract_output_lengths(inputs_dict["input_values"].shape[1]))

        features_shape = (batch_size, feature_seq_length)
Anton Lozhkov's avatar
Anton Lozhkov committed
1423
1424
1425
1426
1427
1428
1429

        torch.manual_seed(0)
        mask_time_indices = _compute_mask_indices(
            features_shape,
            model.config.mask_time_prob,
            model.config.mask_time_length,
            min_masks=2,
1430
1431
        )
        mask_time_indices = torch.from_numpy(mask_time_indices).to(torch_device)
Anton Lozhkov's avatar
Anton Lozhkov committed
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447

        with torch.no_grad():
            outputs = model(
                inputs_dict.input_values.to(torch_device),
                attention_mask=inputs_dict.attention_mask.to(torch_device),
                mask_time_indices=mask_time_indices,
            )

        # compute cosine similarity
        cosine_sim = torch.cosine_similarity(outputs.projected_states, outputs.projected_quantized_states, dim=-1)

        # retrieve cosine sim of masked features
        cosine_sim_masked = cosine_sim[mask_time_indices]

        # ... now compare to randomly initialized model

1448
        config = Wav2Vec2Config.from_pretrained("facebook/wav2vec2-base")
Anton Lozhkov's avatar
Anton Lozhkov committed
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
        model_rand = Wav2Vec2ForPreTraining(config).to(torch_device).eval()

        with torch.no_grad():
            outputs_rand = model_rand(
                inputs_dict.input_values.to(torch_device),
                attention_mask=inputs_dict.attention_mask.to(torch_device),
                mask_time_indices=mask_time_indices,
            )

        # compute cosine similarity
        cosine_sim_rand = torch.cosine_similarity(
            outputs_rand.projected_states, outputs_rand.projected_quantized_states, dim=-1
        )

        # retrieve cosine sim of masked features
        cosine_sim_masked_rand = cosine_sim_rand[mask_time_indices]

        # a pretrained wav2vec2 model has learned to predict the quantized latent states
        # => the cosine similarity between quantized states and predicted states > 0.5
        # a random wav2vec2 model has not learned to predict the quantized latent states
        # => the cosine similarity between quantized states and predicted states is very likely < 0.1
        self.assertTrue(cosine_sim_masked.mean().item() - 5 * cosine_sim_masked_rand.mean().item() > 0)

1472
    @unittest.skipIf(torch_device != "cpu", "cannot make deterministic on GPU")
Anton Lozhkov's avatar
Anton Lozhkov committed
1473
1474
    def test_loss_pretraining(self):
        model = Wav2Vec2ForPreTraining.from_pretrained(
1475
            "facebook/wav2vec2-base",
Anton Lozhkov's avatar
Anton Lozhkov committed
1476
1477
1478
1479
1480
1481
1482
1483
            attention_dropout=0.0,
            feat_proj_dropout=0.0,
            hidden_dropout=0.0,
            layerdrop=0.0,
        )
        model.to(torch_device).train()

        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
1484
            "facebook/wav2vec2-base", return_attention_mask=True
Anton Lozhkov's avatar
Anton Lozhkov committed
1485
1486
1487
1488
1489
        )
        input_speech = self._load_datasamples(2)

        inputs_dict = feature_extractor(input_speech, return_tensors="pt", padding=True)

1490
1491
1492
1493
        batch_size = inputs_dict["input_values"].shape[0]
        feature_seq_length = int(model._get_feat_extract_output_lengths(inputs_dict["input_values"].shape[1]))

        features_shape = (batch_size, feature_seq_length)
Anton Lozhkov's avatar
Anton Lozhkov committed
1494
1495

        torch.manual_seed(0)
1496
1497
        np.random.seed(0)

Anton Lozhkov's avatar
Anton Lozhkov committed
1498
1499
1500
1501
1502
        mask_time_indices = _compute_mask_indices(
            features_shape,
            model.config.mask_time_prob,
            model.config.mask_time_length,
            min_masks=2,
1503
1504
1505
1506
1507
1508
1509
        )
        sampled_negative_indices = _sample_negative_indices(
            mask_time_indices.shape, model.config.num_negatives, mask_time_indices
        )

        mask_time_indices = torch.from_numpy(mask_time_indices).to(torch_device)
        sampled_negative_indices = torch.from_numpy(sampled_negative_indices).to(torch_device)
Anton Lozhkov's avatar
Anton Lozhkov committed
1510
1511
1512
1513
1514
1515

        with torch.no_grad():
            outputs = model(
                inputs_dict.input_values.to(torch_device),
                attention_mask=inputs_dict.attention_mask.to(torch_device),
                mask_time_indices=mask_time_indices,
1516
                sampled_negative_indices=sampled_negative_indices,
Anton Lozhkov's avatar
Anton Lozhkov committed
1517
1518
1519
1520
1521
            )

        # check diversity loss
        num_codevectors = model.config.num_codevectors_per_group * model.config.num_codevector_groups
        diversity_loss = (num_codevectors - outputs.codevector_perplexity) / num_codevectors
1522
        self.assertTrue(abs(diversity_loss.item() - 0.9538) < 1e-3)
Anton Lozhkov's avatar
Anton Lozhkov committed
1523
1524

        # check overall loss (contrastive loss + diversity loss)
1525
        expected_loss = 116.7094
Anton Lozhkov's avatar
Anton Lozhkov committed
1526
1527

        self.assertTrue(abs(outputs.loss.item() - expected_loss) < 1e-3)
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616

    def test_inference_keyword_spotting(self):
        model = Wav2Vec2ForSequenceClassification.from_pretrained("superb/wav2vec2-base-superb-ks").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/wav2vec2-base-superb-ks")
        input_data = self._load_superb("ks", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)
        predicted_logits, predicted_ids = torch.max(outputs.logits, dim=-1)

        expected_labels = [7, 6, 10, 9]
        # s3prl logits for the same batch
        expected_logits = torch.tensor([6.1186, 11.8961, 10.2931, 6.0898], device=torch_device)

        self.assertListEqual(predicted_ids.tolist(), expected_labels)
        self.assertTrue(torch.allclose(predicted_logits, expected_logits, atol=1e-2))

    def test_inference_intent_classification(self):
        model = Wav2Vec2ForSequenceClassification.from_pretrained("superb/wav2vec2-base-superb-ic").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/wav2vec2-base-superb-ic")
        input_data = self._load_superb("ic", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)

        predicted_logits_action, predicted_ids_action = torch.max(outputs.logits[:, :6], dim=-1)
        predicted_logits_object, predicted_ids_object = torch.max(outputs.logits[:, 6:20], dim=-1)
        predicted_logits_location, predicted_ids_location = torch.max(outputs.logits[:, 20:24], dim=-1)

        expected_labels_action = [0, 0, 2, 3]
        expected_logits_action = torch.tensor([0.4568, 11.0848, 1.6621, 9.3841], device=torch_device)
        expected_labels_object = [3, 10, 3, 4]
        expected_logits_object = torch.tensor([1.5322, 10.7094, 5.2469, 22.1318], device=torch_device)
        expected_labels_location = [0, 0, 0, 1]
        expected_logits_location = torch.tensor([1.5335, 6.5096, 10.5704, 11.0569], device=torch_device)

        self.assertListEqual(predicted_ids_action.tolist(), expected_labels_action)
        self.assertListEqual(predicted_ids_object.tolist(), expected_labels_object)
        self.assertListEqual(predicted_ids_location.tolist(), expected_labels_location)

        self.assertTrue(torch.allclose(predicted_logits_action, expected_logits_action, atol=1e-2))
        self.assertTrue(torch.allclose(predicted_logits_object, expected_logits_object, atol=1e-2))
        self.assertTrue(torch.allclose(predicted_logits_location, expected_logits_location, atol=1e-2))

    def test_inference_speaker_identification(self):
        model = Wav2Vec2ForSequenceClassification.from_pretrained("superb/wav2vec2-base-superb-sid").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/wav2vec2-base-superb-sid")
        input_data = self._load_superb("si", 4)

        output_logits = []
        with torch.no_grad():
            for example in input_data["speech"]:
                input = processor(example, return_tensors="pt", padding=True)
                output = model(input.input_values.to(torch_device), attention_mask=None)
                output_logits.append(output.logits[0])
        output_logits = torch.stack(output_logits)
        predicted_logits, predicted_ids = torch.max(output_logits, dim=-1)

        expected_labels = [251, 1, 1, 3]
        # s3prl logits for the same batch
        expected_logits = torch.tensor([37.5627, 71.6362, 64.2419, 31.7778], device=torch_device)

        self.assertListEqual(predicted_ids.tolist(), expected_labels)
        self.assertTrue(torch.allclose(predicted_logits, expected_logits, atol=1e-2))

    def test_inference_emotion_recognition(self):
        model = Wav2Vec2ForSequenceClassification.from_pretrained("superb/wav2vec2-base-superb-er").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/wav2vec2-base-superb-er")
        input_data = self._load_superb("er", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)
        predicted_logits, predicted_ids = torch.max(outputs.logits, dim=-1)

        expected_labels = [1, 1, 2, 2]
        # s3prl logits for the same batch
        expected_logits = torch.tensor([2.1722, 3.0779, 8.0287, 6.6797], device=torch_device)

        self.assertListEqual(predicted_ids.tolist(), expected_labels)
        self.assertTrue(torch.allclose(predicted_logits, expected_logits, atol=1e-2))
1617

1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
    def test_phoneme_recognition(self):
        model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-lv-60-espeak-cv-ft").to(torch_device)
        processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-lv-60-espeak-cv-ft")

        input_speech = self._load_datasamples(4)

        inputs = processor(input_speech, return_tensors="pt", padding=True)

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        with torch.no_grad():
            logits = model(input_values, attention_mask=attention_mask).logits

        predicted_ids = torch.argmax(logits, dim=-1)
        predicted_trans = processor.batch_decode(predicted_ids)

        EXPECTED_TRANSCRIPTIONS = [
            "ɐ m æ n s ɛ d t ə ð ə j uː n ɪ v ɚ s s ɚ aɪ ɛ ɡ z ɪ s t",
Sylvain Gugger's avatar
Sylvain Gugger committed
1637
1638
1639
1640
1641
            "s w ɛ t k ʌ v ɚ d b ɹ iː ɔ n z b ɑː d i t ɹ ɪ k l ɪ ŋ ɪ n t ə ð ə t aɪ t l oɪ n k l ɑː θ ð æ w ʌ z ð ɪ oʊ"
            " n l i ɡ ɑːɹ m ə n t h iː w ɔːɹ",
            "ð ə k aɪ t ɔ n h ɪ z tʃ ɛ s t s t ɪ l d ɹ ɪ p ɪ ŋ b l ʌ d ð ɪ eɪ k ʌ v h ɪ z oʊ v ɚ s t ɹ eɪ n d aɪ z iː"
            " v ə n ð ə s ɔːɹ ɹ ɪ ŋ ɐ ɹ iː n ɐ ɚ ɹ aʊ n d h ɪ m w ɪ ð ə θ aʊ z ə n d z ʌ v s p ɛ k t eɪ ɾ ɚ z w ɜː t ɹ"
            " ɪ v ɪ æ l ᵻ ɾ i z n ɑː t w ɜː θ θ ɪ ŋ k ɪ ŋ ɐ b aʊ t",
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
            "h ɪ z ɪ n s t ə n t v p æ n ɪ k w ʌ z f ɑː l oʊ d b aɪ ɐ s m ɔː l ʃ ɑːɹ p b l oʊ h aɪ ɔ n h ɪ z tʃ ɛ s t",
        ]
        # should correspond to =>:
        # [
        # "a man said to the universe sir i exist",
        # "sweat covered brion's body trickling into the tight loin cloth that was the only garment he wore",
        # "the cut on his chest still dripping blood the ache of his overstrained eyes even the soaring arena around him with the thousands of spectators were trivialities not worth thinking about",
        # "his instant panic was followed by a small sharp blow high on his chest",
        # ]
        self.assertListEqual(predicted_trans, EXPECTED_TRANSCRIPTIONS)

1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
    @require_pyctcdecode
    @require_torchaudio
    def test_wav2vec2_with_lm(self):
        ds = load_dataset("common_voice", "es", split="test", streaming=True)
        sample = next(iter(ds))

        resampled_audio = torchaudio.functional.resample(
            torch.tensor(sample["audio"]["array"]), 48_000, 16_000
        ).numpy()

        model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm").to(
            torch_device
        )
        processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")

        input_values = processor(resampled_audio, return_tensors="pt").input_values

        with torch.no_grad():
            logits = model(input_values.to(torch_device)).logits

        transcription = processor.batch_decode(logits.cpu().numpy()).text

        self.assertEqual(transcription[0], "bien y qué regalo vas a abrir primero")
1676

1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
    @require_pyctcdecode
    @require_torchaudio
    def test_wav2vec2_with_lm_pool(self):
        ds = load_dataset("common_voice", "es", split="test", streaming=True)
        sample = next(iter(ds))

        resampled_audio = torchaudio.functional.resample(
            torch.tensor(sample["audio"]["array"]), 48_000, 16_000
        ).numpy()

        model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm").to(
            torch_device
        )
        processor = Wav2Vec2ProcessorWithLM.from_pretrained("patrickvonplaten/wav2vec2-large-xlsr-53-spanish-with-lm")

        input_values = processor(resampled_audio, return_tensors="pt").input_values

        with torch.no_grad():
            logits = model(input_values.to(torch_device)).logits

        # test user-managed pool
        with multiprocessing.get_context("fork").Pool(2) as pool:
1699
            transcription = processor.batch_decode(logits.cpu().numpy(), pool).text
1700
1701
1702
1703
1704
1705
1706

        self.assertEqual(transcription[0], "bien y qué regalo vas a abrir primero")

        # user-managed pool + num_processes should trigger a warning
        with CaptureLogger(processing_wav2vec2_with_lm.logger) as cl, multiprocessing.get_context("fork").Pool(
            2
        ) as pool:
1707
            transcription = processor.batch_decode(logits.cpu().numpy(), pool, num_processes=2).text
1708
1709
1710
1711
1712
1713
1714
1715
1716

        self.assertIn("num_process", cl.out)
        self.assertIn("it will be ignored", cl.out)

        self.assertEqual(transcription[0], "bien y qué regalo vas a abrir primero")

    @require_pyctcdecode
    @require_torchaudio
    def test_wav2vec2_with_lm_invalid_pool(self):
1717
        run_test_in_subprocess(test_case=self, target_func=_test_wav2vec2_with_lm_invalid_pool, inputs=None)
1718

1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
    def test_inference_diarization(self):
        model = Wav2Vec2ForAudioFrameClassification.from_pretrained("anton-l/wav2vec2-base-superb-sd").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("anton-l/wav2vec2-base-superb-sd")
        input_data = self._load_superb("sd", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True, sampling_rate=16_000)

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)
        # labels is a one-hot array of shape (num_frames, num_speakers)
        labels = (outputs.logits > 0).long()

        # s3prl logits for the same batch
        expected_logits = torch.tensor(
            [
                [[-5.2807, -5.1272], [-5.4059, -4.7757], [-5.2764, -4.9621], [-5.0117, -4.5851]],
                [[-1.7643, -0.5462], [-1.7369, -0.2649], [-1.5066, -0.6200], [-4.5703, -2.4863]],
                [[-0.8656, -0.4783], [-0.8899, -0.3289], [-0.9267, -0.5781], [-0.7817, -0.4619]],
                [[-4.8625, -2.5316], [-5.2339, -2.2155], [-4.9835, -2.0344], [-4.4727, -1.8421]],
            ],
            device=torch_device,
        )
        self.assertEqual(labels[0, :, 0].sum(), 555)
        self.assertEqual(labels[0, :, 1].sum(), 299)
1744
1745
        # TODO: update the tolerance after the CI moves to torch 1.10
        self.assertTrue(torch.allclose(outputs.logits[:, :4], expected_logits, atol=1e-2))
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768

    def test_inference_speaker_verification(self):
        model = Wav2Vec2ForXVector.from_pretrained("anton-l/wav2vec2-base-superb-sv").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("anton-l/wav2vec2-base-superb-sv")
        input_data = self._load_superb("si", 4)

        inputs = processor(input_data["speech"], return_tensors="pt", padding=True, sampling_rate=16_000)
        labels = torch.tensor([5, 1, 1, 3], device=torch_device).T

        with torch.no_grad():
            input_values = inputs.input_values.to(torch_device)
            attention_mask = inputs.attention_mask.to(torch_device)
            outputs = model(input_values, attention_mask=attention_mask, labels=labels)
        embeddings = torch.nn.functional.normalize(outputs.embeddings, dim=-1).cpu()

        cosine_sim = torch.nn.CosineSimilarity(dim=-1)
        # id10002 vs id10002
        self.assertAlmostEqual(cosine_sim(embeddings[1], embeddings[2]).numpy(), 0.9758, 3)
        # id10006 vs id10002
        self.assertAlmostEqual(cosine_sim(embeddings[0], embeddings[1]).numpy(), 0.7579, 3)
        # id10002 vs id10004
        self.assertAlmostEqual(cosine_sim(embeddings[2], embeddings[3]).numpy(), 0.7594, 3)

1769
1770
        # TODO: update the tolerance after the CI moves to torch 1.10
        self.assertAlmostEqual(outputs.loss.item(), 17.7963, 2)