diff --git a/docs/source/index.mdx b/docs/source/index.mdx index f6fc659ed..9ee437711 100644 --- a/docs/source/index.mdx +++ b/docs/source/index.mdx @@ -227,7 +227,7 @@ Flax), PyTorch, and/or TensorFlow. | SEW | ❌ | ❌ | ✅ | ❌ | ❌ | | SEW-D | ❌ | ❌ | ✅ | ❌ | ❌ | | Speech Encoder decoder | ❌ | ❌ | ✅ | ❌ | ❌ | -| Speech2Text | ✅ | ❌ | ✅ | ❌ | ❌ | +| Speech2Text | ✅ | ❌ | ✅ | ✅ | ❌ | | Speech2Text2 | ✅ | ❌ | ❌ | ❌ | ❌ | | Splinter | ✅ | ✅ | ✅ | ❌ | ❌ | | SqueezeBERT | ✅ | ✅ | ✅ | ❌ | ❌ | diff --git a/docs/source/model_doc/auto.mdx b/docs/source/model_doc/auto.mdx index 3588445b0..21bda4319 100644 --- a/docs/source/model_doc/auto.mdx +++ b/docs/source/model_doc/auto.mdx @@ -202,6 +202,10 @@ Likewise, if your `NewModel` is a subclass of [`PreTrainedModel`], make sure its [[autodoc]] TFAutoModelForVision2Seq +## TFAutoModelForSpeechSeq2Seq + +[[autodoc]] TFAutoModelForSpeechSeq2Seq + ## FlaxAutoModel [[autodoc]] FlaxAutoModel diff --git a/docs/source/model_doc/speech_to_text.mdx b/docs/source/model_doc/speech_to_text.mdx index 1a8cce293..2e86c497c 100644 --- a/docs/source/model_doc/speech_to_text.mdx +++ b/docs/source/model_doc/speech_to_text.mdx @@ -144,3 +144,13 @@ See the [model hub](https://huggingface.co/models?filter=speech_to_text) to look [[autodoc]] Speech2TextForConditionalGeneration - forward + +## TFSpeech2TextModel + +[[autodoc]] TFSpeech2TextModel + - call + +## TFSpeech2TextForConditionalGeneration + +[[autodoc]] TFSpeech2TextForConditionalGeneration + - call diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index bb022ba36..18d7c35b6 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -1621,6 +1621,7 @@ if is_tf_available(): "TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING", "TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING", "TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING", + "TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING", "TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING", "TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING", "TF_MODEL_FOR_VISION_2_SEQ_MAPPING", @@ -1635,6 +1636,7 @@ if is_tf_available(): "TFAutoModelForQuestionAnswering", "TFAutoModelForSeq2SeqLM", "TFAutoModelForSequenceClassification", + "TFAutoModelForSpeechSeq2Seq", "TFAutoModelForTableQuestionAnswering", "TFAutoModelForTokenClassification", "TFAutoModelForVision2Seq", @@ -1946,6 +1948,14 @@ if is_tf_available(): "TFRoFormerPreTrainedModel", ] ) + _import_structure["models.speech_to_text"].extend( + [ + "TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST", + "TFSpeech2TextForConditionalGeneration", + "TFSpeech2TextModel", + "TFSpeech2TextPreTrainedModel", + ] + ) _import_structure["models.t5"].extend( [ "TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST", @@ -3588,6 +3598,7 @@ if TYPE_CHECKING: TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING, TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING, TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, TF_MODEL_FOR_VISION_2_SEQ_MAPPING, @@ -3602,6 +3613,7 @@ if TYPE_CHECKING: TFAutoModelForQuestionAnswering, TFAutoModelForSeq2SeqLM, TFAutoModelForSequenceClassification, + TFAutoModelForSpeechSeq2Seq, TFAutoModelForTableQuestionAnswering, TFAutoModelForTokenClassification, TFAutoModelForVision2Seq, @@ -3850,6 +3862,12 @@ if TYPE_CHECKING: TFRoFormerModel, TFRoFormerPreTrainedModel, ) + from .models.speech_to_text import ( + TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST, + TFSpeech2TextForConditionalGeneration, + TFSpeech2TextModel, + TFSpeech2TextPreTrainedModel, + ) from .models.t5 import ( TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST, TFT5EncoderModel, diff --git a/src/transformers/generation_tf_utils.py b/src/transformers/generation_tf_utils.py index f880cb368..d5de492bd 100644 --- a/src/transformers/generation_tf_utils.py +++ b/src/transformers/generation_tf_utils.py @@ -394,9 +394,12 @@ class TFGenerationMixin: Parameters: - input_ids (`tf.Tensor` of `dtype=tf.int32` and shape `(batch_size, sequence_length)`, *optional*): - The sequence used as a prompt for the generation. If `None` the method initializes it with - `bos_token_id` and a batch size of 1. + input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, `(batch_size, sequence_length, + feature_dim)` or `(batch_size, num_channels, height, width)`, *optional*): + The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the + method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs` + should of in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of + `input_ids`, `input_values`, `input_features`, or `pixel_values`. max_length (`int`, *optional*, defaults to 20): The maximum length of the sequence to be generated. min_length (`int`, *optional*, defaults to 10): @@ -657,11 +660,12 @@ class TFGenerationMixin: ), "Greedy beam search decoding cannot return more sequences than it has beams. Please set num_beams >= num_return_sequences" # create attention mask if necessary - # TODO (PVP): this should later be handled by the forward fn() in each model in the future see PR 3140 - if (attention_mask is None) and (pad_token_id is not None) and (pad_token_id in input_ids.numpy()): - attention_mask = tf.cast(tf.math.not_equal(input_ids, pad_token_id), dtype=tf.int32) - elif attention_mask is None: - attention_mask = tf.ones_like(input_ids) + accepts_attention_mask = "attention_mask" in set(inspect.signature(self.call).parameters.keys()) + if accepts_attention_mask: + if (attention_mask is None) and (pad_token_id is not None) and (pad_token_id in input_ids.numpy()): + attention_mask = tf.cast(tf.math.not_equal(input_ids, pad_token_id), dtype=tf.int32) + elif attention_mask is None: + attention_mask = tf.ones(shape_list(input_ids)[:2], dtype=tf.int32) if pad_token_id is None and eos_token_id is not None: logger.warning(f"Setting `pad_token_id` to {eos_token_id} (first `eos_token_id`) to generate sequence") @@ -697,16 +701,12 @@ class TFGenerationMixin: encoder = self.get_encoder() encoder_kwargs = { - "attention_mask": attention_mask, "output_attentions": output_attentions, "output_hidden_states": output_hidden_states, "return_dict": return_dict_in_generate, } - - # vision models don't use `attention_mask`. - signature = dict(inspect.signature(encoder.call).parameters) - if "attention_mask" not in signature: - encoder_kwargs.pop("attention_mask") + if accepts_attention_mask: + encoder_kwargs["attention_mask"] = attention_mask encoder_outputs = encoder(input_ids, **encoder_kwargs) if return_dict_in_generate: @@ -715,23 +715,15 @@ class TFGenerationMixin: if output_hidden_states: model_kwargs["encoder_hidden_states"] = encoder_outputs.hidden_states - # The condition `len(shape_list(input_ids)) == 2` is to make this block treats only text inputs. - # (vision inputs might occur when the model is an encoder-decoder model) - # Expand input ids if num_beams > 1 or num_return_sequences > 1 - if len(shape_list(input_ids)) == 2 and (num_return_sequences > 1 or num_beams > 1): - input_ids_len = shape_list(input_ids)[-1] - input_ids = tf.broadcast_to( - tf.expand_dims(input_ids, 1), (batch_size, effective_batch_mult * num_beams, input_ids_len) - ) - attention_mask = tf.broadcast_to( - tf.expand_dims(attention_mask, 1), (batch_size, effective_batch_mult * num_beams, input_ids_len) - ) - input_ids = tf.reshape( - input_ids, (effective_batch_size * num_beams, input_ids_len) - ) # shape: (batch_size * num_return_sequences * num_beams, cur_len) - attention_mask = tf.reshape( - attention_mask, (effective_batch_size * num_beams, input_ids_len) - ) # shape: (batch_size * num_return_sequences * num_beams, cur_len) + expanded_batch_idxs = tf.reshape( + tf.repeat(tf.expand_dims(tf.range(batch_size), -1), repeats=num_beams * effective_batch_mult, axis=1), + shape=(-1,), + ) + # prepares text-based inputs + if len(shape_list(input_ids)) == 2: + input_ids = tf.gather(input_ids, expanded_batch_idxs, axis=0) + if accepts_attention_mask: + attention_mask = tf.gather(attention_mask, expanded_batch_idxs, axis=0) if self.config.is_encoder_decoder: @@ -749,11 +741,6 @@ class TFGenerationMixin: batch_size == encoder_outputs[0].shape[0] ), f"expected encoder_outputs[0] to have 1st dimension bs={batch_size}, got {encoder_outputs[0].shape[0]} " - # expand batch_idx to assign correct encoder output for expanded input_ids (due to num_beams > 1 and num_return_sequences > 1) - expanded_batch_idxs = tf.reshape( - tf.repeat(tf.expand_dims(tf.range(batch_size), -1), repeats=num_beams * effective_batch_mult, axis=1), - shape=(-1,), - ) # expand encoder_outputs encoder_outputs = (tf.gather(encoder_outputs[0], expanded_batch_idxs, axis=0),) else: @@ -851,7 +838,8 @@ class TFGenerationMixin: unfinished_sents = tf.ones_like(input_ids[:, 0]) sent_lengths = tf.ones_like(input_ids[:, 0]) * max_length - past = encoder_outputs # defined for encoder-decoder models, None for decoder-only models + # defined for encoder-decoder models, None for decoder-only models + past = encoder_outputs # init attention / hidden states / scores tuples scores = () if (return_dict_in_generate and kwargs["output_scores"]) else None @@ -871,7 +859,11 @@ class TFGenerationMixin: while cur_len < max_length: model_inputs = self.prepare_inputs_for_generation( - input_ids, past=past, attention_mask=attention_mask, use_cache=use_cache, **kwargs + input_ids, + past=past, + attention_mask=attention_mask, + use_cache=use_cache, + **kwargs, ) outputs = self( **model_inputs, @@ -1132,7 +1124,11 @@ class TFGenerationMixin: while cur_len < max_length: model_inputs = self.prepare_inputs_for_generation( - input_ids, past=past, attention_mask=attention_mask, use_cache=use_cache, **kwargs + input_ids, + past=past, + attention_mask=attention_mask, + use_cache=use_cache, + **kwargs, ) outputs = self( **model_inputs, diff --git a/src/transformers/modeling_tf_pytorch_utils.py b/src/transformers/modeling_tf_pytorch_utils.py index 037525979..94e442bee 100644 --- a/src/transformers/modeling_tf_pytorch_utils.py +++ b/src/transformers/modeling_tf_pytorch_utils.py @@ -35,6 +35,7 @@ class TransposeType(ExplicitEnum): NO = "no" SIMPLE = "simple" + CONV1D = "conv1d" CONV2D = "conv2d" @@ -68,8 +69,9 @@ def convert_tf_weight_name_to_pt_weight_name(tf_name, start_prefix_to_remove="", # When should we transpose the weights if tf_name[-1] == "kernel" and tf_weight_shape is not None and tf_weight_shape.rank == 4: - # A simple heuristic to detect conv layer using weight array shape transpose = TransposeType.CONV2D + elif tf_name[-1] == "kernel" and tf_weight_shape is not None and tf_weight_shape.rank == 3: + transpose = TransposeType.CONV1D elif bool( tf_name[-1] in ["kernel", "pointwise_kernel", "depthwise_kernel"] or "emb_projs" in tf_name @@ -194,7 +196,6 @@ def load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs=None, a # authorized missing keys don't have to be loaded if any(re.search(pat, name) is not None for pat in tf_model._keys_to_ignore_on_load_missing): continue - raise AttributeError(f"{name} not found in PyTorch model") array = pt_state_dict[name].numpy() @@ -204,6 +205,11 @@ def load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs=None, a # PT: (num_out_channel, num_in_channel, kernel[0], kernel[1]) # -> TF: (kernel[0], kernel[1], num_in_channel, num_out_channel) array = numpy.transpose(array, axes=(2, 3, 1, 0)) + elif transpose is TransposeType.CONV1D: + # Conv1D weight: + # PT: (num_out_channel, num_in_channel, kernel) + # -> TF: (kernel, num_in_channel, num_out_channel) + array = numpy.transpose(array, axes=(2, 1, 0)) elif transpose is TransposeType.SIMPLE: array = numpy.transpose(array) @@ -355,7 +361,6 @@ def load_tf2_weights_in_pytorch_model(pt_model, tf_weights, allow_missing_keys=F all_tf_weights = set(list(tf_weights_map.keys())) loaded_pt_weights_data_ptr = {} missing_keys_pt = [] - for pt_weight_name, pt_weight in current_pt_params_dict.items(): # Handle PyTorch shared weight ()not duplicated in TF 2.0 if pt_weight.data_ptr() in loaded_pt_weights_data_ptr: @@ -377,6 +382,11 @@ def load_tf2_weights_in_pytorch_model(pt_model, tf_weights, allow_missing_keys=F # TF: (kernel[0], kernel[1], num_in_channel, num_out_channel) # -> PT: (num_out_channel, num_in_channel, kernel[0], kernel[1]) array = numpy.transpose(array, axes=(3, 2, 0, 1)) + elif transpose is TransposeType.CONV1D: + # Conv1D weight: + # TF: (kernel, num_in_channel, num_out_channel) + # -> PT: (num_out_channel, num_in_channel, kernel) + array = numpy.transpose(array, axes=(2, 1, 0)) elif transpose is TransposeType.SIMPLE: array = numpy.transpose(array) diff --git a/src/transformers/models/auto/__init__.py b/src/transformers/models/auto/__init__.py index 7798a51a8..9e39f38d0 100644 --- a/src/transformers/models/auto/__init__.py +++ b/src/transformers/models/auto/__init__.py @@ -87,6 +87,7 @@ if is_tf_available(): "TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING", "TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING", "TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING", + "TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING", "TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING", "TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING", "TF_MODEL_FOR_VISION_2_SEQ_MAPPING", @@ -101,6 +102,7 @@ if is_tf_available(): "TFAutoModelForQuestionAnswering", "TFAutoModelForSeq2SeqLM", "TFAutoModelForSequenceClassification", + "TFAutoModelForSpeechSeq2Seq", "TFAutoModelForTableQuestionAnswering", "TFAutoModelForTokenClassification", "TFAutoModelForVision2Seq", @@ -201,6 +203,7 @@ if TYPE_CHECKING: TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING, TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING, TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, TF_MODEL_FOR_VISION_2_SEQ_MAPPING, @@ -215,6 +218,7 @@ if TYPE_CHECKING: TFAutoModelForQuestionAnswering, TFAutoModelForSeq2SeqLM, TFAutoModelForSequenceClassification, + TFAutoModelForSpeechSeq2Seq, TFAutoModelForTableQuestionAnswering, TFAutoModelForTokenClassification, TFAutoModelForVision2Seq, diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 80190a9d2..fe472cc65 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -801,7 +801,7 @@ class AutoModelForSpeechSeq2Seq(_BaseAutoModelClass): AutoModelForSpeechSeq2Seq = auto_class_update( - AutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeing" + AutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeling" ) diff --git a/src/transformers/models/auto/modeling_tf_auto.py b/src/transformers/models/auto/modeling_tf_auto.py index 4cfbf9ba9..cd4158bc7 100644 --- a/src/transformers/models/auto/modeling_tf_auto.py +++ b/src/transformers/models/auto/modeling_tf_auto.py @@ -29,6 +29,7 @@ logger = logging.get_logger(__name__) TF_MODEL_MAPPING_NAMES = OrderedDict( [ # Base model mapping + ("speech_to_text", "TFSpeech2TextModel"), ("clip", "TFCLIPModel"), ("deberta-v2", "TFDebertaV2Model"), ("deberta", "TFDebertaModel"), @@ -103,6 +104,7 @@ TF_MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict( TF_MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict( [ # Model with LM heads mapping + ("speech_to_text", "TFSpeech2TextForConditionalGeneration"), ("rembert", "TFRemBertForMaskedLM"), ("roformer", "TFRoFormerForMaskedLM"), ("convbert", "TFConvBertForMaskedLM"), @@ -204,6 +206,12 @@ TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict( ] ) +TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict( + [ + ("speech_to_text", "TFSpeech2TextForConditionalGeneration"), + ] +) + TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( [ # Model for Sequence Classification mapping @@ -340,6 +348,9 @@ TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = _LazyAutoMapping( TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = _LazyAutoMapping( CONFIG_MAPPING_NAMES, TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES ) +TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = _LazyAutoMapping( + CONFIG_MAPPING_NAMES, TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES +) TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping( CONFIG_MAPPING_NAMES, TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES ) @@ -468,6 +479,15 @@ TFAutoModelForNextSentencePrediction = auto_class_update( ) +class TFAutoModelForSpeechSeq2Seq(_BaseAutoModelClass): + _model_mapping = TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING + + +TFAutoModelForSpeechSeq2Seq = auto_class_update( + TFAutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeling" +) + + class TFAutoModelWithLMHead(_TFAutoModelWithLMHead): @classmethod def from_config(cls, config): diff --git a/src/transformers/models/bart/modeling_tf_bart.py b/src/transformers/models/bart/modeling_tf_bart.py index 4b6999b70..bfe27588c 100644 --- a/src/transformers/models/bart/modeling_tf_bart.py +++ b/src/transformers/models/bart/modeling_tf_bart.py @@ -147,7 +147,11 @@ class TFBartAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder @@ -296,11 +300,11 @@ class TFBartEncoderLayer(tf.keras.layers.Layer): def call(self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, layer_head_mask: tf.Tensor, training=False): """ Args: - hidden_states (`tf.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`tf.Tensor`): attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size - *(encoder_attention_heads,)* + `(encoder_attention_heads,)` """ residual = hidden_states hidden_states, self_attn_weights, _ = self.self_attn( @@ -372,17 +376,17 @@ class TFBartDecoderLayer(tf.keras.layers.Layer): ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: """ Args: - hidden_states (`tf.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`tf.Tensor`): attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`tf.Tensor`): - cross attention input to the layer of shape *(seq_len, batch, embed_dim)* + cross attention input to the layer of shape `(seq_len, batch, embed_dim)` encoder_attention_mask (`tf.Tensor`): encoder attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size - *(decoder_attention_heads,)* + `(decoder_attention_heads,)` cross_attn_layer_head_mask (`tf.Tensor`): mask for heads of the cross-attention module. - *(decoder_attention_heads,)* + `(decoder_attention_heads,)` past_key_value (`Tuple(tf.Tensor)`): cached past key and value projection states """ residual = hidden_states diff --git a/src/transformers/models/blenderbot/modeling_tf_blenderbot.py b/src/transformers/models/blenderbot/modeling_tf_blenderbot.py index bccb24362..fc70c7c5a 100644 --- a/src/transformers/models/blenderbot/modeling_tf_blenderbot.py +++ b/src/transformers/models/blenderbot/modeling_tf_blenderbot.py @@ -150,7 +150,11 @@ class TFBlenderbotAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder diff --git a/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py b/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py index 43fb57c8e..ffb057c6a 100644 --- a/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py +++ b/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py @@ -149,7 +149,11 @@ class TFBlenderbotSmallAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder @@ -299,11 +303,11 @@ class TFBlenderbotSmallEncoderLayer(tf.keras.layers.Layer): def call(self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, layer_head_mask: tf.Tensor, training=False): """ Args: - hidden_states (`tf.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`tf.Tensor`): attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size - *(encoder_attention_heads,)* + `(encoder_attention_heads,)` """ residual = hidden_states hidden_states, self_attn_weights, _ = self.self_attn( @@ -376,17 +380,17 @@ class TFBlenderbotSmallDecoderLayer(tf.keras.layers.Layer): ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: """ Args: - hidden_states (`tf.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`tf.Tensor`): attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`tf.Tensor`): - cross attention input to the layer of shape *(seq_len, batch, embed_dim)* + cross attention input to the layer of shape `(seq_len, batch, embed_dim)` encoder_attention_mask (`tf.Tensor`): encoder attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size - *(decoder_attention_heads,)* + `(decoder_attention_heads,)` cross_attn_layer_head_mask (`tf.Tensor`): mask for heads of the cross-attention module. - *(decoder_attention_heads,)* + `(decoder_attention_heads,)` past_key_value (`Tuple(tf.Tensor)`): cached past key and value projection states """ residual = hidden_states diff --git a/src/transformers/models/hubert/modeling_tf_hubert.py b/src/transformers/models/hubert/modeling_tf_hubert.py index 79ae4ad93..a3a34b5e7 100644 --- a/src/transformers/models/hubert/modeling_tf_hubert.py +++ b/src/transformers/models/hubert/modeling_tf_hubert.py @@ -736,7 +736,11 @@ class TFHubertAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder diff --git a/src/transformers/models/marian/modeling_tf_marian.py b/src/transformers/models/marian/modeling_tf_marian.py index ba6b39d45..dad78bb3f 100644 --- a/src/transformers/models/marian/modeling_tf_marian.py +++ b/src/transformers/models/marian/modeling_tf_marian.py @@ -189,7 +189,11 @@ class TFMarianAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder @@ -339,11 +343,11 @@ class TFMarianEncoderLayer(tf.keras.layers.Layer): def call(self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, layer_head_mask: tf.Tensor, training=False): """ Args: - hidden_states (`tf.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`tf.Tensor`): attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size - *(encoder_attention_heads,)* + `(encoder_attention_heads,)` """ residual = hidden_states hidden_states, self_attn_weights, _ = self.self_attn( @@ -416,17 +420,17 @@ class TFMarianDecoderLayer(tf.keras.layers.Layer): ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: """ Args: - hidden_states (`tf.Tensor`): input to the layer of shape *(seq_len, batch, embed_dim)* + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` attention_mask (`tf.Tensor`): attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`tf.Tensor`): - cross attention input to the layer of shape *(seq_len, batch, embed_dim)* + cross attention input to the layer of shape `(seq_len, batch, embed_dim)` encoder_attention_mask (`tf.Tensor`): encoder attention mask of size - *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values. + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size - *(decoder_attention_heads,)* + `(decoder_attention_heads,)` cross_attn_layer_head_mask (`tf.Tensor`): mask for heads of the cross-attention module. - *(decoder_attention_heads,)* + `(decoder_attention_heads,)` past_key_value (`Tuple(tf.Tensor)`): cached past key and value projection states """ residual = hidden_states diff --git a/src/transformers/models/mbart/modeling_tf_mbart.py b/src/transformers/models/mbart/modeling_tf_mbart.py index 666ee4bea..91f9c003d 100644 --- a/src/transformers/models/mbart/modeling_tf_mbart.py +++ b/src/transformers/models/mbart/modeling_tf_mbart.py @@ -149,7 +149,11 @@ class TFMBartAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder diff --git a/src/transformers/models/pegasus/modeling_tf_pegasus.py b/src/transformers/models/pegasus/modeling_tf_pegasus.py index 73c3c13d7..853954683 100644 --- a/src/transformers/models/pegasus/modeling_tf_pegasus.py +++ b/src/transformers/models/pegasus/modeling_tf_pegasus.py @@ -190,7 +190,11 @@ class TFPegasusAttention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder diff --git a/src/transformers/models/speech_to_text/__init__.py b/src/transformers/models/speech_to_text/__init__.py index c0cb8f9a8..a61724c61 100644 --- a/src/transformers/models/speech_to_text/__init__.py +++ b/src/transformers/models/speech_to_text/__init__.py @@ -17,7 +17,13 @@ # limitations under the License. from typing import TYPE_CHECKING -from ...file_utils import _LazyModule, is_sentencepiece_available, is_speech_available, is_torch_available +from ...file_utils import ( + _LazyModule, + is_sentencepiece_available, + is_speech_available, + is_tf_available, + is_torch_available, +) _import_structure = { @@ -36,6 +42,14 @@ if is_speech_available(): if is_sentencepiece_available(): _import_structure["processing_speech_to_text"] = ["Speech2TextProcessor"] +if is_tf_available(): + _import_structure["modeling_tf_speech_to_text"] = [ + "TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST", + "TFSpeech2TextForConditionalGeneration", + "TFSpeech2TextModel", + "TFSpeech2TextPreTrainedModel", + ] + if is_torch_available(): _import_structure["modeling_speech_to_text"] = [ "SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST", @@ -57,6 +71,14 @@ if TYPE_CHECKING: if is_sentencepiece_available(): from .processing_speech_to_text import Speech2TextProcessor + if is_tf_available(): + from .modeling_tf_speech_to_text import ( + TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST, + TFSpeech2TextForConditionalGeneration, + TFSpeech2TextModel, + TFSpeech2TextPreTrainedModel, + ) + if is_torch_available(): from .modeling_speech_to_text import ( SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST, diff --git a/src/transformers/models/speech_to_text/modeling_tf_speech_to_text.py b/src/transformers/models/speech_to_text/modeling_tf_speech_to_text.py new file mode 100755 index 000000000..a6876bf9e --- /dev/null +++ b/src/transformers/models/speech_to_text/modeling_tf_speech_to_text.py @@ -0,0 +1,1615 @@ +# coding=utf-8 +# Copyright 2021 The Fairseq Authors and 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. +""" TensorFlow Speech2Text model.""" + + +import random +from typing import Dict, Optional, Tuple + +import tensorflow as tf + +from ...activations_tf import get_tf_activation, glu +from ...file_utils import ( + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ...modeling_tf_outputs import ( + TFBaseModelOutput, + TFBaseModelOutputWithPastAndCrossAttentions, + TFSeq2SeqLMOutput, + TFSeq2SeqModelOutput, +) +from ...modeling_tf_utils import ( + TFCausalLanguageModelingLoss, + TFPreTrainedModel, + TFSharedEmbeddings, + input_processing, + keras_serializable, + shape_list, +) +from ...utils import logging +from .configuration_speech_to_text import Speech2TextConfig + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "Speech2TextConfig" +_TOKENIZER_FOR_DOC = "Speech2TextTokenizer" +_CHECKPOINT_FOR_DOC = "facebook/s2t-small-librispeech-asr" + + +TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "facebook/s2t-small-librispeech-asr", + # See all Speech2Text models at https://huggingface.co/models?filter=speech_to_text +] + + +LARGE_NEGATIVE = -1e8 + + +# Copied from transformers.models.bart.modeling_tf_bart.shift_tokens_right +def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int): + pad_token_id = tf.cast(pad_token_id, input_ids.dtype) + decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype) + start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id) + shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1) + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids = tf.where( + shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids + ) + + if tf.executing_eagerly(): + # "Verify that `labels` has only positive values and -100" + assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype)) + + # Make sure the assertion op is called by wrapping the result in an identity no-op + with tf.control_dependencies([assert_gte0]): + shifted_input_ids = tf.identity(shifted_input_ids) + + return shifted_input_ids + + +# Copied from transformers.models.bart.modeling_tf_bart._make_causal_mask +def _make_causal_mask(input_ids_shape: tf.TensorShape, past_key_values_length: int = 0): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = tf.ones((tgt_len, tgt_len)) * LARGE_NEGATIVE + mask_cond = tf.range(shape_list(mask)[-1]) + + mask = tf.where(mask_cond < tf.reshape(mask_cond + 1, (shape_list(mask)[-1], 1)), 0.0, mask) + + if past_key_values_length > 0: + mask = tf.concat([tf.zeros((tgt_len, past_key_values_length)), mask], axis=-1) + + return tf.tile(mask[None, None, :, :], (bsz, 1, 1, 1)) + + +# Copied from transformers.models.bart.modeling_tf_bart._expand_mask +def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None, past_key_values_length: int = 0): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + src_len = shape_list(mask)[1] + tgt_len = tgt_len if tgt_len is not None else src_len + one_cst = tf.constant(1.0) + mask = tf.cast(mask, dtype=one_cst.dtype) + expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1)) + + return (one_cst - expanded_mask) * LARGE_NEGATIVE + + +class TFConv1dSubsampler(tf.keras.layers.Layer): + """ + Convolutional subsampler: a stack of 1D convolution (along temporal dimension) followed by non-linear activation + via gated linear units (https://arxiv.org/abs/1911.08460) + """ + + def __init__(self, config: Speech2TextConfig, **kwargs): + super().__init__(**kwargs) + self.config = config + self.num_layers = config.num_conv_layers + self.in_channels = config.input_feat_per_channel * config.input_channels + self.mid_channels = config.conv_channels + self.out_channels = config.d_model + self.kernel_sizes = config.conv_kernel_sizes + + self.conv_layers = [ + tf.keras.layers.Conv1D( + filters=self.mid_channels if i < self.num_layers - 1 else self.out_channels * 2, + kernel_size=k, + strides=2, + name=f"conv_layers.{i}", + ) + for i, k in enumerate(self.kernel_sizes) + ] + + def call(self, input_features: tf.Tensor) -> tf.Tensor: + hidden_states = tf.identity(input_features) # TF Conv1D assumes Batch x Time x Channels, same as the input + for i, conv in enumerate(self.conv_layers): + # equivalent to `padding=k // 2` on PT's `nn.Conv1d` + pad_len = self.kernel_sizes[i] // 2 + hidden_shapes = shape_list(hidden_states) + hidden_states = tf.concat( + ( + tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])), + hidden_states, + tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])), + ), + axis=1, + ) + + hidden_states = conv(hidden_states) + hidden_states = glu(hidden_states, axis=2) # GLU over the Channel dimension + return hidden_states + + +class TFSpeech2TextSinusoidalPositionalEmbedding(tf.keras.layers.Layer): + """This module produces sinusoidal positional embeddings of any length.""" + + def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None, **kwargs): + super().__init__(**kwargs) + self.offset = 2 + self.embedding_dim = embedding_dim + self.padding_idx = padding_idx + self.embedding_weights = self._get_embedding(num_positions + self.offset, embedding_dim, padding_idx) + + @staticmethod + def _get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None) -> tf.Tensor: + """ + Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the + description in Section 3.5 of "Attention Is All You Need". + """ + half_dim = embedding_dim // 2 + emb = tf.math.log(10000.0) / (half_dim - 1) + emb = tf.math.exp(tf.range(half_dim, dtype=tf.float32) * -emb) + emb = tf.expand_dims(tf.range(num_embeddings, dtype=tf.float32), axis=1) * tf.expand_dims(emb, axis=0) + emb = tf.reshape(tf.concat([tf.math.sin(emb), tf.math.cos(emb)], axis=1), shape=[num_embeddings, -1]) + if embedding_dim % 2 == 1: + # zero pad + emb = tf.concat([emb, tf.zeros(num_embeddings, 1)], axis=1) + if padding_idx is not None: + emb = tf.concat([emb[:padding_idx, :], tf.zeros((1, emb.shape[1])), emb[padding_idx + 1 :, :]], axis=0) + return emb + + def _resize_embeddings(self): + """Recreates (and effectivelly resizes) the sinusoidal embeddings""" + self.embeddings = self.add_weight( + name="weights", # name also used in PT + shape=self.embedding_weights.shape, + ) + self.embeddings.assign(self.embedding_weights) + + def build(self, input_shape: tf.TensorShape): + """ + Build shared token embedding layer Shared weights logic adapted from + https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24 + """ + self._resize_embeddings() + super().build(input_shape) + + def call(self, input_ids: tf.Tensor, past_key_values_length: int = 0) -> tf.Tensor: + bsz, seq_len = shape_list(input_ids) + # Create the position ids from the input token ids. Any padded tokens remain padded. + position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length) + + # expand embeddings if needed + max_pos = self.padding_idx + 1 + seq_len + if max_pos > shape_list(self.embeddings)[0]: + self.embedding_weights = self._get_embedding(max_pos + self.offset, self.embedding_dim, self.padding_idx) + self._resize_embeddings() + return tf.reshape(tf.gather(self.embeddings, tf.reshape(position_ids, (-1,)), axis=0), (bsz, seq_len, -1)) + + @staticmethod + def create_position_ids_from_input_ids( + input_ids: tf.Tensor, padding_idx: int, past_key_values_length: Optional[int] = 0 + ) -> tf.Tensor: + """ + Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding + symbols are ignored. This is modified from fairseq's `utils.make_positions`. + + Args: + x: tf.Tensor x: + Returns: tf.Tensor + """ + mask = tf.cast(tf.math.not_equal(input_ids, padding_idx), dtype=tf.int32) + incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask + return tf.cast(incremental_indices, dtype=tf.int64) + padding_idx + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartAttention with Bart->Speech2Text +class TFSpeech2TextAttention(tf.keras.layers.Layer): + """Multi-headed attention from "Attention Is All You Need""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + self.embed_dim = embed_dim + + self.num_heads = num_heads + self.dropout = tf.keras.layers.Dropout(dropout) + self.head_dim = embed_dim // num_heads + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim ** -0.5 + self.is_decoder = is_decoder + + self.k_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="k_proj") + self.q_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="q_proj") + self.v_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="v_proj") + self.out_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="out_proj") + + def _shape(self, tensor: tf.Tensor, seq_len: int, bsz: int): + return tf.transpose(tf.reshape(tensor, (bsz, seq_len, self.num_heads, self.head_dim)), (0, 2, 1, 3)) + + def call( + self, + hidden_states: tf.Tensor, + key_value_states: Optional[tf.Tensor] = None, + past_key_value: Optional[Tuple[Tuple[tf.Tensor]]] = None, + attention_mask: Optional[tf.Tensor] = None, + layer_head_mask: Optional[tf.Tensor] = None, + training=False, + ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + bsz, tgt_len, embed_dim = shape_list(hidden_states) + + # get query proj + query_states = self.q_proj(hidden_states) * self.scaling + # get key, value proj + if is_cross_attention and past_key_value is not None: + # reuse k,v, cross_attentions + key_states = past_key_value[0] + value_states = past_key_value[1] + elif is_cross_attention: + # cross_attentions + key_states = self._shape(self.k_proj(key_value_states), -1, bsz) + value_states = self._shape(self.v_proj(key_value_states), -1, bsz) + elif past_key_value is not None: + # reuse k, v, self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + key_states = tf.concat([past_key_value[0], key_states], axis=2) + value_states = tf.concat([past_key_value[1], value_states], axis=2) + else: + # self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + + if self.is_decoder: + # if cross_attention save Tuple(tf.Tensor, tf.Tensor) of all cross attention key/value_states. + # Further calls to cross_attention layer can then reuse all cross-attention + # key/value_states (first "if" case) + # if uni-directional self-attention (decoder) save Tuple(tf.Tensor, tf.Tensor) of + # all previous decoder key/value_states. Further calls to uni-directional self-attention + # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) + # if encoder bi-directional self-attention `past_key_value` is always `None` + past_key_value = (key_states, value_states) + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = tf.reshape(self._shape(query_states, tgt_len, bsz), proj_shape) + key_states = tf.reshape(key_states, proj_shape) + value_states = tf.reshape(value_states, proj_shape) + + src_len = shape_list(key_states)[1] + attn_weights = tf.matmul(query_states, key_states, transpose_b=True) + + # The tf.debugging asserts are not compliant with XLA then they + # have to be disabled in other modes than eager. + if tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(attn_weights), + [bsz * self.num_heads, tgt_len, src_len], + message=f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is {shape_list(attn_weights)}", + ) + + if attention_mask is not None: + # The tf.debugging asserts are not compliant with XLA then they + # have to be disabled in other modes than eager. + if tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(attention_mask), + [bsz, 1, tgt_len, src_len], + message=f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {shape_list(attention_mask)}", + ) + + attention_mask = tf.cast(attention_mask, dtype=attn_weights.dtype) + attn_weights = tf.reshape(attn_weights, (bsz, self.num_heads, tgt_len, src_len)) + attention_mask + attn_weights = tf.reshape(attn_weights, (bsz * self.num_heads, tgt_len, src_len)) + + attn_weights = tf.nn.softmax(attn_weights, axis=-1) + + if layer_head_mask is not None: + # The tf.debugging asserts are not compliant with XLA then they + # have to be disabled in other modes than eager. + if tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(layer_head_mask), + [self.num_heads], + message=f"Head mask for a single layer should be of size {(self.num_heads)}, but is {shape_list(layer_head_mask)}", + ) + + attn_weights = tf.reshape(layer_head_mask, (1, -1, 1, 1)) * tf.reshape( + attn_weights, (bsz, self.num_heads, tgt_len, src_len) + ) + attn_weights = tf.reshape(attn_weights, (bsz * self.num_heads, tgt_len, src_len)) + + attn_probs = self.dropout(attn_weights, training=training) + attn_output = tf.matmul(attn_probs, value_states) + + # The tf.debugging asserts are not compliant with XLA then they + # have to be disabled in other modes than eager. + if tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(attn_output), + [bsz * self.num_heads, tgt_len, self.head_dim], + message=f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is {shape_list(attn_output)}", + ) + + attn_output = tf.transpose( + tf.reshape(attn_output, (bsz, self.num_heads, tgt_len, self.head_dim)), (0, 2, 1, 3) + ) + attn_output = tf.reshape(attn_output, (bsz, tgt_len, embed_dim)) + + attn_output = self.out_proj(attn_output) + attn_weights: tf.Tensor = tf.reshape(attn_weights, (bsz, self.num_heads, tgt_len, src_len)) + + return attn_output, attn_weights, past_key_value + + +class TFSpeech2TextEncoderLayer(tf.keras.layers.Layer): + def __init__(self, config: Speech2TextConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFSpeech2TextAttention( + self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout, name="self_attn" + ) + self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm") + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.activation_fn = get_tf_activation(config.activation_function) + self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout) + self.fc1 = tf.keras.layers.Dense(config.encoder_ffn_dim, name="fc1") + self.fc2 = tf.keras.layers.Dense(self.embed_dim, name="fc2") + self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm") + + def call( + self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, layer_head_mask: tf.Tensor, training: bool = False + ): + """ + Args: + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size + `(encoder_attention_heads,)` + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states, self_attn_weights, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + training=training, + ) + + # The tf.debugging asserts are not compliant with XLA then they + # have to be disabled in other modes than eager. + if tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(hidden_states), + shape_list(residual), + message=f"Self attn modified the shape of query {shape_list(residual)} to {shape_list(hidden_states)}", + ) + + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = self.activation_dropout(hidden_states, training=training) + hidden_states = self.fc2(hidden_states) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + return hidden_states, self_attn_weights + + +class TFSpeech2TextDecoderLayer(tf.keras.layers.Layer): + def __init__(self, config: Speech2TextConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + + self.self_attn = TFSpeech2TextAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + name="self_attn", + is_decoder=True, + ) + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.activation_fn = get_tf_activation(config.activation_function) + self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout) + + self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm") + self.encoder_attn = TFSpeech2TextAttention( + self.embed_dim, + config.decoder_attention_heads, + dropout=config.attention_dropout, + name="encoder_attn", + is_decoder=True, + ) + self.encoder_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="encoder_attn_layer_norm") + self.fc1 = tf.keras.layers.Dense(config.decoder_ffn_dim, name="fc1") + self.fc2 = tf.keras.layers.Dense(self.embed_dim, name="fc2") + self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm") + + def call( + self, + hidden_states, + attention_mask: Optional[tf.Tensor] = None, + encoder_hidden_states: Optional[tf.Tensor] = None, + encoder_attention_mask: Optional[tf.Tensor] = None, + layer_head_mask: Optional[tf.Tensor] = None, + cross_attn_layer_head_mask: Optional[tf.Tensor] = None, + past_key_value: Optional[Tuple[tf.Tensor]] = None, + training=False, + ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: + """ + Args: + hidden_states (`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + encoder_hidden_states (`tf.Tensor`): + cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (`tf.Tensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size + `(decoder_attention_heads,)` + cross_attn_layer_head_mask (`tf.Tensor`): mask for heads of the cross-attention module. + `(decoder_attention_heads,)` + past_key_value (`Tuple(tf.Tensor)`): cached past key and value projection states + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Self Attention + # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 + self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None + # add present self-attn cache to positions 1,2 of present_key_value tuple + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + past_key_value=self_attn_past_key_value, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + training=training, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + # Cross-Attention Block + cross_attn_present_key_value = None + cross_attn_weights = None + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states = self.encoder_attn_layer_norm(hidden_states) + + # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple + cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None + hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + layer_head_mask=cross_attn_layer_head_mask, + past_key_value=cross_attn_past_key_value, + training=training, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + # add cross-attn to positions 3,4 of present_key_value tuple + present_key_value = present_key_value + cross_attn_present_key_value + + # Fully Connected + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = self.activation_dropout(hidden_states, training=training) + hidden_states = self.fc2(hidden_states) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + return ( + hidden_states, + self_attn_weights, + cross_attn_weights, + present_key_value, + ) + + +class TFSpeech2TextPreTrainedModel(TFPreTrainedModel): + config_class = Speech2TextConfig + base_model_prefix = "model" + main_input_name = "input_features" + + # Overwritten property due to different expected input shape and type + @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, + random.randint(1, self.config.max_source_positions), # time + self.config.input_feat_per_channel * self.config.input_channels, # input channels + ] + ), + "decoder_input_ids": tf.constant([[2, 3]], dtype=tf.int32), + } + + def _get_feat_extract_output_lengths(self, input_lengths: tf.Tensor): + """ + Computes the output length of the convolutional layers + """ + for _ in range(self.config.num_conv_layers): + input_lengths = (input_lengths - 1) // 2 + 1 + + return input_lengths + + @tf.function( + input_signature=[ + { + "input_ids": tf.TensorSpec((None, None), tf.int32, name="input_ids"), + "attention_mask": tf.TensorSpec((None, None), tf.int32, name="attention_mask"), + "decoder_input_ids": tf.TensorSpec((None, None), tf.int32, name="decoder_input_ids"), + "decoder_attention_mask": tf.TensorSpec((None, None), tf.int32, name="decoder_attention_mask"), + } + ] + ) + def serving(self, inputs): + output = self.call(inputs) + + return self.serving_output(output) + + +SPEECH_TO_TEXT_START_DOCSTRING = r""" + This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it + as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and + behavior. + + + + TF 2.0 models accepts two formats as inputs: + + - having all inputs as keyword arguments (like PyTorch models), or + - having all inputs as a list, tuple or dict in the first positional arguments. + + This second option is useful when using [`tf.keras.Model.fit`] method which currently requires having all the + tensors in the first argument of the model call function: `model(inputs)`. + + If you choose this second option, there are three possibilities you can use to gather all the input Tensors in the + first positional argument : + + - a single Tensor with `input_ids` only and nothing else: `model(input_ids)` + - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: + `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])` + - a dictionary with one or several input Tensors associated to the input names given in the docstring: + `model({"input_ids": input_ids, "token_type_ids": token_type_ids})` + + + + Parameters: + config ([`Speech2TextConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~TFPreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +SPEECH_TO_TEXT_INPUTS_DOCSTRING = r""" + Args: + input_features (`tf.Tensor` of shape `(batch_size, sequence_length, feature_size)`): + Float values of fbank features extracted from the raw speech waveform. Raw speech waveform can be obtained + by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* + via the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the + [`Speech2TextTokenizer`] should be used for extracting the fbank features, padding and conversion into a + tensor of floats. See [`~Speech2TextTokenizer.__call__`] + attention_mask (`tf.Tensor` of shape `({0})`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + decoder_input_ids (`tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`Speech2TextTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + Bart uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` + is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). + + For translation and summarization training, `decoder_input_ids` should be provided. If no + `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right + for denoising pre-training following the paper. + decoder_attention_mask (`tf.Tensor` of shape `(batch_size, target_sequence_length)`, *optional*): + will be made by default and ignore pad tokens. It is not recommended to set this for most use cases. + head_mask (`tf.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules in the encoder. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + decoder_head_mask (`tf.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules in the decoder. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + cross_attn_head_mask (`tf.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + encoder_outputs (`tf.FloatTensor`, *optional*): + hidden states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. + of shape `(batch_size, sequence_length, hidden_size)` is a sequence of + past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers`) + contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the + config will be used instead. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. This argument can be used only in eager mode, in graph mode the value in the config will be + used instead. + return_dict (`bool`, *optional*): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. This argument can be used + in eager mode, in graph mode the value will always be set to True. + training (`bool`, *optional*, defaults to `False`): + Whether or not to use the model in training mode (some modules like dropout modules have different + behaviors between training and evaluation). +""" + + +@keras_serializable +class TFSpeech2TextEncoder(tf.keras.layers.Layer): + config_class = Speech2TextConfig + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + [`TFSpeech2TextEncoderLayer`]. + + Args: + config: Speech2TextConfig + """ + + def __init__(self, config: Speech2TextConfig, **kwargs): + super().__init__(**kwargs) + self.config = config + + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.layerdrop = config.encoder_layerdrop + + embed_dim = config.d_model + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_source_positions + self.embed_scale = tf.math.sqrt(float(embed_dim)) if config.scale_embedding else 1.0 + + self.conv = TFConv1dSubsampler(config, name="conv") + + self.embed_positions = TFSpeech2TextSinusoidalPositionalEmbedding( + num_positions=config.max_source_positions, + embedding_dim=embed_dim, + padding_idx=self.padding_idx, + name="embed_positions", + ) + self.layers = [TFSpeech2TextEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)] + self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm") + + def _get_feat_extract_output_lengths(self, input_lengths: tf.Tensor): + """ + Computes the output length of the convolutional layers + """ + for _ in range(self.config.num_conv_layers): + input_lengths = (input_lengths - 1) // 2 + 1 + + return input_lengths + + def _get_feature_vector_attention_mask(self, feature_vector_length, attention_mask): + # generate creates 3D attention mask, because of the shape of input_features + # convert it to 2D if thats the case + if len(attention_mask.shape) > 2: + attention_mask = attention_mask[:, :, -1] + + subsampled_lengths = self._get_feat_extract_output_lengths(tf.math.reduce_sum(attention_mask, -1)) + bsz = shape_list(attention_mask)[0] + indices = tf.concat( + ( + tf.expand_dims(tf.range(bsz, dtype=attention_mask.dtype), -1), + tf.expand_dims(subsampled_lengths - 1, -1), + ), + axis=-1, + ) + + attention_mask = tf.scatter_nd(indices=indices, updates=tf.ones(bsz), shape=[bsz, feature_vector_length]) + attention_mask = tf.cast(tf.reverse(tf.math.cumsum(tf.reverse(attention_mask, [-1]), -1), [-1]), tf.int64) + return attention_mask + + def call( + self, + input_features=None, + attention_mask=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + """ + Args: + input_features (`tf.Tensor` of shape `(batch_size, sequence_length, feature_size)`): + Float values of fbank features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`Speech2TextTokenizer`] should be used for extracting the fbank features, + padding and conversion into a tensor of floats. See [`~Speech2TextTokenizer.__call__`] + attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + head_mask (`tf.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, `optional): + Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_features, + attention_mask=attention_mask, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + training=training, + kwargs_call=kwargs, + ) + if "input_ids" in inputs: + inputs["input_features"] = inputs.pop("input_ids") + + if inputs["input_features"] is None: + raise ValueError("You have to specify input_features") + + inputs_embeds = self.conv(inputs["input_features"]) + inputs_embeds = self.embed_scale * inputs_embeds + + # subsample attention mask if necessary + if inputs["attention_mask"] is not None: + inputs["attention_mask"] = self._get_feature_vector_attention_mask( + inputs_embeds.shape[1], inputs["attention_mask"] + ) + padding_mask = tf.cast(tf.math.not_equal(inputs["attention_mask"], 1), tf.int64) + else: + padding_mask = tf.zeros(inputs_embeds.shape[:-1], dtype=tf.int64) + + embed_pos = self.embed_positions(padding_mask) + + hidden_states = inputs_embeds + embed_pos + hidden_states = self.dropout(hidden_states, training=inputs["training"]) + + # check attention mask and invert + if inputs["attention_mask"] is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + inputs["attention_mask"] = _expand_mask(inputs["attention_mask"]) + + encoder_states = () if inputs["output_hidden_states"] else None + all_attentions = () if inputs["output_attentions"] else None + + # check if head_mask has a correct number of layers specified if desired + # The tf.debugging asserts are not compliant with XLA then they have to be disabled in other modes than eager. + if inputs["head_mask"] is not None and tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(inputs["head_mask"])[0], + len(self.layers), + message=f"The head_mask should be specified for {len(self.layers)} layers, but it is for {shape_list(inputs['head_mask'])[0]}.", + ) + + for idx, encoder_layer in enumerate(self.layers): + if inputs["output_hidden_states"]: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + dropout_probability = random.uniform(0, 1) + if inputs["training"] and (dropout_probability < self.layerdrop): # skip the layer + continue + + hidden_states, attn = encoder_layer( + hidden_states, + inputs["attention_mask"], + inputs["head_mask"][idx] if inputs["head_mask"] is not None else None, + training=training, + ) + + if inputs["output_attentions"]: + all_attentions += (attn,) + + hidden_states = self.layer_norm(hidden_states) + if inputs["output_hidden_states"]: + encoder_states = encoder_states + (hidden_states,) + + if not inputs["return_dict"]: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return TFBaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions + ) + + +@keras_serializable +class TFSpeech2TextDecoder(tf.keras.layers.Layer): + config_class = Speech2TextConfig + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`TFSpeech2TextDecoderLayer`] + + Args: + config: Speech2TextConfig + """ + + def __init__(self, config: Speech2TextConfig, **kwargs): + super().__init__(**kwargs) + self.config = config + self.layerdrop = config.decoder_layerdrop + self.padding_idx = config.pad_token_id + self.max_target_positions = config.max_target_positions + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + + self.embed_tokens = TFSharedEmbeddings(config.vocab_size, config.d_model, name="embed_tokens") + + self.embed_positions = TFSpeech2TextSinusoidalPositionalEmbedding( + num_positions=config.max_target_positions, + embedding_dim=config.d_model, + padding_idx=self.padding_idx, + name="embed_positions", + ) + + self.layers = [TFSpeech2TextDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)] + self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm") + + self.dropout = tf.keras.layers.Dropout(config.dropout) + + def get_embed_tokens(self): + return self.embed_tokens + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length): + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + combined_attention_mask = None + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length) + + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = _expand_mask(attention_mask, tgt_len=input_shape[-1]) + combined_attention_mask = ( + expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask + ) + + return combined_attention_mask + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + head_mask=None, + cross_attn_head_mask=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + r""" + Args: + input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you + provide it. + + Indices can be obtained using [`Speech2TextTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + encoder_hidden_states (`tf.Tensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention + of the decoder. + encoder_attention_mask (`tf.Tensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): + Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values + selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + head_mask (`tf.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + cross_attn_head_mask (`tf.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden-states of the attention blocks. Can be used to speed up + decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those + that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of + all ``decoder_input_ids``` of shape `(batch_size, sequence_length)`. inputs_embeds (`tf.Tensor` of + shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing + `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more + control over how to convert `input_ids` indices into associated vectors than the model's internal + embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + head_mask=head_mask, + cross_attn_head_mask=cross_attn_head_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + training=training, + kwargs_call=kwargs, + ) + + if inputs["input_ids"] is not None and inputs["inputs_embeds"] is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif inputs["input_ids"] is not None: + input_shape = shape_list(inputs["input_ids"]) + elif inputs["inputs_embeds"] is not None: + input_shape = shape_list(inputs["inputs_embeds"])[:-1] + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + # past_key_values_length + past_key_values_length = ( + shape_list(inputs["past_key_values"][0][0])[2] if inputs["past_key_values"] is not None else 0 + ) + + if inputs["inputs_embeds"] is None: + inputs_embeds = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + else: + inputs_embeds = inputs["inputs_embeds"] + + attention_mask = self._prepare_decoder_attention_mask( + inputs["attention_mask"], input_shape, inputs_embeds, past_key_values_length + ) + + # expand encoder attention mask + if inputs["encoder_hidden_states"] is not None and inputs["encoder_attention_mask"] is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + inputs["encoder_attention_mask"] = _expand_mask(inputs["encoder_attention_mask"], tgt_len=input_shape[-1]) + + # embed positions + positions = self.embed_positions(inputs["input_ids"], past_key_values_length=past_key_values_length) + + hidden_states = inputs_embeds + positions + hidden_states = self.dropout(hidden_states, training=inputs["training"]) + + # decoder layers + all_hidden_states = () if inputs["output_hidden_states"] else None + all_self_attns = () if inputs["output_attentions"] else None + all_cross_attns = () if (inputs["output_attentions"] and inputs["encoder_hidden_states"] is not None) else None + next_decoder_cache = () if inputs["use_cache"] else None + + # check if head_mask and cross_attn_head_mask have a correct number of layers specified if desired + # The tf.debugging asserts are not compliant with XLA then they have to be disabled in other modes than eager. + for attn_mask in ["head_mask", "cross_attn_head_mask"]: + if inputs[attn_mask] is not None and tf.executing_eagerly(): + tf.debugging.assert_equal( + shape_list(inputs[attn_mask])[0], + len(self.layers), + message=f"The {attn_mask} should be specified for {len(self.layers)} layers, but it is for {shape_list(inputs[attn_mask])[0]}.", + ) + + for idx, decoder_layer in enumerate(self.layers): + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + dropout_probability = random.uniform(0, 1) + if inputs["training"] and (dropout_probability < self.layerdrop): + continue + + past_key_value = inputs["past_key_values"][idx] if inputs["past_key_values"] is not None else None + cross_attn_layer_head_mask = ( + inputs["cross_attn_head_mask"][idx] if inputs["cross_attn_head_mask"] is not None else None + ) + + hidden_states, layer_self_attn, layer_cross_attn, present_key_value = decoder_layer( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=inputs["encoder_attention_mask"], + layer_head_mask=inputs["head_mask"][idx] if inputs["head_mask"] is not None else None, + cross_attn_layer_head_mask=cross_attn_layer_head_mask, + past_key_value=past_key_value, + ) + + if inputs["use_cache"]: + next_decoder_cache += (present_key_value,) + + if inputs["output_attentions"]: + all_self_attns += (layer_self_attn,) + + if inputs["encoder_hidden_states"] is not None: + all_cross_attns += (layer_cross_attn,) + + hidden_states = self.layer_norm(hidden_states) + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + + next_cache = (inputs["encoder_hidden_states"], next_decoder_cache) if use_cache else None + + if not inputs["return_dict"]: + return hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attns + else: + return TFBaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + cross_attentions=all_cross_attns, + ) + + +@keras_serializable +class TFSpeech2TextMainLayer(tf.keras.layers.Layer): + config_class = Speech2TextConfig + + def __init__(self, config: Speech2TextConfig, **kwargs): + super().__init__(**kwargs) + self.config = config + + self.encoder = TFSpeech2TextEncoder(config, name="encoder") + self.decoder = TFSpeech2TextDecoder(config, name="decoder") + + def get_input_embeddings(self): + return self.decoder.embed_tokens + + def set_input_embeddings(self, new_embeddings): + self.decoder.embed_tokens = new_embeddings + + def call( + self, + input_features=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + head_mask=None, + decoder_head_mask=None, + cross_attn_head_mask=None, + encoder_outputs=None, + past_key_values=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + head_mask=head_mask, + decoder_head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + training=training, + kwargs_call=kwargs, + ) + if "input_ids" in inputs: + inputs["input_features"] = inputs.pop("input_ids") + + # if the attribute is not set, fetch it from the config + for attr in ("output_attentions", "output_hidden_states", "use_cache"): + if inputs[attr] is None: + inputs[attr] = getattr(self.config, attr) + inputs["return_dict"] = ( + inputs["return_dict"] if inputs["return_dict"] is not None else self.config.use_return_dict + ) + + if inputs["encoder_outputs"] is None: + inputs["encoder_outputs"] = self.encoder( + input_features=inputs["input_features"], + attention_mask=inputs["attention_mask"], + head_mask=inputs["head_mask"], + output_attentions=inputs["output_attentions"], + output_hidden_states=inputs["output_hidden_states"], + return_dict=inputs["return_dict"], + training=inputs["training"], + ) + # If the user passed a tuple for encoder_outputs, we wrap it in a TFBaseModelOutput when return_dict=True + elif inputs["return_dict"] and not isinstance(inputs["encoder_outputs"], TFBaseModelOutput): + inputs["encoder_outputs"] = TFBaseModelOutput( + last_hidden_state=inputs["encoder_outputs"][0], + hidden_states=inputs["encoder_outputs"][1] if len(inputs["encoder_outputs"]) > 1 else None, + attentions=inputs["encoder_outputs"][2] if len(inputs["encoder_outputs"]) > 2 else None, + ) + # If the user passed a TFBaseModelOutput for encoder_outputs, we wrap it in a tuple when return_dict=False + elif not inputs["return_dict"] and not isinstance(inputs["encoder_outputs"], tuple): + inputs["encoder_outputs"] = inputs["encoder_outputs"].to_tuple() + + # downsample encoder attention mask + if inputs["attention_mask"] is not None: + inputs["encoder_attention_mask"] = self.encoder._get_feature_vector_attention_mask( + inputs["encoder_outputs"][0].shape[1], inputs["attention_mask"] + ) + else: + inputs["encoder_attention_mask"] = None + + # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn) + decoder_outputs = self.decoder( + input_ids=inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=inputs["encoder_outputs"][0], + encoder_attention_mask=inputs["encoder_attention_mask"], + head_mask=inputs["decoder_head_mask"], + cross_attn_head_mask=inputs["cross_attn_head_mask"], + past_key_values=inputs["past_key_values"], + inputs_embeds=inputs["decoder_inputs_embeds"], + use_cache=inputs["use_cache"], + output_attentions=inputs["output_attentions"], + output_hidden_states=inputs["output_hidden_states"], + return_dict=inputs["return_dict"], + training=inputs["training"], + ) + + if not inputs["return_dict"]: + return decoder_outputs + inputs["encoder_outputs"] + + return TFSeq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=inputs["encoder_outputs"].last_hidden_state, + encoder_hidden_states=inputs["encoder_outputs"].hidden_states, + encoder_attentions=inputs["encoder_outputs"].attentions, + ) + + +@add_start_docstrings( + "The bare Speech2Text Model outputting raw hidden-states without any specific head on top.", + SPEECH_TO_TEXT_START_DOCSTRING, +) +class TFSpeech2TextModel(TFSpeech2TextPreTrainedModel): + def __init__(self, config: Speech2TextConfig, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + + self.model = TFSpeech2TextMainLayer(config, name="model") + + def get_encoder(self): + return self.model.encoder + + def get_decoder(self): + return self.model.decoder + + @add_start_docstrings_to_model_forward(SPEECH_TO_TEXT_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + processor_class=_TOKENIZER_FOR_DOC, + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=TFSeq2SeqModelOutput, + config_class=_CONFIG_FOR_DOC, + ) + def call( + self, + input_features=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + head_mask=None, + decoder_head_mask=None, + cross_attn_head_mask=None, + encoder_outputs=None, + past_key_values=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + head_mask=head_mask, + decoder_head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + training=training, + kwargs_call=kwargs, + ) + if "input_ids" in inputs: + inputs["input_features"] = inputs.pop("input_ids") + + outputs = self.model( + input_features=inputs["input_features"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + decoder_attention_mask=inputs["decoder_attention_mask"], + head_mask=inputs["head_mask"], + decoder_head_mask=inputs["decoder_head_mask"], + cross_attn_head_mask=inputs["cross_attn_head_mask"], + encoder_outputs=inputs["encoder_outputs"], + past_key_values=inputs["past_key_values"], + decoder_inputs_embeds=inputs["decoder_inputs_embeds"], + use_cache=inputs["use_cache"], + output_attentions=inputs["output_attentions"], + output_hidden_states=inputs["output_hidden_states"], + return_dict=inputs["return_dict"], + training=inputs["training"], + ) + + return outputs + + def serving_output(self, output): + pkv = tf.tuple(output.past_key_values)[1] if self.config.use_cache else None + dec_hs = tf.convert_to_tensor(output.decoder_hidden_states) if self.config.output_hidden_states else None + dec_attns = tf.convert_to_tensor(output.decoder_attentions) if self.config.output_attentions else None + cross_attns = tf.convert_to_tensor(output.cross_attentions) if self.config.output_attentions else None + enc_hs = tf.convert_to_tensor(output.encoder_hidden_states) if self.config.output_hidden_states else None + enc_attns = tf.convert_to_tensor(output.encoder_attentions) if self.config.output_attentions else None + + return TFSeq2SeqModelOutput( + last_hidden_state=output.last_hidden_state, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + cross_attentions=cross_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + +@add_start_docstrings( + "The Speech2Text Model with a language modeling head. Can be used for summarization.", + SPEECH_TO_TEXT_START_DOCSTRING, +) +class TFSpeech2TextForConditionalGeneration(TFSpeech2TextPreTrainedModel, TFCausalLanguageModelingLoss): + def __init__(self, config: Speech2TextConfig): + super().__init__(config) + self.model = TFSpeech2TextMainLayer(config, name="model") + self.lm_head = tf.keras.layers.Dense(self.config.vocab_size, use_bias=False, name="lm_head") + + def get_encoder(self): + return self.model.encoder + + def get_decoder(self): + return self.model.decoder + + def resize_token_embeddings(self, new_num_tokens: int) -> tf.Variable: + new_embeddings = super().resize_token_embeddings(new_num_tokens) + return new_embeddings + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + @add_start_docstrings_to_model_forward(SPEECH_TO_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + def call( + self, + input_features=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + head_mask=None, + decoder_head_mask=None, + cross_attn_head_mask=None, + encoder_outputs=None, + past_key_values=None, + decoder_inputs_embeds=None, + labels=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + r""" + labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> import tensorflow as tf + >>> from transformers import Speech2TextProcessor, TFSpeech2TextForConditionalGeneration + >>> from datasets import load_dataset + >>> import soundfile as sf + + >>> model = TFSpeech2TextForConditionalGeneration.from_pretrained( + ... "facebook/s2t-small-librispeech-asr", from_pt=True + ... ) + >>> processor = Speech2TextProcessor.from_pretrained("facebook/s2t-small-librispeech-asr") + + + >>> def map_to_array(batch): + ... speech, _ = sf.read(batch["file"]) + ... batch["speech"] = speech + ... return batch + + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> ds = ds.map(map_to_array) + >>> ds.set_format(type="tf") + + >>> input_features = processor( + ... ds["speech"][0], sampling_rate=16000, return_tensors="tf" + >>> ).input_features # Batch size 1 + >>> generated_ids = model.generate(input_features) + + >>> transcription = processor.batch_decode(generated_ids) + ```""" + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + head_mask=head_mask, + decoder_head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + labels=labels, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + training=training, + kwargs_call=kwargs, + ) + if "input_ids" in inputs: + inputs["input_features"] = inputs.pop("input_ids") + + inputs["return_dict"] = ( + inputs["return_dict"] if inputs["return_dict"] is not None else self.config.use_return_dict + ) + + if inputs["labels"] is not None: + if inputs["decoder_input_ids"] is None: + inputs["decoder_input_ids"] = shift_tokens_right( + inputs["labels"], self.config.pad_token_id, self.config.decoder_start_token_id + ) + + outputs = self.model( + input_features=inputs["input_features"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + encoder_outputs=inputs["encoder_outputs"], + decoder_attention_mask=inputs["decoder_attention_mask"], + head_mask=inputs["head_mask"], + decoder_head_mask=inputs["decoder_head_mask"], + cross_attn_head_mask=inputs["cross_attn_head_mask"], + past_key_values=inputs["past_key_values"], + decoder_inputs_embeds=inputs["decoder_inputs_embeds"], + use_cache=inputs["use_cache"], + output_attentions=inputs["output_attentions"], + output_hidden_states=inputs["output_hidden_states"], + return_dict=inputs["return_dict"], + training=inputs["training"], + ) + lm_logits = self.lm_head(outputs[0]) + masked_lm_loss = None if inputs["labels"] is None else self.hf_compute_loss(inputs["labels"], lm_logits) + + if not inputs["return_dict"]: + output = (lm_logits,) + outputs[1:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return TFSeq2SeqLMOutput( + loss=masked_lm_loss, + logits=lm_logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + def serving_output(self, output): + pkv = tf.tuple(output.past_key_values)[1] if self.config.use_cache else None + dec_hs = tf.convert_to_tensor(output.decoder_hidden_states) if self.config.output_hidden_states else None + dec_attns = tf.convert_to_tensor(output.decoder_attentions) if self.config.output_attentions else None + cross_attns = tf.convert_to_tensor(output.cross_attentions) if self.config.output_attentions else None + enc_hs = tf.convert_to_tensor(output.encoder_hidden_states) if self.config.output_hidden_states else None + enc_attns = tf.convert_to_tensor(output.encoder_attentions) if self.config.output_attentions else None + + return TFSeq2SeqModelOutput( + last_hidden_state=output.last_hidden_state, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + cross_attentions=cross_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + def prepare_inputs_for_generation( + self, + decoder_input_ids, + past=None, + attention_mask=None, + head_mask=None, + decoder_head_mask=None, + cross_attn_head_mask=None, + use_cache=None, + **kwargs + ): + if past is not None and len(past) <= 2: + if not isinstance(past[0], tf.Tensor): + raise ValueError(f"`past[0]` has to be of type `tf.Tensor`, but is {type(past[0])}") + encoder_outputs = TFBaseModelOutput(last_hidden_state=past[0]) + if len(past) == 1: + past_key_values = None + else: + past_key_values = past[1] + if not past_key_values: + raise ValueError(f"decoder cached states must be truthy, got {past_key_values}") + decoder_input_ids = decoder_input_ids[:, -1:] + else: + raise ValueError(f"`past` must be an iterable with length 1 or 2, got {past}") + + return { + "input_features": None, # needs to be passed to make Keras.layer.__call__ happy + "encoder_outputs": encoder_outputs, + "past_key_values": past_key_values, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "head_mask": head_mask, + "decoder_head_mask": decoder_head_mask, + "cross_attn_head_mask": cross_attn_head_mask, + "use_cache": use_cache, # change this to avoid caching (presumably for debugging) + } + + @staticmethod + def _reorder_cache(past, beam_idx): + if len(past) == 1: + return past + + past_key_values = past[1] + + reordered_past = () + for layer_past_key_values in past_key_values: + reordered_past += ( + tuple(tf.gather(layer_past_key_value, beam_idx) for layer_past_key_value in layer_past_key_values[:2]) + + layer_past_key_values[2:], + ) + return (past[0], reordered_past) diff --git a/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py b/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py index fb44424a1..a493ee1eb 100644 --- a/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py +++ b/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py @@ -765,7 +765,11 @@ class TFWav2Vec2Attention(tf.keras.layers.Layer): self.num_heads = num_heads self.dropout = tf.keras.layers.Dropout(dropout) self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder diff --git a/src/transformers/utils/dummy_tf_objects.py b/src/transformers/utils/dummy_tf_objects.py index 777d01afc..02b401ef3 100644 --- a/src/transformers/utils/dummy_tf_objects.py +++ b/src/transformers/utils/dummy_tf_objects.py @@ -198,6 +198,9 @@ TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = None TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = None +TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = None + + TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING = None @@ -276,6 +279,13 @@ class TFAutoModelForSequenceClassification(metaclass=DummyObject): requires_backends(self, ["tf"]) +class TFAutoModelForSpeechSeq2Seq(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + class TFAutoModelForTableQuestionAnswering(metaclass=DummyObject): _backends = ["tf"] @@ -1678,6 +1688,30 @@ class TFRoFormerPreTrainedModel(metaclass=DummyObject): requires_backends(self, ["tf"]) +TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST = None + + +class TFSpeech2TextForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSpeech2TextModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSpeech2TextPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST = None diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 22cbd66b8..40885e1d2 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -1478,6 +1478,8 @@ class ModelTesterMixin: tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.float32) elif key == "pixel_values": tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.float32) + elif key == "input_features": + tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.float32) else: tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.int32) @@ -1529,6 +1531,8 @@ class ModelTesterMixin: tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.float32) elif key == "pixel_values": tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.float32) + elif key == "input_features": + tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.float32) else: tf_inputs_dict[key] = tf.convert_to_tensor(tensor.numpy(), dtype=tf.int32) diff --git a/tests/test_modeling_tf_common.py b/tests/test_modeling_tf_common.py index b100ac3ff..95c953a6e 100644 --- a/tests/test_modeling_tf_common.py +++ b/tests/test_modeling_tf_common.py @@ -57,6 +57,7 @@ if is_tf_available(): TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING, TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, BertConfig, TFAutoModel, @@ -140,6 +141,7 @@ class TFModelTesterMixin: *get_values(TF_MODEL_FOR_MASKED_LM_MAPPING), *get_values(TF_MODEL_FOR_PRETRAINING_MAPPING), *get_values(TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING), + *get_values(TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING), ]: inputs_dict["labels"] = tf.zeros( (self.model_tester.batch_size, self.model_tester.seq_length), dtype=tf.int32 @@ -358,7 +360,6 @@ class TFModelTesterMixin: pt_model = pt_model_class(config) # Check we can load pt model in tf and vice-versa with model => model functions - tf_model = transformers.load_pytorch_model_in_tf2_model( tf_model, pt_model, tf_inputs=self._prepare_for_class(inputs_dict, model_class) ) @@ -374,6 +375,8 @@ class TFModelTesterMixin: pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.float32) elif name == "pixel_values": pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.float32) + elif name == "input_features": + pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.float32) else: pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.long) @@ -416,6 +419,8 @@ class TFModelTesterMixin: pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.float32) elif name == "pixel_values": pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.float32) + elif name == "input_features": + pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.float32) else: pt_inputs_dict[name] = torch.from_numpy(key.numpy()).to(torch.long) @@ -443,7 +448,24 @@ class TFModelTesterMixin: metric = tf.keras.metrics.SparseCategoricalAccuracy("accuracy") for model_class in self.all_model_classes: - if self.is_encoder_decoder: + if model_class.__name__ in ["TFSpeech2TextModel", "TFSpeech2TextForConditionalGeneration"]: + inputs = { + "decoder_input_ids": tf.keras.Input( + batch_shape=(2, max_input), + name="decoder_input_ids", + dtype="int32", + ), + "input_features": tf.keras.Input( + batch_shape=( + 2, + max_input, + self.model_tester.input_feat_per_channel * self.model_tester.input_channels, + ), + name="input_features", + dtype="float32", + ), + } + elif self.is_encoder_decoder: inputs = { "decoder_input_ids": tf.keras.Input( batch_shape=(2, max_input), @@ -511,10 +533,7 @@ class TFModelTesterMixin: outputs_dict = model(inputs) inputs_keywords = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class)) - input_ids = inputs_keywords.pop("input_ids", None) - if input_ids is None: - input_ids = inputs_keywords.pop("pixel_values", None) - outputs_keywords = model(input_ids, **inputs_keywords) + outputs_keywords = model(**inputs_keywords) output_dict = outputs_dict[0].numpy() output_keywords = outputs_keywords[0].numpy() @@ -699,23 +718,28 @@ class TFModelTesterMixin: def test_model_common_attributes(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() - list_lm_models = ( + text_in_text_out_models = ( get_values(TF_MODEL_FOR_CAUSAL_LM_MAPPING) + get_values(TF_MODEL_FOR_MASKED_LM_MAPPING) + get_values(TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING) ) + speech_in_text_out_models = get_values(TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING) for model_class in self.all_model_classes: model = model_class(config) assert isinstance(model.get_input_embeddings(), tf.keras.layers.Layer) - - if model_class in list_lm_models: + if model_class in text_in_text_out_models: x = model.get_output_embeddings() assert isinstance(x, tf.keras.layers.Layer) name = model.get_bias() assert isinstance(name, dict) for k, v in name.items(): assert isinstance(v, tf.Variable) + elif model_class in speech_in_text_out_models: + x = model.get_output_embeddings() + assert isinstance(x, tf.keras.layers.Layer) + name = model.get_bias() + assert name is None else: x = model.get_output_embeddings() assert x is None @@ -922,13 +946,13 @@ class TFModelTesterMixin: model = model_class(config) if config.bos_token_id is None: - # if bos token id is not defined mobel needs input_ids + # if bos token id is not defined model needs input_ids with self.assertRaises(AssertionError): model.generate(do_sample=True, max_length=5) # num_return_sequences = 1 self._check_generated_ids(model.generate(input_ids, do_sample=True)) - else: - # num_return_sequences = 1 + elif model_class.__name__ not in ["TFSpeech2TextForConditionalGeneration"]: + # Models with non-text inputs won't work here; num_return_sequences = 1 self._check_generated_ids(model.generate(do_sample=True, max_length=5)) with self.assertRaises(AssertionError): @@ -952,6 +976,8 @@ class TFModelTesterMixin: def test_lm_head_model_no_beam_search_generate_dict_outputs(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() input_ids = inputs_dict.get("input_ids", None) + if input_ids is None: + input_ids = inputs_dict.get("input_features", None) # iterate over all generative models for model_class in self.all_generative_model_classes: @@ -988,7 +1014,7 @@ class TFModelTesterMixin: model = model_class(config) if config.bos_token_id is None: - # if bos token id is not defined mobel needs input_ids, num_return_sequences = 1 + # if bos token id is not defined model needs input_ids, num_return_sequences = 1 self._check_generated_ids(model.generate(input_ids, do_sample=True, num_beams=2)) else: # num_return_sequences = 1 @@ -1023,6 +1049,8 @@ class TFModelTesterMixin: def test_lm_head_model_beam_search_generate_dict_outputs(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() input_ids = inputs_dict.get("input_ids", None) + if input_ids is None: + input_ids = inputs_dict.get("input_features", None) # iterate over all generative models for model_class in self.all_generative_model_classes: @@ -1072,10 +1100,11 @@ class TFModelTesterMixin: # Test that model correctly compute the loss with kwargs prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True) - input_name = "input_ids" if "input_ids" in prepared_for_class else "pixel_values" - input_ids = prepared_for_class.pop(input_name) + possible_input_names = {"input_ids", "pixel_values", "input_features"} + input_name = possible_input_names.intersection(set(prepared_for_class)).pop() + model_input = prepared_for_class.pop(input_name) - loss = model(input_ids, **prepared_for_class)[0] + loss = model(model_input, **prepared_for_class)[0] self.assertEqual(loss.shape, [loss_size]) # Test that model correctly compute the loss with a dict diff --git a/tests/test_modeling_tf_speech_to_text.py b/tests/test_modeling_tf_speech_to_text.py new file mode 100644 index 000000000..e34892bf1 --- /dev/null +++ b/tests/test_modeling_tf_speech_to_text.py @@ -0,0 +1,605 @@ +# 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 TensorFlow Speech2Text model. """ + +import inspect +import unittest + +from transformers import Speech2TextConfig +from transformers.file_utils import cached_property, is_tf_available +from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow + +from .test_configuration_common import ConfigTester +from .test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor + + +if is_tf_available(): + import tensorflow as tf + + from transformers import Speech2TextProcessor, TFSpeech2TextForConditionalGeneration, TFSpeech2TextModel + + +def prepare_speech_to_text_inputs_dict( + config, + input_features, + decoder_input_ids, + attention_mask=None, + decoder_attention_mask=None, + head_mask=None, + decoder_head_mask=None, + cross_attn_head_mask=None, +): + if attention_mask is None: + attention_mask = tf.math.not_equal(input_features, 0) + if decoder_attention_mask is None: + decoder_attention_mask = tf.math.not_equal(decoder_input_ids, config.pad_token_id) + if head_mask is None: + head_mask = tf.ones((config.encoder_layers, config.encoder_attention_heads)) + if decoder_head_mask is None: + decoder_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads)) + if cross_attn_head_mask is None: + cross_attn_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads)) + return { + "input_features": input_features, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "decoder_attention_mask": attention_mask, + "head_mask": head_mask, + "decoder_head_mask": decoder_head_mask, + "cross_attn_head_mask": cross_attn_head_mask, + } + + +@require_tf +class TFSpeech2TextModelTester: + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=16, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=4, + num_conv_layers=2, + conv_kernel_sizes=(5, 5), + conv_channels=32, + input_feat_per_channel=24, + input_channels=1, + hidden_act="relu", + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=20, + max_source_positions=20, + max_target_positions=20, + eos_token_id=2, + pad_token_id=1, + bos_token_id=0, + ): + self.parent = parent + self.batch_size = batch_size + self.seq_length = seq_length + self.is_training = is_training + self.use_labels = use_labels + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.intermediate_size = intermediate_size + self.num_conv_layers = num_conv_layers + self.conv_kernel_sizes = conv_kernel_sizes + self.conv_channels = conv_channels + self.input_feat_per_channel = input_feat_per_channel + self.input_channels = input_channels + self.hidden_act = hidden_act + self.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + self.max_source_positions = max_source_positions + self.max_target_positions = max_target_positions + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + + def prepare_config_and_inputs(self): + input_features = floats_tensor( + [self.batch_size, self.seq_length, self.input_feat_per_channel], self.vocab_size + ) + attention_mask = tf.ones([self.batch_size, self.seq_length], dtype=tf.int64) + decoder_input_ids = tf.math.maximum(ids_tensor([self.batch_size, self.seq_length], self.vocab_size), 2) + + config = self.get_config() + inputs_dict = prepare_speech_to_text_inputs_dict( + config, + input_features=input_features, + decoder_input_ids=decoder_input_ids, + attention_mask=attention_mask, + ) + return config, inputs_dict + + def get_config(self): + return Speech2TextConfig( + vocab_size=self.vocab_size, + d_model=self.hidden_size, + encoder_layers=self.num_hidden_layers, + decoder_layers=self.num_hidden_layers, + encoder_attention_heads=self.num_attention_heads, + decoder_attention_heads=self.num_attention_heads, + encoder_ffn_dim=self.intermediate_size, + decoder_ffn_dim=self.intermediate_size, + num_conv_layers=self.num_conv_layers, + conv_kernel_sizes=self.conv_kernel_sizes, + conv_channels=self.conv_channels, + input_feat_per_channel=self.input_feat_per_channel, + input_channels=self.input_channels, + dropout=self.hidden_dropout_prob, + attention_dropout=self.attention_probs_dropout_prob, + max_position_embeddings=self.max_position_embeddings, + max_source_positions=self.max_source_positions, + max_target_positions=self.max_target_positions, + eos_token_id=self.eos_token_id, + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + ) + + def prepare_config_and_inputs_for_common(self): + config, inputs_dict = self.prepare_config_and_inputs() + return config, inputs_dict + + def get_subsampled_output_lengths(self, input_lengths): + """ + Computes the output length of the convolutional layers + """ + + for _ in range(self.num_conv_layers): + input_lengths = (input_lengths - 1) // 2 + 1 + + return input_lengths + + def create_and_check_decoder_model_past_large_inputs(self, config, inputs_dict): + model = TFSpeech2TextModel(config=config).get_decoder() + input_ids = inputs_dict["decoder_input_ids"] + attention_mask = inputs_dict["decoder_attention_mask"] + + # first forward pass + outputs = model(input_ids, attention_mask=attention_mask, use_cache=True) + + _, (_, past_key_values) = outputs.to_tuple() + + # create hypothetical multiple next token and extent to next_input_ids + next_tokens = tf.math.maximum(ids_tensor((self.batch_size, 3), config.vocab_size), 2) + next_attn_mask = ids_tensor((self.batch_size, 3), 2, dtype=tf.int64) + + # append to next input_ids and + next_input_ids = tf.concat([input_ids, next_tokens], axis=-1) + next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1) + + output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)["last_hidden_state"] + output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[ + "last_hidden_state" + ] + + # select random slice + random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1])) + output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx] + output_from_past_slice = output_from_past[:, :, random_slice_idx] + + self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]) + + # test that outputs are equal for slice + tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, atol=1e-2) + + +@require_tf +class TFSpeech2TextModelTest(TFModelTesterMixin, unittest.TestCase): + all_model_classes = (TFSpeech2TextModel, TFSpeech2TextForConditionalGeneration) if is_tf_available() else () + all_generative_model_classes = (TFSpeech2TextForConditionalGeneration,) if is_tf_available() else () + is_encoder_decoder = True + test_pruning = False + test_missing_keys = False + test_onnx = False + + input_name = "input_ids" + + def setUp(self): + self.model_tester = TFSpeech2TextModelTester(self) + self.config_tester = ConfigTester(self, config_class=Speech2TextConfig) + self.maxDiff = 3000 + + def test_config(self): + self.config_tester.run_common_tests() + + def test_decoder_model_past_with_large_inputs(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs) + + # not implemented currently + def test_inputs_embeds(self): + pass + + # training is not supported yet + def test_training(self): + pass + + def test_training_gradient_checkpointing(self): + pass + + def test_generate_fp16(self): + pass + + def test_hidden_states_output(self): + def check_hidden_states_output(inputs_dict, config, model_class): + model = model_class(config) + outputs = model(**self._prepare_for_class(inputs_dict, model_class)) + + hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states + + expected_num_layers = getattr( + self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1 + ) + self.assertEqual(len(hidden_states), expected_num_layers) + + if hasattr(self.model_tester, "encoder_seq_length"): + seq_length = self.model_tester.encoder_seq_length + else: + seq_length = self.model_tester.seq_length + + subsampled_seq_length = model._get_feat_extract_output_lengths(seq_length) + + self.assertListEqual( + list(hidden_states[0].shape[-2:]), + [subsampled_seq_length, self.model_tester.hidden_size], + ) + + if config.is_encoder_decoder: + hidden_states = outputs.decoder_hidden_states + + self.assertIsInstance(hidden_states, (list, tuple)) + self.assertEqual(len(hidden_states), expected_num_layers) + seq_len = getattr(self.model_tester, "seq_length", None) + decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len) + + self.assertListEqual( + list(hidden_states[0].shape[-2:]), + [decoder_seq_length, self.model_tester.hidden_size], + ) + + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + inputs_dict["output_hidden_states"] = True + check_hidden_states_output(inputs_dict, config, model_class) + + # check that output_hidden_states also work using config + del inputs_dict["output_hidden_states"] + config.output_hidden_states = True + + check_hidden_states_output(inputs_dict, config, model_class) + + def test_attention_outputs(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.return_dict = True + + seq_len = getattr(self.model_tester, "seq_length", None) + decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len) + encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", seq_len) + decoder_key_length = getattr(self.model_tester, "decoder_key_length", decoder_seq_length) + encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length) + + for model_class in self.all_model_classes: + inputs_dict["output_attentions"] = True + inputs_dict["output_hidden_states"] = False + config.return_dict = True + model = model_class(config) + + subsampled_encoder_seq_length = model._get_feat_extract_output_lengths(encoder_seq_length) + subsampled_encoder_key_length = model._get_feat_extract_output_lengths(encoder_key_length) + + outputs = model(**self._prepare_for_class(inputs_dict, model_class)) + attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions + self.assertEqual(len(attentions), self.model_tester.num_hidden_layers) + + # check that output_attentions also work using config + del inputs_dict["output_attentions"] + config.output_attentions = True + model = model_class(config) + + outputs = model(**self._prepare_for_class(inputs_dict, model_class)) + attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions + self.assertEqual(len(attentions), self.model_tester.num_hidden_layers) + + self.assertListEqual( + list(attentions[0].shape[-3:]), + [self.model_tester.num_attention_heads, subsampled_encoder_seq_length, subsampled_encoder_key_length], + ) + out_len = len(outputs) + + correct_outlen = 5 + + # loss is at first position + if "labels" in inputs_dict: + correct_outlen += 1 # loss is added to beginning + if "past_key_values" in outputs: + correct_outlen += 1 # past_key_values have been returned + + self.assertEqual(out_len, correct_outlen) + + # decoder attentions + decoder_attentions = outputs.decoder_attentions + self.assertIsInstance(decoder_attentions, (list, tuple)) + self.assertEqual(len(decoder_attentions), self.model_tester.num_hidden_layers) + self.assertListEqual( + list(decoder_attentions[0].shape[-3:]), + [self.model_tester.num_attention_heads, decoder_seq_length, decoder_key_length], + ) + + # cross attentions + cross_attentions = outputs.cross_attentions + self.assertIsInstance(cross_attentions, (list, tuple)) + self.assertEqual(len(cross_attentions), self.model_tester.num_hidden_layers) + self.assertListEqual( + list(cross_attentions[0].shape[-3:]), + [ + self.model_tester.num_attention_heads, + decoder_seq_length, + subsampled_encoder_key_length, + ], + ) + + # Check attention is always last and order is fine + inputs_dict["output_attentions"] = True + inputs_dict["output_hidden_states"] = True + model = model_class(config) + + outputs = model(**self._prepare_for_class(inputs_dict, model_class)) + + added_hidden_states = 2 + self.assertEqual(out_len + added_hidden_states, len(outputs)) + + self_attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions + + self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers) + self.assertListEqual( + list(self_attentions[0].shape[-3:]), + [self.model_tester.num_attention_heads, subsampled_encoder_seq_length, subsampled_encoder_key_length], + ) + + def test_resize_token_embeddings(self): + # Overwritten method from parent; see `test_resize_embeddings_untied` + pass + + def test_resize_tokens_embeddings(self): + # see `test_resize_embeddings_untied` + pass + + def test_resize_embeddings_untied(self): + # TODO: copy test from PT. Not working at the moment because the test relies on `model.resize_token_embeddings`, + # whose TF implementation assumes the use of `TFWrappedEmbeddings`. But with a `TFWrappedEmbeddings` we can't + # load the weights from PT (also, it induces TF1 behavior, so we might want to rework how + # `model.resize_token_embeddings` operates). + pass + + def test_generate_without_input_ids(self): + pass + + @staticmethod + def _get_encoder_outputs( + model, input_ids, attention_mask, output_attentions=None, output_hidden_states=None, num_interleave=1 + ): + encoder = model.get_encoder() + encoder_outputs = encoder( + input_ids, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + encoder_outputs["last_hidden_state"] = tf.repeat(encoder_outputs.last_hidden_state, num_interleave, axis=0) + + input_ids = input_ids[:, :, 0] + input_ids = tf.zeros_like(input_ids[:, :1], dtype=tf.int64) + model._get_decoder_start_token_id() + attention_mask = None + return encoder_outputs, input_ids, attention_mask + + def _check_outputs(self, output, input_ids, config, use_cache=False, num_return_sequences=1): + batch_size, seq_length = input_ids.shape[:2] + subsampled_seq_length = self.model_tester.get_subsampled_output_lengths(seq_length) + num_sequences_in_output = batch_size * num_return_sequences + gen_len = ( + output.sequences.shape[-1] - 1 if config.is_encoder_decoder else output.sequences.shape[-1] - seq_length + ) + + # scores + self._check_scores(num_sequences_in_output, output.scores, length=gen_len, config=config) + + # Attentions + # encoder + self._check_encoder_attention_for_generate( + output.encoder_attentions, batch_size, config, subsampled_seq_length + ) + # decoder + self._check_attentions_for_generate( + num_sequences_in_output, + output.decoder_attentions, + min_length=1, + max_length=output.sequences.shape[-1], + config=config, + use_cache=use_cache, + ) + + # Hidden States + # encoder + self._check_encoder_hidden_states_for_generate( + output.encoder_hidden_states, batch_size, config, subsampled_seq_length + ) + + # decoder + self._check_hidden_states_for_generate( + num_sequences_in_output, + output.decoder_hidden_states, + min_length=1, + max_length=output.sequences.shape[-1], + config=config, + use_cache=use_cache, + ) + + # overwritten from parent due to the inability to work when non-text inputs are not passed AND because the input is + # `input_features` + def test_lm_head_model_random_no_beam_search_generate(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + input_features = inputs_dict.get("input_features", None) + + # iterate over all generative models + for model_class in self.all_generative_model_classes: + model = model_class(config) + + if config.bos_token_id is None: + # if bos token id is not defined model needs input_features + with self.assertRaises(AssertionError): + model.generate(do_sample=True, max_length=5) + # num_return_sequences = 1 + self._check_generated_ids(model.generate(input_features, do_sample=True)) + + with self.assertRaises(AssertionError): + # generating multiple sequences when no beam search generation + # is not allowed as it would always generate the same sequences + model.generate(input_features, do_sample=False, num_return_sequences=2) + + # num_return_sequences > 1, sample + self._check_generated_ids(model.generate(input_features, do_sample=True, num_return_sequences=2)) + + # check bad words tokens language generation + # create list of 1-seq bad token and list of 2-seq of bad tokens + bad_words_ids = [self._generate_random_bad_tokens(1, model), self._generate_random_bad_tokens(2, model)] + output_tokens = model.generate( + input_features, do_sample=True, bad_words_ids=bad_words_ids, num_return_sequences=2 + ) + # only count generated tokens + generated_ids = output_tokens[:, input_features.shape[-1] :] + self.assertFalse(self._check_match_tokens(generated_ids.numpy().tolist(), bad_words_ids)) + + # overwritten from parent due to the inability to work when non-text inputs are not passed AND because the input is + # `input_features` + def test_lm_head_model_random_beam_search_generate(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + input_features = inputs_dict.get("input_features", None) + + for model_class in self.all_generative_model_classes: + model = model_class(config) + + if config.bos_token_id is None: + # if bos token id is not defined model needs input_ids, num_return_sequences = 1 + self._check_generated_ids(model.generate(input_features, do_sample=True, num_beams=2)) + + with self.assertRaises(AssertionError): + # generating more sequences than having beams leads is not possible + model.generate(input_features, do_sample=False, num_return_sequences=3, num_beams=2) + + # num_return_sequences > 1, sample + self._check_generated_ids( + model.generate( + input_features, + do_sample=True, + num_beams=2, + num_return_sequences=2, + ) + ) + # num_return_sequences > 1, greedy + self._check_generated_ids( + model.generate(input_features, do_sample=False, num_beams=2, num_return_sequences=2) + ) + + # check bad words tokens language generation + # create list of 1-seq bad token and list of 2-seq of bad tokens + bad_words_ids = [self._generate_random_bad_tokens(1, model), self._generate_random_bad_tokens(2, model)] + output_tokens = model.generate( + input_features, do_sample=False, bad_words_ids=bad_words_ids, num_beams=2, num_return_sequences=2 + ) + # only count generated tokens + generated_ids = output_tokens[:, input_features.shape[-1] :] + self.assertFalse(self._check_match_tokens(generated_ids.numpy().tolist(), bad_words_ids)) + + # overwritten from parent -- the input is `input_features`, not `input_ids` + def test_forward_signature(self): + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + signature = inspect.signature(model.call) + # signature.parameters is an OrderedDict => so arg_names order is deterministic + arg_names = [*signature.parameters.keys()] + + expected_arg_names = [ + "input_features", + "attention_mask", + "decoder_input_ids", + "decoder_attention_mask", + ] + self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names) + + +@require_tf +@require_sentencepiece +@require_tokenizers +@slow +class TFSpeech2TextModelIntegrationTests(unittest.TestCase): + @cached_property + def default_processor(self): + return Speech2TextProcessor.from_pretrained("facebook/s2t-small-librispeech-asr") + + def _load_datasamples(self, num_samples): + from datasets import load_dataset + + ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + # automatic decoding with librispeech + speech_samples = ds.sort("id").select(range(num_samples))[:num_samples]["audio"] + + return [x["array"] for x in speech_samples] + + def test_generation_librispeech(self): + model = TFSpeech2TextForConditionalGeneration.from_pretrained("facebook/s2t-small-librispeech-asr") + processor = self.default_processor + + input_speech = self._load_datasamples(1) + + input_features = processor(input_speech, return_tensors="tf").input_features + + generated_ids = model.generate(input_features) + generated_transcript = processor.batch_decode(generated_ids, skip_special_tokens=True) + + EXPECTED_TRANSCRIPTIONS = [ + "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel" + ] + self.assertListEqual(generated_transcript, EXPECTED_TRANSCRIPTIONS) + + def test_generation_librispeech_batched(self): + model = TFSpeech2TextForConditionalGeneration.from_pretrained("facebook/s2t-small-librispeech-asr") + processor = self.default_processor + + input_speech = self._load_datasamples(4) + + inputs = processor(input_speech, return_tensors="tf", padding=True) + generated_ids = model.generate(inputs.input_features, attention_mask=inputs.attention_mask) + generated_transcripts = processor.batch_decode(generated_ids, skip_special_tokens=True) + + EXPECTED_TRANSCRIPTIONS = [ + "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel", + "nor is mister cultar's manner less interesting than his matter", + "he tells us that at this festive season of the year with christmas and roast beef looming before us similes drawn from eating and its results occur most readily to the mind", + "he has grave doubts whether sir frederick leyton's work is really greek after all and can discover in it but little of rocky ithaca", + ] + self.assertListEqual(generated_transcripts, EXPECTED_TRANSCRIPTIONS)