Add tensor flow whisper model for audio classification

docs/source/en/model_doc/whisper.md

@@ -99,6 +99,10 @@ The original code can be found [here](https://github.com/openai/whisper).
 [[autodoc]] TFWhisperForConditionalGeneration
     - call
 
+## TFWhisperForAudioClassification
+
+[[autodoc]] TFWhisperForAudioClassification
+    - call
 
 ## FlaxWhisperModel
 

src/transformers/__init__.py

@@ -3705,6 +3705,7 @@
     _import_structure["models.whisper"].extend(
         [
             "TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
+            "TFWhisperForAudioClassification",
             "TFWhisperForConditionalGeneration",
             "TFWhisperModel",
             "TFWhisperPreTrainedModel",
@@ -7124,6 +7125,7 @@
         )
         from .models.whisper import (
             TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
+            TFWhisperForAudioClassification,
             TFWhisperForConditionalGeneration,
             TFWhisperModel,
             TFWhisperPreTrainedModel,

src/transformers/models/auto/modeling_tf_auto.py

@@ -357,7 +357,13 @@
         ("xlnet", "TFXLNetForQuestionAnsweringSimple"),
     ]
 )
-TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict([("wav2vec2", "TFWav2Vec2ForSequenceClassification")])
+
+TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+    [
+        ("wav2vec2", "TFWav2Vec2ForSequenceClassification"),
+        ("whisper", "TFWhisperForAudioClassification"),
+    ]
+)
 
 TF_MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
     [
@@ -461,6 +467,10 @@
 )
 
 TF_MODEL_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, TF_MODEL_MAPPING_NAMES)
+TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+    CONFIG_MAPPING_NAMES, TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
+)
+
 TF_MODEL_FOR_PRETRAINING_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, TF_MODEL_FOR_PRETRAINING_MAPPING_NAMES)
 TF_MODEL_WITH_LM_HEAD_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, TF_MODEL_WITH_LM_HEAD_MAPPING_NAMES)
 TF_MODEL_FOR_CAUSAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, TF_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)

src/transformers/models/whisper/__init__.py

@@ -63,6 +63,7 @@
         "TFWhisperForConditionalGeneration",
         "TFWhisperModel",
         "TFWhisperPreTrainedModel",
+        "TFWhisperForAudioClassification",
     ]
 
 try:
@@ -115,6 +116,7 @@
     else:
         from .modeling_tf_whisper import (
             TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
+            TFWhisperForAudioClassification,
             TFWhisperForConditionalGeneration,
             TFWhisperModel,
             TFWhisperPreTrainedModel,

src/transformers/models/whisper/modeling_tf_whisper.py

@@ -32,6 +32,7 @@
     TFBaseModelOutputWithPastAndCrossAttentions,
     TFSeq2SeqLMOutput,
     TFSeq2SeqModelOutput,
+    TFSequenceClassifierOutput,
 )
 from ...modeling_tf_utils import (
     TFCausalLanguageModelingLoss,
@@ -133,7 +134,7 @@ def build(self, input_shape):
 
     def call(self, input_ids, past_key_values_length=0):
         past_key_values_length = tf.cast(past_key_values_length, tf.int32)
-        gather_indices = tf.range(tf.shape(input_ids)[1], delta=1) + past_key_values_length
+        gather_indices = tf.range(tf.shape(input_ids)[-1], delta=1) + past_key_values_length
         return tf.gather(self.weight, gather_indices)
 
 
@@ -627,6 +628,12 @@ def __init__(self, config: WhisperConfig, **kwargs):
         self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm")
 
         self.dropout = tf.keras.layers.Dropout(config.dropout)
+
+    def get_input_embeddings(self):
+        return self.conv1
+
+    def set_input_embeddings(self, value):
+        self.conv1 = value
 
     @unpack_inputs
     def call(
@@ -1599,3 +1606,94 @@ def prepare_inputs_for_generation(
             "decoder_attention_mask": decoder_attention_mask,
             "decoder_position_ids": decoder_position_ids,
         }
+
+
+class TFWhisperForAudioClassification(TFWhisperPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.encoder = TFWhisperEncoder(config)
+        num_layers = config.num_hidden_layers + 1
+        if config.use_weighted_layer_sum:
+            self.layer_weights = tf.Variable(tf.ones(shape=(num_layers,)) / num_layers)
+        self.projector = tf.keras.layers.Dense(units=config.classifier_proj_size, input_shape=(config.hidden_size,))
+        self.classifier = tf.keras.layers.Dense(
+            units=config.num_labels, input_shape=(config.classifier_proj_size,), activation=None
+        )
+
+    def get_input_embeddings(self):
+        return self.encoder.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.encoder.set_input_embeddings(value)
+
+    @property
+    def dummy_inputs(self) -> Dict[str, tf.Tensor]:
+        """
+        Dummy inputs to build the network.
+
+        Returns:
+            `Dict[str, tf.Tensor]`: The dummy inputs.
+        """
+        return {
+            self.main_input_name: tf.random.uniform(
+                [1, self.config.num_mel_bins, self.config.max_source_positions * 2 - 1], dtype=tf.float32
+            ),
+        }
+
+    @property
+    def input_signature(self):
+        return {
+            "input_features": tf.TensorSpec((None, self.config.num_mel_bins, None), tf.float32, name="input_features"),
+        }
+
+    @unpack_inputs
+    def call(
+        self,
+        input_features: Optional[tf.Tensor] = None,
+        head_mask: Optional[tf.Tensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[tf.Tensor]]] = None,
+        labels: Optional[tf.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
+
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_features,
+                head_mask=head_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+        if self.config.use_weighted_layer_sum:
+            hidden_states = tf.stack(encoder_outputs, axis=1)
+            norm_weights = tf.nn.softmax(self.layer_weights, axis=-1)
+            hidden_states = tf.reduce_sum(hidden_states * tf.reshape(norm_weights, [-1, 1, 1]), axis=1)
+        else:
+            hidden_states = encoder_outputs[0]
+
+        hidden_states = self.projector(hidden_states)
+        pooled_output = tf.reduce_mean(hidden_states, axis=1)
+
+        logits = self.classifier(pooled_output)
+
+        loss = None
+
+        if labels is not None:
+            loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+            loss = loss_fn(tf.reshape(labels, [-1]), tf.reshape(logits, [-1, self.config.num_labels]))
+
+        if not return_dict:
+            output = (logits,) + encoder_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TFSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

src/transformers/utils/dummy_tf_objects.py

@@ -2699,6 +2699,13 @@ def __init__(self, *args, **kwargs):
 TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST = None
 
 
+class TFWhisperForAudioClassification(metaclass=DummyObject):
+    _backends = ["tf"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["tf"])
+
+
 class TFWhisperForConditionalGeneration(metaclass=DummyObject):
     _backends = ["tf"]