Adding TFWav2Vec2Model (#11617)

* [WIP] Add TFWav2Vec2Model

Work in progress for adding a tensorflow version of Wav2Vec2

* feedback changes

* small fix

* Test Feedback Round 1

* Add SpecAugment and CTC Loss

* correct spec augment mask creation

* docstring and correct copyright

* correct bugs

* remove bogus file

* finish tests correction

* del unnecessary layers

* Update src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

* make style

* correct final bug

* Feedback Changes

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Author: will-rice

docs/source/index.rst

@@ -399,7 +399,7 @@ Flax), PyTorch, and/or TensorFlow.
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
 |         VisualBert          |       ❌       |       ❌       |       ✅        |         ❌         |      ❌      |
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
-|          Wav2Vec2           |       ✅       |       ❌       |       ✅        |         ❌         |      ❌      |
+|          Wav2Vec2           |       ✅       |       ❌       |       ✅        |         ✅         |      ❌      |
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
 |             XLM             |       ✅       |       ❌       |       ✅        |         ✅         |      ❌      |
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+

docs/source/model_doc/wav2vec2.rst

@@ -80,9 +80,22 @@ Wav2Vec2ForCTC
 .. autoclass:: transformers.Wav2Vec2ForCTC
     :members: forward
 
-
 Wav2Vec2ForPreTraining
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.Wav2Vec2ForPreTraining
     :members: forward
+
+
+TFWav2Vec2Model
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFWav2Vec2Model
+    :members: call
+
+
+TFWav2Vec2ForCTC
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFWav2Vec2ForCTC
+    :members: call

src/transformers/__init__.py

@@ -1430,6 +1430,14 @@
             "TFTransfoXLPreTrainedModel",
         ]
     )
+    _import_structure["models.wav2vec2"].extend(
+        [
+            "TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST",
+            "TFWav2Vec2ForCTC",
+            "TFWav2Vec2Model",
+            "TFWav2Vec2PreTrainedModel",
+        ]
+    )
     _import_structure["models.xlm"].extend(
         [
             "TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -2743,6 +2751,12 @@
             TFTransfoXLModel,
             TFTransfoXLPreTrainedModel,
         )
+        from .models.wav2vec2 import (
+            TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
+            TFWav2Vec2ForCTC,
+            TFWav2Vec2Model,
+            TFWav2Vec2PreTrainedModel,
+        )
         from .models.xlm import (
             TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
             TFXLMForMultipleChoice,

src/transformers/convert_pytorch_checkpoint_to_tf2.py

@@ -37,6 +37,7 @@
     ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
     T5_PRETRAINED_CONFIG_ARCHIVE_MAP,
     TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
+    WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
     WEIGHTS_NAME,
     XLM_PRETRAINED_CONFIG_ARCHIVE_MAP,
     XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
@@ -79,10 +80,13 @@
     TFRobertaForSequenceClassification,
     TFT5ForConditionalGeneration,
     TFTransfoXLLMHeadModel,
+    TFWav2Vec2Model,
     TFXLMRobertaForMaskedLM,
     TFXLMWithLMHeadModel,
     TFXLNetLMHeadModel,
     TransfoXLConfig,
+    Wav2Vec2Config,
+    Wav2Vec2Model,
     XLMConfig,
     XLMRobertaConfig,
     XLNetConfig,
@@ -287,6 +291,12 @@
         ElectraForPreTraining,
         ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP,
     ),
+    "wav2vec2": (
+        Wav2Vec2Config,
+        TFWav2Vec2Model,
+        Wav2Vec2Model,
+        WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
+    ),
 }
 
 

src/transformers/models/auto/modeling_tf_auto.py

@@ -163,6 +163,7 @@
     TFTransfoXLLMHeadModel,
     TFTransfoXLModel,
 )
+from ..wav2vec2.modeling_tf_wav2vec2 import TFWav2Vec2Model
 from ..xlm.modeling_tf_xlm import (
     TFXLMForMultipleChoice,
     TFXLMForQuestionAnsweringSimple,
@@ -218,6 +219,7 @@
     RoFormerConfig,
     T5Config,
     TransfoXLConfig,
+    Wav2Vec2Config,
     XLMConfig,
     XLMRobertaConfig,
     XLNetConfig,
@@ -263,6 +265,7 @@
         (PegasusConfig, TFPegasusModel),
         (BlenderbotConfig, TFBlenderbotModel),
         (BlenderbotSmallConfig, TFBlenderbotSmallModel),
+        (Wav2Vec2Config, TFWav2Vec2Model),
     ]
 )
 

src/transformers/models/wav2vec2/__init__.py

@@ -17,7 +17,7 @@
 # limitations under the License.
 from typing import TYPE_CHECKING
 
-from ...file_utils import _BaseLazyModule, is_tokenizers_available, is_torch_available
+from ...file_utils import _BaseLazyModule, is_tf_available, is_torch_available
 
 
 _import_structure = {
@@ -38,6 +38,15 @@
     ]
 
 
+if is_tf_available():
+    _import_structure["modeling_tf_wav2vec2"] = [
+        "TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST",
+        "TFWav2Vec2ForCTC",
+        "TFWav2Vec2Model",
+        "TFWav2Vec2PreTrainedModel",
+    ]
+
+
 if TYPE_CHECKING:
     from .configuration_wav2vec2 import WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP, Wav2Vec2Config
     from .feature_extraction_wav2vec2 import Wav2Vec2FeatureExtractor
@@ -54,6 +63,14 @@
             Wav2Vec2PreTrainedModel,
         )
 
+    if is_tf_available():
+        from .modeling_tf_wav2vec2 import (
+            TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
+            TFWav2Vec2ForCTC,
+            TFWav2Vec2Model,
+            TFWav2Vec2PreTrainedModel,
+        )
+
 
 else:
     import importlib