Add pretrained weights from wav2vec2.0 and XLSR papers (#1827)

Add pretrained weights from https://github.com/pytorch/fairseq/tree/main/examples/wav2vec#pre-trained-models
- Wav2Vec 2.0 Base / Large / Large (LV-60)
- XLSR-53

Author: mthrok

docs/source/models.rst

@@ -139,6 +139,45 @@ Pre-trained Models
 
    .. autoproperty:: labels
 
+   .. autodata:: WAV2VEC2_BASE
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_BASE_10M
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_BASE_100H
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_BASE_960H
+      :no-value:
+
+   .. autodata:: WAV2VEC2_LARGE
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_LARGE_10M
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_LARGE_100H
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_LARGE_960H
+      :no-value:
+
+   .. autodata:: WAV2VEC2_LARGE_LV60K
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_LARGE_LV60K_10M
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_LARGE_LV60K_100H
+      :no-value:
+
+   .. autodata:: WAV2VEC2_ASR_LARGE_LV60K_960H
+      :no-value:
+
+   .. autodata:: WAV2VEC2_XLSR53
+      :no-value:
+
    .. autodata:: HUBERT_BASE
       :no-value:
 

test/integration_tests/wav2vec2_model_test.py

@@ -1,5 +1,18 @@
 import torchaudio
 from torchaudio.models import (
+    WAV2VEC2_BASE,
+    WAV2VEC2_LARGE,
+    WAV2VEC2_LARGE_LV60K,
+    WAV2VEC2_ASR_BASE_10M,
+    WAV2VEC2_ASR_BASE_100H,
+    WAV2VEC2_ASR_BASE_960H,
+    WAV2VEC2_ASR_LARGE_10M,
+    WAV2VEC2_ASR_LARGE_100H,
+    WAV2VEC2_ASR_LARGE_960H,
+    WAV2VEC2_ASR_LARGE_LV60K_10M,
+    WAV2VEC2_ASR_LARGE_LV60K_100H,
+    WAV2VEC2_ASR_LARGE_LV60K_960H,
+    WAV2VEC2_XLSR53,
     HUBERT_BASE,
     HUBERT_LARGE,
     HUBERT_XLARGE,
@@ -12,6 +25,10 @@
 @pytest.mark.parametrize(
     "bundle",
     [
+        WAV2VEC2_BASE,
+        WAV2VEC2_LARGE,
+        WAV2VEC2_LARGE_LV60K,
+        WAV2VEC2_XLSR53,
         HUBERT_BASE,
         HUBERT_LARGE,
         HUBERT_XLARGE,
@@ -25,6 +42,15 @@ def test_pretraining_models(bundle):
 @pytest.mark.parametrize(
     "bundle,expected",
     [
+        (WAV2VEC2_ASR_BASE_10M, 'I|HAD|THAT|CURIYOSSITY|BESID|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_BASE_100H, 'I|HAD|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_BASE_960H, 'I|HAD|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_LARGE_10M, 'I|HAD|THAT|CURIOUSITY|BESIDE|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_LARGE_100H, 'I|HAD|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_LARGE_960H, 'I|HAD|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_LARGE_LV60K_10M, 'I|HAD|THAT|CURIOUSSITY|BESID|ME|AT|THISS|MOMENT|'),
+        (WAV2VEC2_ASR_LARGE_LV60K_100H, 'I|HAVE|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
+        (WAV2VEC2_ASR_LARGE_LV60K_960H, 'I|HAVE|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
         (HUBERT_ASR_LARGE, 'I|HAVE|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|'),
         (HUBERT_ASR_XLARGE, 'I|HAVE|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|')
     ]

torchaudio/models/__init__.py

@@ -19,6 +19,19 @@
 )
 from .wav2vec2.pretrained import (
     Wav2Vec2PretrainedModelBundle,
+    WAV2VEC2_BASE,
+    WAV2VEC2_LARGE,
+    WAV2VEC2_LARGE_LV60K,
+    WAV2VEC2_ASR_BASE_10M,
+    WAV2VEC2_ASR_BASE_100H,
+    WAV2VEC2_ASR_BASE_960H,
+    WAV2VEC2_ASR_LARGE_10M,
+    WAV2VEC2_ASR_LARGE_100H,
+    WAV2VEC2_ASR_LARGE_960H,
+    WAV2VEC2_ASR_LARGE_LV60K_10M,
+    WAV2VEC2_ASR_LARGE_LV60K_100H,
+    WAV2VEC2_ASR_LARGE_LV60K_960H,
+    WAV2VEC2_XLSR53,
     HUBERT_BASE,
     HUBERT_LARGE,
     HUBERT_XLARGE,
@@ -45,6 +58,19 @@
     'hubert_ft_large',
     'hubert_ft_xlarge',
     'Wav2Vec2PretrainedModelBundle',
+    'WAV2VEC2_BASE',
+    'WAV2VEC2_LARGE',
+    'WAV2VEC2_LARGE_LV60K',
+    'WAV2VEC2_ASR_BASE_10M',
+    'WAV2VEC2_ASR_BASE_100H',
+    'WAV2VEC2_ASR_BASE_960H',
+    'WAV2VEC2_ASR_LARGE_10M',
+    'WAV2VEC2_ASR_LARGE_100H',
+    'WAV2VEC2_ASR_LARGE_960H',
+    'WAV2VEC2_ASR_LARGE_LV60K_10M',
+    'WAV2VEC2_ASR_LARGE_LV60K_100H',
+    'WAV2VEC2_ASR_LARGE_LV60K_960H',
+    'WAV2VEC2_XLSR53',
     'HUBERT_BASE',
     'HUBERT_LARGE',
     'HUBERT_XLARGE',