From 7f28613213b39a9ae831dfc68783aa1718ec0697 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Tue, 12 Jan 2021 02:06:32 +0100 Subject: [PATCH] [TFBart] Split TF-Bart (#9497) * make templates ready * make add_new_model_command_ready * finish tf bart * prepare tf mbart * finish tf bart * add tf mbart * add marian * prep pegasus * add tf pegasus * push blenderbot tf * add blenderbot * add blenderbot small * clean-up * make fix copy * define blend bot tok * fix * up * make style * add to docs * add copy statements * overwrite changes * improve * fix docs * finish * fix last slow test * fix missing git conflict line * fix blenderbot * up * fix blenderbot small * load changes * finish copied from * upload fix --- docs/source/index.rst | 2 +- docs/source/model_doc/blenderbot.rst | 11 +- docs/source/model_doc/blenderbot_small.rst | 14 + docs/source/model_doc/marian.rst | 8 + docs/source/model_doc/mbart.rst | 9 +- docs/source/model_doc/pegasus.rst | 8 + src/transformers/__init__.py | 20 +- .../models/auto/modeling_tf_auto.py | 20 +- .../models/auto/tokenization_auto.py | 5 +- .../models/bart/configuration_bart.py | 10 - .../models/bart/modeling_tf_bart.py | 324 ++-- .../models/blenderbot/__init__.py | 4 +- .../blenderbot/configuration_blenderbot.py | 11 - .../blenderbot/modeling_tf_blenderbot.py | 1304 +++++++++++++++- .../models/blenderbot_small/__init__.py | 10 +- .../modeling_blenderbot_small.py | 1 + .../modeling_tf_blenderbot_small.py | 1295 ++++++++++++++++ .../tokenization_blenderbot_small.py | 15 +- .../models/led/modeling_tf_led.py | 2 +- src/transformers/models/marian/__init__.py | 4 +- .../models/marian/configuration_marian.py | 11 - .../models/marian/modeling_tf_marian.py | 1307 +++++++++++++++- src/transformers/models/mbart/__init__.py | 4 +- .../models/mbart/configuration_mbart.py | 11 - .../models/mbart/modeling_tf_mbart.py | 1295 +++++++++++++++- src/transformers/models/pegasus/__init__.py | 4 +- .../models/pegasus/configuration_pegasus.py | 11 - .../models/pegasus/modeling_tf_pegasus.py | 1322 ++++++++++++++++- src/transformers/utils/dummy_tf_objects.py | 50 + ...tf_{{cookiecutter.lowercase_modelname}}.py | 56 +- ...ng_{{cookiecutter.lowercase_modelname}}.py | 2 +- ...tf_{{cookiecutter.lowercase_modelname}}.py | 21 +- ...ng_{{cookiecutter.lowercase_modelname}}.py | 10 +- tests/test_modeling_tf_bart.py | 167 ++- tests/test_modeling_tf_blenderbot.py | 217 ++- tests/test_modeling_tf_blenderbot_small.py | 328 ++++ tests/test_modeling_tf_marian.py | 198 ++- tests/test_modeling_tf_mbart.py | 204 ++- tests/test_modeling_tf_pegasus.py | 193 ++- 39 files changed, 7883 insertions(+), 605 deletions(-) create mode 100644 src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py create mode 100644 tests/test_modeling_tf_blenderbot_small.py diff --git a/docs/source/index.rst b/docs/source/index.rst index 2798bb9eec..fabed24506 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -225,7 +225,7 @@ TensorFlow and/or Flax. +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+ | Blenderbot | ✅ | ❌ | ✅ | ✅ | ❌ | +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+ -| BlenderbotSmall | ✅ | ❌ | ✅ | ❌ | ❌ | +| BlenderbotSmall | ✅ | ❌ | ✅ | ✅ | ❌ | +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+ | CTRL | ✅ | ❌ | ✅ | ✅ | ❌ | +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+ diff --git a/docs/source/model_doc/blenderbot.rst b/docs/source/model_doc/blenderbot.rst index 2f1ee4b021..43b4fb7a93 100644 --- a/docs/source/model_doc/blenderbot.rst +++ b/docs/source/model_doc/blenderbot.rst @@ -98,10 +98,15 @@ See :obj:`transformers.BartForConditionalGeneration` for arguments to `forward` :members: forward +TFBlenderbotModel +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: transformers.TFBlenderbotModel + :members: call + + TFBlenderbotForConditionalGeneration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -See :obj:`transformers.TFBartForConditionalGeneration` for arguments to `forward` and `generate` - .. autoclass:: transformers.TFBlenderbotForConditionalGeneration - :members: + :members: call diff --git a/docs/source/model_doc/blenderbot_small.rst b/docs/source/model_doc/blenderbot_small.rst index f44b0b73f2..ede238d47b 100644 --- a/docs/source/model_doc/blenderbot_small.rst +++ b/docs/source/model_doc/blenderbot_small.rst @@ -68,3 +68,17 @@ BlenderbotSmallForConditionalGeneration .. autoclass:: transformers.BlenderbotSmallForConditionalGeneration :members: forward + + +TFBlenderbotSmallModel +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: transformers.TFBlenderbotSmallModel + :members: call + + +TFBlenderbotSmallForConditionalGeneration +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: transformers.TFBlenderbotSmallForConditionalGeneration + :members: call diff --git a/docs/source/model_doc/marian.rst b/docs/source/model_doc/marian.rst index b7d97aae99..fba8e7d876 100644 --- a/docs/source/model_doc/marian.rst +++ b/docs/source/model_doc/marian.rst @@ -193,7 +193,15 @@ MarianMTModel :members: forward +TFMarianModel +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: transformers.TFMarianModel + :members: call + + TFMarianMTModel ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. autoclass:: transformers.TFMarianMTModel + :members: call diff --git a/docs/source/model_doc/mbart.rst b/docs/source/model_doc/mbart.rst index 202e498370..0a43d0f12c 100644 --- a/docs/source/model_doc/mbart.rst +++ b/docs/source/model_doc/mbart.rst @@ -124,8 +124,15 @@ MBartForSequenceClassification .. autoclass:: transformers.MBartForSequenceClassification +TFMBartModel +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: transformers.TFMBartModel + :members: call + + TFMBartForConditionalGeneration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. autoclass:: transformers.TFMBartForConditionalGeneration - :members: + :members: call diff --git a/docs/source/model_doc/pegasus.rst b/docs/source/model_doc/pegasus.rst index 9fc6a6a479..b27f7d074b 100644 --- a/docs/source/model_doc/pegasus.rst +++ b/docs/source/model_doc/pegasus.rst @@ -131,7 +131,15 @@ PegasusForConditionalGeneration :members: forward +TFPegasusModel +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: transformers.TFPegasusModel + :members: call + + TFPegasusForConditionalGeneration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. autoclass:: transformers.TFPegasusForConditionalGeneration + :members: call diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index e7d566f345..5ffaa001b5 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -868,7 +868,10 @@ if is_tf_available(): "TFBertPreTrainedModel", ] ) - _import_structure["models.blenderbot"].append("TFBlenderbotForConditionalGeneration") + _import_structure["models.blenderbot"].extend(["TFBlenderbotForConditionalGeneration", "TFBlenderbotModel"]) + _import_structure["models.blenderbot_small"].extend( + ["TFBlenderbotSmallForConditionalGeneration", "TFBlenderbotSmallModel"] + ) _import_structure["models.camembert"].extend( [ "TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST", @@ -986,8 +989,8 @@ if is_tf_available(): "TFLxmertVisualFeatureEncoder", ] ) - _import_structure["models.marian"].append("TFMarianMTModel") - _import_structure["models.mbart"].append("TFMBartForConditionalGeneration") + _import_structure["models.marian"].extend(["TFMarianMTModel", "TFMarianModel"]) + _import_structure["models.mbart"].extend(["TFMBartForConditionalGeneration", "TFMBartModel"]) _import_structure["models.mobilebert"].extend( [ "TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST", @@ -1028,7 +1031,7 @@ if is_tf_available(): "TFOpenAIGPTPreTrainedModel", ] ) - _import_structure["models.pegasus"].append("TFPegasusForConditionalGeneration") + _import_structure["models.pegasus"].extend(["TFPegasusForConditionalGeneration", "TFPegasusModel"]) _import_structure["models.roberta"].extend( [ "TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST", @@ -1855,7 +1858,8 @@ if TYPE_CHECKING: TFBertModel, TFBertPreTrainedModel, ) - from .models.blenderbot import TFBlenderbotForConditionalGeneration + from .models.blenderbot import TFBlenderbotForConditionalGeneration, TFBlenderbotModel + from .models.blenderbot_small import TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel from .models.camembert import ( TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST, TFCamembertForMaskedLM, @@ -1953,8 +1957,8 @@ if TYPE_CHECKING: TFLxmertPreTrainedModel, TFLxmertVisualFeatureEncoder, ) - from .models.marian import TFMarianMTModel - from .models.mbart import TFMBartForConditionalGeneration + from .models.marian import TFMarian, TFMarianMTModel + from .models.mbart import TFMBartForConditionalGeneration, TFMBartModel from .models.mobilebert import ( TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST, TFMobileBertForMaskedLM, @@ -1989,7 +1993,7 @@ if TYPE_CHECKING: TFOpenAIGPTModel, TFOpenAIGPTPreTrainedModel, ) - from .models.pegasus import TFPegasusForConditionalGeneration + from .models.pegasus import TFPegasusForConditionalGeneration, TFPegasusModel from .models.roberta import ( TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST, TFRobertaForMaskedLM, diff --git a/src/transformers/models/auto/modeling_tf_auto.py b/src/transformers/models/auto/modeling_tf_auto.py index 44243f8109..1fadbfcabe 100644 --- a/src/transformers/models/auto/modeling_tf_auto.py +++ b/src/transformers/models/auto/modeling_tf_auto.py @@ -44,7 +44,11 @@ from ..bert.modeling_tf_bert import ( TFBertLMHeadModel, TFBertModel, ) -from ..blenderbot.modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration +from ..blenderbot.modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration, TFBlenderbotModel +from ..blenderbot_small.modeling_tf_blenderbot_small import ( + TFBlenderbotSmallForConditionalGeneration, + TFBlenderbotSmallModel, +) from ..camembert.modeling_tf_camembert import ( TFCamembertForMaskedLM, TFCamembertForMultipleChoice, @@ -100,8 +104,8 @@ from ..longformer.modeling_tf_longformer import ( TFLongformerModel, ) from ..lxmert.modeling_tf_lxmert import TFLxmertForPreTraining, TFLxmertModel -from ..marian.modeling_tf_marian import TFMarianMTModel -from ..mbart.modeling_tf_mbart import TFMBartForConditionalGeneration +from ..marian.modeling_tf_marian import TFMarianModel, TFMarianMTModel +from ..mbart.modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel from ..mobilebert.modeling_tf_mobilebert import ( TFMobileBertForMaskedLM, TFMobileBertForMultipleChoice, @@ -122,7 +126,7 @@ from ..mpnet.modeling_tf_mpnet import ( ) from ..mt5.modeling_tf_mt5 import TFMT5ForConditionalGeneration, TFMT5Model from ..openai.modeling_tf_openai import TFOpenAIGPTForSequenceClassification, TFOpenAIGPTLMHeadModel, TFOpenAIGPTModel -from ..pegasus.modeling_tf_pegasus import TFPegasusForConditionalGeneration +from ..pegasus.modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel from ..roberta.modeling_tf_roberta import ( TFRobertaForMaskedLM, TFRobertaForMultipleChoice, @@ -167,6 +171,7 @@ from .configuration_auto import ( BartConfig, BertConfig, BlenderbotConfig, + BlenderbotSmallConfig, CamembertConfig, CTRLConfig, DistilBertConfig, @@ -225,6 +230,12 @@ TF_MODEL_MAPPING = OrderedDict( (FunnelConfig, TFFunnelModel), (DPRConfig, TFDPRQuestionEncoder), (MPNetConfig, TFMPNetModel), + (BartConfig, TFBartModel), + (MBartConfig, TFMBartModel), + (MarianConfig, TFMarianModel), + (PegasusConfig, TFPegasusModel), + (BlenderbotConfig, TFBlenderbotModel), + (BlenderbotSmallConfig, TFBlenderbotSmallModel), ] ) @@ -328,6 +339,7 @@ TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = OrderedDict( (MBartConfig, TFMBartForConditionalGeneration), (PegasusConfig, TFPegasusForConditionalGeneration), (BlenderbotConfig, TFBlenderbotForConditionalGeneration), + (BlenderbotSmallConfig, TFBlenderbotSmallForConditionalGeneration), (BartConfig, TFBartForConditionalGeneration), ] ) diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index 3a7b3c35b4..f46380ca0a 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -24,6 +24,7 @@ from ..bart.tokenization_bart import BartTokenizer from ..bert.tokenization_bert import BertTokenizer from ..bert_japanese.tokenization_bert_japanese import BertJapaneseTokenizer from ..bertweet.tokenization_bertweet import BertweetTokenizer +from ..blenderbot.tokenization_blenderbot import BlenderbotTokenizer from ..blenderbot_small.tokenization_blenderbot_small import BlenderbotSmallTokenizer from ..ctrl.tokenization_ctrl import CTRLTokenizer from ..deberta.tokenization_deberta import DebertaTokenizer @@ -58,6 +59,7 @@ from .configuration_auto import ( BertConfig, BertGenerationConfig, BlenderbotConfig, + BlenderbotSmallConfig, CamembertConfig, CTRLConfig, DebertaConfig, @@ -201,7 +203,8 @@ TOKENIZER_MAPPING = OrderedDict( (MBartConfig, (MBartTokenizer, MBartTokenizerFast)), (XLMRobertaConfig, (XLMRobertaTokenizer, XLMRobertaTokenizerFast)), (MarianConfig, (MarianTokenizer, None)), - (BlenderbotConfig, (BlenderbotSmallTokenizer, None)), + (BlenderbotSmallConfig, (BlenderbotSmallTokenizer, None)), + (BlenderbotConfig, (BlenderbotTokenizer, None)), (LongformerConfig, (LongformerTokenizer, LongformerTokenizerFast)), (BartConfig, (BartTokenizer, BartTokenizerFast)), (LongformerConfig, (LongformerTokenizer, LongformerTokenizerFast)), diff --git a/src/transformers/models/bart/configuration_bart.py b/src/transformers/models/bart/configuration_bart.py index 1b128df0dc..1bbdfabcd3 100644 --- a/src/transformers/models/bart/configuration_bart.py +++ b/src/transformers/models/bart/configuration_bart.py @@ -170,16 +170,6 @@ class BartConfig(PretrainedConfig): self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True self.force_bos_token_to_be_generated = force_bos_token_to_be_generated # only relevant for CNN - # IMPORTANT - # DELETE ALL OF THE FOLLOWING LINES AS SOON AS TF IS READY - self.extra_pos_embeddings = 2 - self.normalize_before = False - self.add_final_layer_norm = False - self.do_blenderbot_90_layernorm = False - self.normalize_embedding = True - self.static_position_embeddings = False - self.add_bias_logits = False - @property def num_attention_heads(self) -> int: return self.encoder_attention_heads diff --git a/src/transformers/models/bart/modeling_tf_bart.py b/src/transformers/models/bart/modeling_tf_bart.py index 10410d0a87..8578a18ebb 100644 --- a/src/transformers/models/bart/modeling_tf_bart.py +++ b/src/transformers/models/bart/modeling_tf_bart.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team. +# 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. @@ -12,19 +12,18 @@ # 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. -"""TF BART model, ported from the fairseq repo.""" +""" TF 2.0 Bart model. """ + -import math import random -import warnings from typing import Dict, Optional, Tuple, Union -import numpy as np import tensorflow as tf -from ...activations_tf import ACT2FN +from ...activations_tf import get_tf_activation from ...file_utils import ( add_code_sample_docstrings, + add_end_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings, @@ -55,13 +54,14 @@ logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "BartConfig" _TOKENIZER_FOR_DOC = "BartTokenizer" + LARGE_NEGATIVE = -1e8 -def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, eos_token_id: int): +def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int): shifted_input_ids = tf.cast(input_ids, tf.int32) shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1) - start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), eos_token_id) + start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), decoder_start_token_id) shifted_input_ids = tf.concat([start_tokens, shifted_input_ids[:, 1:]], -1) # replace possible -100 values in labels by `pad_token_id` shifted_input_ids = tf.where( @@ -94,7 +94,7 @@ def _make_causal_mask(input_ids_shape: tf.TensorShape, past_key_values_length: i return tf.broadcast_to(mask[None, None, :, :], (bsz, 1, tgt_len, tgt_len + past_key_values_length)) -def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None): +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]`. """ @@ -108,18 +108,15 @@ def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None): class TFBartLearnedPositionalEmbedding(TFSharedEmbeddings): """ - This module learns positional embeddings up to a fixed maximum size. Padding ids are ignored by either offsetting - based on padding_idx or by setting padding_idx to None and ensuring that the appropriate position ids are passed to - the forward function. + This module learns positional embeddings up to a fixed maximum size. """ - def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, offset, **kwargs): + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, **kwargs): + assert padding_idx is not None, "padding_idx cannot be None" # Bart is set up so that if padding_idx is specified then offset the embedding ids by 2 # and adjust num_embeddings appropriately. Other models dont have this hack - self.offset = offset - assert padding_idx is not None, "padding_idx cannot be None" - num_embeddings += offset - super().__init__(num_embeddings, embedding_dim, **kwargs) + self.offset = 2 + super().__init__(num_embeddings + self.offset, embedding_dim, **kwargs) def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): """Input is expected to be of size [bsz x seqlen].""" @@ -128,56 +125,7 @@ class TFBartLearnedPositionalEmbedding(TFSharedEmbeddings): positions = tf.range( past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" ) - return super().call(positions + self.offset) # super object is not callable for some reason - - -class TFBartSinusoidalPositionalEmbedding(tf.keras.layers.Embedding): - """This module produces sinusoidal positional embeddings of any length.""" - - def __init__(self, num_positions: int, embedding_dim: int, **kwargs): - - if embedding_dim % 2 != 0: - raise NotImplementedError(f"odd embedding_dim {embedding_dim} not supported") - super().__init__( - num_positions, - embedding_dim, - **kwargs, - ) - - 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 - """ - super().build(input_shape) # Instantiates self.weight so it can be loaded - weight: np.ndarray = self._init_weight(self.input_dim, self.output_dim) - self.set_weights([weight]) # overwrite self.weight to correct value - - @staticmethod - def _init_weight(n_pos: int, dim: int): - """ - Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in - the 2nd half of the vector. [dim // 2:] - """ - position_enc = np.array( - [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)] - ) - # index 0 is all zero - position_enc[:, 0 : dim // 2] = np.sin(position_enc[:, 0::2]) - position_enc[:, dim // 2 :] = np.cos(position_enc[:, 1::2]) - # convert to tensor - table = tf.convert_to_tensor(position_enc, dtype=tf.float32) - tf.stop_gradient(table) - return table - - def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): - """Input is expected to be of size [bsz x seqlen].""" - bsz, seq_len = input_shape[:2] - - positions = tf.range( - past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" - ) - return super().call(positions) + return super().call(positions + self.offset) class TFBartAttention(tf.keras.layers.Layer): @@ -310,10 +258,9 @@ class TFBartEncoderLayer(tf.keras.layers.Layer): self.self_attn = TFBartAttention( self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout, name="self_attn" ) - self.normalize_before = config.normalize_before 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 = ACT2FN[config.activation_function] + 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") @@ -327,8 +274,6 @@ class TFBartEncoderLayer(tf.keras.layers.Layer): `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. """ residual = hidden_states - if self.normalize_before: - 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 ) @@ -339,19 +284,15 @@ class TFBartEncoderLayer(tf.keras.layers.Layer): ) hidden_states = self.dropout(hidden_states, training=training) hidden_states = residual + hidden_states - if not self.normalize_before: - hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states = self.self_attn_layer_norm(hidden_states) residual = hidden_states - if self.normalize_before: - 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 - if not self.normalize_before: - hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.final_layer_norm(hidden_states) return hidden_states, self_attn_weights @@ -368,9 +309,8 @@ class TFBartDecoderLayer(tf.keras.layers.Layer): is_decoder=True, ) self.dropout = tf.keras.layers.Dropout(config.dropout) - self.activation_fn = ACT2FN[config.activation_function] + self.activation_fn = get_tf_activation(config.activation_function) self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout) - self.normalize_before = config.normalize_before self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm") self.encoder_attn = TFBartAttention( @@ -405,8 +345,6 @@ class TFBartDecoderLayer(tf.keras.layers.Layer): past_key_value (:obj:`Tuple(tf.Tensor)`): cached past key and value projection states """ residual = hidden_states - if self.normalize_before: - 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 @@ -419,15 +357,12 @@ class TFBartDecoderLayer(tf.keras.layers.Layer): ) hidden_states = self.dropout(hidden_states, training=training) hidden_states = residual + hidden_states - if not self.normalize_before: - hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states = self.self_attn_layer_norm(hidden_states) # Cross-Attention Block cross_attn_present_key_value = None if encoder_hidden_states is not None: residual = hidden_states - if self.normalize_before: - 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 @@ -439,24 +374,19 @@ class TFBartDecoderLayer(tf.keras.layers.Layer): ) hidden_states = self.dropout(hidden_states, training=training) hidden_states = residual + hidden_states - if not self.normalize_before: - hidden_states = self.encoder_attn_layer_norm(hidden_states) + hidden_states = self.encoder_attn_layer_norm(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 - if self.normalize_before: - 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 - - if not self.normalize_before: - hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.final_layer_norm(hidden_states) return ( hidden_states, @@ -472,8 +402,8 @@ class TFBartPretrainedModel(TFPreTrainedModel): @property def dummy_inputs(self): pad_token = 1 - input_ids = tf.cast(tf.constant(DUMMY_INPUTS), tf.int32) - decoder_input_ids = tf.cast(tf.constant(DUMMY_INPUTS), tf.int32) + input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) dummy_inputs = { "decoder_input_ids": decoder_input_ids, "attention_mask": tf.math.not_equal(input_ids, pad_token), @@ -520,14 +450,6 @@ class TFBartPretrainedModel(TFPreTrainedModel): return self.serving_output(output) -class TFPretrainedBartModel(TFBartPretrainedModel): - def __init_subclass__(self): - warnings.warn( - "The class `TFPretrainedBartModel` has been deprecated, please use `TFBartPretrainedModel` instead.", - FutureWarning, - ) - - BART_START_DOCSTRING = r""" This model inherits from :class:`~transformers.TFPreTrainedModel`. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input @@ -563,6 +485,36 @@ BART_START_DOCSTRING = r""" model weights. """ + +BART_GENERATION_EXAMPLE = r""" + Summarization example:: + + >>> from transformers import BartTokenizer, TFBartForConditionalGeneration, BartConfig + + >>> model = TFBartForConditionalGeneration.from_pretrained('facebook/bart-large') + >>> tokenizer = BartTokenizer.from_pretrained('facebook/bart-large') + + >>> ARTICLE_TO_SUMMARIZE = "My friends are cool but they eat too many carbs." + >>> inputs = tokenizer([ARTICLE_TO_SUMMARIZE], max_length=1024, return_tensors='tf') + + >>> # Generate Summary + >>> summary_ids = model.generate(inputs['input_ids'], num_beams=4, max_length=5, early_stopping=True) + >>> print([tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids]) + + Mask filling example:: + + >>> from transformers import BartTokenizer, TFBartForConditionalGeneration + >>> tokenizer = BartTokenizer.from_pretrained('facebook/bart-large') + >>> TXT = "My friends are but they eat too many carbs." + + >>> model = TFBartForConditionalGeneration.from_pretrained('facebook/bart-large') + >>> input_ids = tokenizer([TXT], return_tensors='tf')['input_ids'] + >>> logits = model(input_ids).logits + >>> probs = tf.nn.softmax(logits[0]) + >>> # probs[5] is associated with the mask token +""" + + BART_INPUTS_DOCSTRING = r""" Args: input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): @@ -581,8 +533,21 @@ BART_INPUTS_DOCSTRING = r""" `What are attention masks? <../glossary.html#attention-mask>`__ decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): - Provide for translation and summarization training. By default, the model will create this tensor by - shifting the input_ids right, following the paper. + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.BartTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + Bart uses the :obj:`eos_token_id` as the starting token for :obj:`decoder_input_ids` generation. If + :obj:`past_key_values` is used, optionally only the last :obj:`decoder_input_ids` have to be input (see + :obj:`past_key_values`). + + For translation and summarization training, :obj:`decoder_input_ids` should be provided. If no + :obj:`decoder_input_ids` is provided, the model will create this tensor by shifting the :obj:`input_ids` to + the right for denoising pre-training following the paper. decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. encoder_outputs (:obj:`tf.FloatTensor`, `optional`): @@ -603,7 +568,7 @@ BART_INPUTS_DOCSTRING = r""" Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for more detail. return_dict (:obj:`bool`, `optional`): - Whether or not to return a :class:`~transformers.file_utils.TFModelOutput` instead of a plain tuple. + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. training (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether or not to use the model in training mode (some modules like dropout modules have different behaviors between training and evaluation). @@ -626,36 +591,19 @@ class TFBartEncoder(tf.keras.layers.Layer): self.config = config self.dropout = tf.keras.layers.Dropout(config.dropout) self.layerdrop = config.encoder_layerdrop - self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 self.padding_idx = config.pad_token_id self.max_source_positions = config.max_position_embeddings + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 self.embed_tokens = embed_tokens - if config.static_position_embeddings: - self.embed_positions = TFBartSinusoidalPositionalEmbedding( - config.max_position_embeddings, - config.d_model, - name="embed_positions", - ) - else: - self.embed_positions = TFBartLearnedPositionalEmbedding( - config.max_position_embeddings, - config.d_model, - self.padding_idx, - config.extra_pos_embeddings, - name="embed_positions", - ) + self.embed_positions = TFBartLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) self.layers = [TFBartEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)] - self.layernorm_embedding = ( - tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") - if config.normalize_embedding - else tf.keras.layers.Layer() - ) - self.layer_norm = ( - tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm") - if config.add_final_layer_norm - else None - ) + self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") def set_embed_tokens(self, embed_tokens): self.embed_tokens = embed_tokens @@ -725,11 +673,7 @@ class TFBartEncoder(tf.keras.layers.Layer): raise ValueError("You have to specify either input_ids or inputs_embeds") if inputs["inputs_embeds"] is None: - inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) - else: - inputs["inputs_embeds"] = inputs["inputs_embeds"] - - inputs["inputs_embeds"] = inputs["inputs_embeds"] * self.embed_scale + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale embed_pos = self.embed_positions(input_shape) hidden_states = inputs["inputs_embeds"] + embed_pos @@ -739,7 +683,9 @@ class TFBartEncoder(tf.keras.layers.Layer): # 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"]) + attention_mask = _expand_mask(inputs["attention_mask"]) + else: + attention_mask = None encoder_states = () if inputs["output_hidden_states"] else None all_attentions = () if inputs["output_attentions"] else None @@ -754,12 +700,11 @@ class TFBartEncoder(tf.keras.layers.Layer): if inputs["training"] and (dropout_probability < self.layerdrop): # skip the layer continue - hidden_states, attn = encoder_layer(hidden_states, inputs["attention_mask"]) + hidden_states, attn = encoder_layer(hidden_states, attention_mask) if inputs["output_attentions"]: all_attentions += (attn,) - if self.layer_norm: - hidden_states = self.layer_norm(hidden_states) + if inputs["output_hidden_states"]: encoder_states = encoder_states + (hidden_states,) @@ -786,36 +731,18 @@ class TFBartDecoder(tf.keras.layers.Layer): self.config = config self.padding_idx = config.pad_token_id self.embed_tokens = embed_tokens - self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 self.layerdrop = config.decoder_layerdrop - if config.static_position_embeddings: - self.embed_positions = TFBartSinusoidalPositionalEmbedding( - config.max_position_embeddings, - config.d_model, - name="embed_positions", - ) - else: - self.embed_positions = TFBartLearnedPositionalEmbedding( - config.max_position_embeddings, - config.d_model, - self.padding_idx, - config.extra_pos_embeddings, - name="embed_positions", - ) + self.embed_positions = TFBartLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 self.layers = [TFBartDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)] - self.layernorm_embedding = ( - tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") - if config.normalize_embedding - else tf.keras.layers.Layer() - ) - self.layer_norm = ( - tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm") - if config.add_final_layer_norm - else None - ) + self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") self.dropout = tf.keras.layers.Dropout(config.dropout) - self.do_blenderbot_90_layernorm = config.do_blenderbot_90_layernorm def set_embed_tokens(self, embed_tokens): self.embed_tokens = embed_tokens @@ -912,16 +839,16 @@ class TFBartDecoder(tf.keras.layers.Layer): raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") past_key_values_length = ( - inputs["past_key_values"][0][0].shape[2] if inputs["past_key_values"] is not None else 0 + shape_list(inputs["past_key_values"][0][0])[2] if inputs["past_key_values"] is not None else 0 ) # embed positions positions = self.embed_positions(input_shape, past_key_values_length) if inputs["inputs_embeds"] is None: - inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale - hidden_states = inputs["inputs_embeds"] * self.embed_scale + hidden_states = inputs["inputs_embeds"] # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] if input_shape[-1] > 1: @@ -931,35 +858,16 @@ class TFBartDecoder(tf.keras.layers.Layer): tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] ) - if inputs["attention_mask"] is None and inputs["input_ids"] is not None and input_shape[-1] > 1: - inputs["attention_mask"] = tf.cast( - tf.math.not_equal(inputs["input_ids"], self.config.pad_token_id), inputs["input_ids"].dtype + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask( + inputs["attention_mask"], tgt_len=input_shape[-1] ) - inputs["attention_mask"] = tf.concat( - [ - tf.ones((input_shape[0], past_key_values_length), dtype=inputs["attention_mask"].dtype), - inputs["attention_mask"], - ], - axis=-1, - ) - else: - inputs["attention_mask"] = tf.ones( - (input_shape[0], input_shape[1] + past_key_values_length), dtype=tf.int32 - ) - - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - combined_attention_mask = combined_attention_mask + _expand_mask( - inputs["attention_mask"], tgt_len=input_shape[-1] - ) 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]) - if self.do_blenderbot_90_layernorm: - hidden_states = self.layernorm_embedding(hidden_states) + positions - else: - hidden_states = self.layernorm_embedding(hidden_states + positions) + hidden_states = self.layernorm_embedding(hidden_states + positions) hidden_states = self.dropout(hidden_states, training=inputs["training"]) # decoder layers @@ -991,10 +899,6 @@ class TFBartDecoder(tf.keras.layers.Layer): if inputs["output_attentions"]: all_self_attns += (layer_self_attn,) - if self.layer_norm is not None: # same as if config.add_final_layer_norm - hidden_states = self.layer_norm(hidden_states) - - # Convert to standard output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim) if inputs["output_hidden_states"]: all_hidden_states += (hidden_states,) else: @@ -1002,7 +906,7 @@ class TFBartDecoder(tf.keras.layers.Layer): all_self_attns = list(all_self_attns) if inputs["output_attentions"] else None - present_key_values = (inputs["encoder_hidden_states"], present_key_values) if inputs["use_cache"] else None + present_key_values = (encoder_hidden_states, present_key_values) if inputs["use_cache"] else None if not inputs["return_dict"]: return hidden_states, present_key_values, all_hidden_states, all_self_attns @@ -1098,7 +1002,7 @@ class TFBartModel(TFBartPretrainedModel): if inputs["decoder_input_ids"] is None and inputs["input_ids"] is not None: inputs["decoder_input_ids"] = shift_tokens_right( - inputs["input_ids"], self.config.pad_token_id, self.config.eos_token_id + inputs["input_ids"], self.config.pad_token_id, self.config.decoder_start_token_id ) if inputs["encoder_outputs"] is None: @@ -1206,6 +1110,7 @@ class TFBartForConditionalGeneration(TFBartPretrainedModel): @add_start_docstrings_to_model_forward(BART_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + @add_end_docstrings(BART_GENERATION_EXAMPLE) def call( self, input_ids=None, @@ -1224,22 +1129,14 @@ class TFBartForConditionalGeneration(TFBartPretrainedModel): training=False, **kwargs, ): - """ + r""" + labels (:obj:`tf.Tensor` of shape :obj:`(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: - Examples:: - - # Mask filling only works for bart-large - from transformers import BartTokenizer, TFBartForConditionalGeneration - import tensorflow as tf - mname = 'facebook/bart-large' - tokenizer = BartTokenizer.from_pretrained(mname) - TXT = "My friends are but they eat too many carbs." - model = TFBartForConditionalGeneration.from_pretrained(mname) - batch = tokenizer([TXT], return_tensors='tf') - logits = model(inputs=batch.input_ids).logits - probs = tf.nn.softmax(logits[0]) - # probs[5] is associated with the mask token """ inputs = input_processing( func=self.call, @@ -1265,7 +1162,7 @@ class TFBartForConditionalGeneration(TFBartPretrainedModel): inputs["use_cache"] = False if inputs["decoder_input_ids"] is None: inputs["decoder_input_ids"] = shift_tokens_right( - inputs["labels"], self.config.pad_token_id, self.config.eos_token_id + inputs["labels"], self.config.pad_token_id, self.config.decoder_start_token_id ) outputs = self.model( @@ -1363,7 +1260,8 @@ class TFBartForConditionalGeneration(TFBartPretrainedModel): 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), + 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/blenderbot/__init__.py b/src/transformers/models/blenderbot/__init__.py index 1c2565b404..b418d3e9b3 100644 --- a/src/transformers/models/blenderbot/__init__.py +++ b/src/transformers/models/blenderbot/__init__.py @@ -36,7 +36,7 @@ if is_torch_available(): if is_tf_available(): - _import_structure["modeling_tf_blenderbot"] = ["TFBlenderbotForConditionalGeneration"] + _import_structure["modeling_tf_blenderbot"] = ["TFBlenderbotForConditionalGeneration", "TFBlenderbotModel"] if TYPE_CHECKING: @@ -52,7 +52,7 @@ if TYPE_CHECKING: ) if is_tf_available(): - from .modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration + from .modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration, TFBlenderbotModel else: import importlib diff --git a/src/transformers/models/blenderbot/configuration_blenderbot.py b/src/transformers/models/blenderbot/configuration_blenderbot.py index 1b48f8952b..242d5d36f5 100644 --- a/src/transformers/models/blenderbot/configuration_blenderbot.py +++ b/src/transformers/models/blenderbot/configuration_blenderbot.py @@ -161,17 +161,6 @@ class BlenderbotConfig(PretrainedConfig): self.gradient_checkpointing = gradient_checkpointing self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True - # IMPORTANT - # DELETE ALL OF THE FOLLOWING LINES AS SOON AS TF IS READY - self.extra_pos_embeddings = 0 - self.normalize_before = True - self.add_final_layer_norm = True - self.do_blenderbot_90_layernorm = True - self.normalize_embedding = False - self.static_position_embeddings = False - self.add_bias_logits = False - self.force_bos_token_to_be_generated = False - @property def num_attention_heads(self) -> int: return self.encoder_attention_heads diff --git a/src/transformers/models/blenderbot/modeling_tf_blenderbot.py b/src/transformers/models/blenderbot/modeling_tf_blenderbot.py index ba51e87a1c..a7e0de3efd 100644 --- a/src/transformers/models/blenderbot/modeling_tf_blenderbot.py +++ b/src/transformers/models/blenderbot/modeling_tf_blenderbot.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team. +# Copyright 2021 The Facebook, Inc 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. @@ -12,35 +12,1301 @@ # 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. -"""TF BlenderBot model, ported from the fairseq repo.""" +""" TF 2.0 Blenderbot model. """ + + +import os +import random +import warnings +from typing import Dict, Optional, Tuple, Union import tensorflow as tf -from ...file_utils import add_start_docstrings +from ...activations_tf import get_tf_activation +from ...file_utils import ( + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ...modeling_tf_outputs import ( + TFBaseModelOutput, + TFBaseModelOutputWithPast, + TFSeq2SeqLMOutput, + TFSeq2SeqModelOutput, +) + +# Public API +from ...modeling_tf_utils import ( + DUMMY_INPUTS, + TFPreTrainedModel, + TFSharedEmbeddings, + TFWrappedEmbeddings, + input_processing, + keras_serializable, + shape_list, +) from ...utils import logging -from ..bart.modeling_tf_bart import BART_START_DOCSTRING, LARGE_NEGATIVE, TFBartForConditionalGeneration +from ..blenderbot_small import TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel from .configuration_blenderbot import BlenderbotConfig -_CONFIG_FOR_DOC = "BlenderbotConfig" - -START_DOCSTRING = BART_START_DOCSTRING.replace( - "inherits from :class:`~transformers.TFPreTrainedModel`", - "inherits from :class:`~transformers.TFBartForConditionalGeneration`", -).replace("BartConfig", _CONFIG_FOR_DOC) - - logger = logging.get_logger(__name__) +_CONFIG_FOR_DOC = "BlenderbotConfig" +_TOKENIZER_FOR_DOC = "BlenderbotTokenizer" -@add_start_docstrings("Blenderbot model for open domain dialogue", START_DOCSTRING) -class TFBlenderbotForConditionalGeneration(TFBartForConditionalGeneration): + +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): + shifted_input_ids = tf.cast(input_ids, tf.int32) + shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1) + start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), decoder_start_token_id) + shifted_input_ids = tf.concat([start_tokens, shifted_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 + ) + + # "Verify that `labels` has only positive values and -100" + assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, tf.int32)) + + # 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), dtype=tf.float32) * 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) + mask = tf.cast(mask, tf.float32) + + if past_key_values_length > 0: + mask = tf.concat([tf.zeros((tgt_len, past_key_values_length), dtype=tf.float32), mask], axis=-1) + return tf.broadcast_to(mask[None, None, :, :], (bsz, 1, tgt_len, tgt_len + past_key_values_length)) + + +# 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]`. + """ + bsz, src_len = shape_list(mask) + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = tf.cast(tf.broadcast_to(mask[:, None, None, :], (bsz, 1, tgt_len, src_len)), tf.float32) + + return (1.0 - expanded_mask) * LARGE_NEGATIVE + + +class TFBlenderbotLearnedPositionalEmbedding(TFSharedEmbeddings): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, **kwargs): + assert padding_idx is not None, "padding_idx cannot be None" + super().__init__(num_embeddings, embedding_dim, **kwargs) + + def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): + """Input is expected to be of size [bsz x seqlen].""" + bsz, seq_len = input_shape[:2] + + positions = tf.range( + past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" + ) + return super().call(positions) + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartAttention with Bart->Blenderbot +class TFBlenderbotAttention(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 + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by 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, + 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) + + 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: + 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)}", + ) + 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) + + attn_probs = self.dropout(attn_weights, training=training) + + attn_output = tf.matmul(attn_probs, value_states) + + 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 + + +# Copied from transformers.models.mbart.modeling_tf_mbart.TFMBartEncoderLayer with MBart->Blenderbot +class TFBlenderbotEncoderLayer(tf.keras.layers.Layer): + def __init__(self, config: BlenderbotConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFBlenderbotAttention( + 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, training=False): + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + """ + 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 + ) + 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 + + +# Copied from transformers.models.mbart.modeling_tf_mbart.TFMBartDecoderLayer with MBart->Blenderbot +class TFBlenderbotDecoderLayer(tf.keras.layers.Layer): + def __init__(self, config: BlenderbotConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFBlenderbotAttention( + 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 = TFBlenderbotAttention( + 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, + past_key_value: Optional[Tuple[tf.Tensor]] = None, + training=False, + ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`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 (:obj:`tf.Tensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (:obj:`tf.Tensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + past_key_value (:obj:`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, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + # Cross-Attention Block + cross_attn_present_key_value = 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_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_value=cross_attn_past_key_value, + ) + 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, + present_key_value, + ) + + +class TFBlenderbotPreTrainedModel(TFPreTrainedModel): config_class = BlenderbotConfig + base_model_prefix = "model" + + @property + def dummy_inputs(self): + pad_token = 1 + input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + dummy_inputs = { + "decoder_input_ids": decoder_input_ids, + "attention_mask": tf.math.not_equal(input_ids, pad_token), + "input_ids": input_ids, + } + return dummy_inputs + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.get_input_embeddings + def get_input_embeddings(self): + base_model = getattr(self, self.base_model_prefix, self) + + return base_model.shared + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.set_input_embeddings + def set_input_embeddings(self, value): + base_model = getattr(self, self.base_model_prefix, self) + + try: + base_model.shared.weight = value + except AttributeError: + self(self.dummy_inputs) + base_model.shared.weight = value + + base_model.shared.vocab_size = shape_list(base_model.shared.weight)[0] + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + embed_tokens = TFWrappedEmbeddings(base_model.shared, abs_scope_name=shared_abs_scope_name) + base_model.encoder.set_embed_tokens(embed_tokens) + base_model.decoder.set_embed_tokens(embed_tokens) + + @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"), + } + ] + ) + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.serving + def serving(self, inputs): + output = self.call(inputs) + + return self.serving_output(output) + + +BLENDERBOT_START_DOCSTRING = r""" + This model inherits from :class:`~transformers.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 `__ 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. + + .. note:: + + 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 :meth:`tf.keras.Model.fit` method which currently requires having all + the tensors in the first argument of the model call function: :obj:`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 :obj:`input_ids` only and nothing else: :obj:`model(input_ids)` + - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: + :obj:`model([input_ids, attention_mask])` or :obj:`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: + :obj:`model({"input_ids": input_ids, "token_type_ids": token_type_ids})` + + Args: + config (:class:`~transformers.BlenderbotConfig`): 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 :meth:`~transformers.TFPreTrainedModel.from_pretrained` method to load the + model weights. +""" + +BLENDERBOT_GENERATION_EXAMPLE = r""" + Conversation example:: + + >>> from transformers import BlenderbotTokenizer, TFBlenderbotForConditionalGeneration + >>> mname = 'facebook/blenderbot-400M-distill' + >>> model = TFBlenderbotForConditionalGeneration.from_pretrained(mname) + >>> tokenizer = BlenderbotTokenizer.from_pretrained(mname) + >>> UTTERANCE = "My friends are cool but they eat too many carbs." + >>> print("Human: ", UTTERANCE) + >>> inputs = tokenizer([UTTERANCE], return_tensors='tf') + >>> reply_ids = model.generate(**inputs) + >>> print("Bot: ", tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]) + + >>> REPLY = "I'm not sure" + >>> print("Human: ", REPLY) + >>> NEXT_UTTERANCE = ( + ... "My friends are cool but they eat too many carbs. That's unfortunate. " + ... "Are they trying to lose weight or are they just trying to be healthier? " + ... " I'm not sure." + ... ) + >>> inputs = tokenizer([NEXT_UTTERANCE], return_tensors='tf') + >>> next_reply_ids = model.generate(**inputs) + >>> print("Bot: ", tokenizer.batch_decode(next_reply_ids, skip_special_tokens=True)[0]) +""" + +BLENDERBOT_INPUTS_DOCSTRING = r""" + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.BlenderbotTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`({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.html#attention-mask>`__ + decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.BlenderbotTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + Blenderbot uses the :obj:`bos_token_id` as the starting token for :obj:`decoder_input_ids` generation. If + :obj:`past_key_values` is used, optionally only the last :obj:`decoder_input_ids` have to be input (see + :obj:`past_key_values`). + decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. + encoder_outputs (:obj:`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 :obj:`(batch_size, sequence_length, hidden_size)` is a sequence of + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers`) + contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). Set to :obj:`False` during training, :obj:`True` during generation + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for + more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. + training (:obj:`bool`, `optional`, defaults to :obj:`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 TFBlenderbotEncoder(tf.keras.layers.Layer): + config_class = BlenderbotConfig + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + :class:`TFBlenderbotEncoderLayer`. + + Args: + config: BlenderbotConfig + """ + + def __init__(self, config: BlenderbotConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.layerdrop = config.encoder_layerdrop + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_position_embeddings + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + + self.embed_tokens = embed_tokens + self.embed_positions = TFBlenderbotLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.layers = [TFBlenderbotEncoderLayer(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 set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + """ + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.BlenderbotTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + inputs_embeds=inputs_embeds, + 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 input_ids and 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 input_ids or inputs_embeds") + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + embed_pos = self.embed_positions(input_shape) + hidden_states = inputs["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] + attention_mask = _expand_mask(inputs["attention_mask"]) + else: + attention_mask = None + + encoder_states = () if inputs["output_hidden_states"] else None + all_attentions = () if inputs["output_attentions"] else None + + # encoder layers + for encoder_layer in 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, attention_mask) + + 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 TFBlenderbotDecoder(tf.keras.layers.Layer): + config_class = BlenderbotConfig + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a :class:`TFBlenderbotDecoderLayer` + + Args: + config: BlenderbotConfig + embed_tokens: output embedding + """ + + def __init__(self, config: BlenderbotConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.padding_idx = config.pad_token_id + self.embed_tokens = embed_tokens + self.layerdrop = config.decoder_layerdrop + self.embed_positions = TFBlenderbotLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + self.layers = [TFBlenderbotDecoderLayer(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 set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_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 (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.BlenderbotTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + encoder_hidden_states (:obj:`tf.Tensor` of shape :obj:`(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 (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape :obj:`(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 :obj:`past_key_values` are used, the user can optionally input only the last + :obj:`decoder_input_ids` (those that don't have their past key value states given to this model) of + shape :obj:`(batch_size, 1)` instead of all :obj:`decoder_input_ids`` of shape :obj:`(batch_size, + sequence_length)`. + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + 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 = ( + shape_list(inputs["past_key_values"][0][0])[2] if inputs["past_key_values"] is not None else 0 + ) + + # embed positions + positions = self.embed_positions(input_shape, past_key_values_length) + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + hidden_states = inputs["inputs_embeds"] + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length) + else: + combined_attention_mask = _expand_mask( + tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] + ) + + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask( + inputs["attention_mask"], tgt_len=input_shape[-1] + ) + + 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]) + + hidden_states = hidden_states + positions + hidden_states = self.dropout(hidden_states, training=inputs["training"]) + + # decoder layers + all_hidden_states = () + all_self_attns = () + present_key_values = () + 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 + + hidden_states, layer_self_attn, present_key_value = decoder_layer( + hidden_states, + attention_mask=combined_attention_mask, + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=inputs["encoder_attention_mask"], + past_key_value=past_key_value, + ) + + if inputs["use_cache"]: + present_key_values += (present_key_value,) + + if inputs["output_attentions"]: + all_self_attns += (layer_self_attn,) + + hidden_states = self.layer_norm(hidden_states) + + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + else: + all_hidden_states = None + + all_self_attns = list(all_self_attns) if inputs["output_attentions"] else None + + present_key_values = (encoder_hidden_states, present_key_values) if inputs["use_cache"] else None + + if not inputs["return_dict"]: + return hidden_states, present_key_values, all_hidden_states, all_self_attns + else: + return TFBaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=present_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +@add_start_docstrings( + "The bare BLENDERBOT Model outputting raw hidden-states without any specific head on top.", + BLENDERBOT_START_DOCSTRING, +) +@keras_serializable +class TFBlenderbotModel(TFBlenderbotPreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config: BlenderbotConfig, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.shared = TFSharedEmbeddings(config.vocab_size, config.d_model, config.pad_token_id, name="model.shared") + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + # Wraps layer to avoid problems with weight restoring and ensuring we're in the correct TF scope. + embed_tokens = TFWrappedEmbeddings(self.shared, abs_scope_name=shared_abs_scope_name) + embed_tokens.vocab_size = self.shared.vocab_size + embed_tokens.hidden_size = self.shared.hidden_size + + self.encoder = TFBlenderbotEncoder(config, embed_tokens, name="encoder") + self.decoder = TFBlenderbotDecoder(config, embed_tokens, name="decoder") + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs): + if pretrained_model_name_or_path == "facebook/blenderbot-90M": + warnings.warn( + "The checkpoint `facebook/blenderbot-90M` is deprecated. In the future, please use the identical checkpoint `facebook/small_blenderbot-90M` with `TFBlenderbotSmallModel.from_pretrained('facebook/small_blenderbot-90M')` instead.", + FutureWarning, + ) + return TFBlenderbotSmallModel.from_pretrained(pretrained_model_name_or_path) + + return super(TFBlenderbotModel, cls).from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(BLENDERBOT_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + @replace_return_docstrings(output_type=TFSeq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + r""" + Returns: + + Example:: + + >>> from transformers import BlenderbotTokenizer, TFBlenderbotModel + + >>> model = TFBlenderbotModel.from_pretrained("facebook/blenderbot-400M-distill") + >>> tokenizer = BlenderbotTokenizer.from_pretrained("facebook/blenderbot-400M-distill") + + >>> input_ids = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="tf").input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="tf").input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + + >>> last_hidden_states = outputs.last_hidden_state + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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, + ) + + inputs["output_hidden_states"] = ( + inputs["output_hidden_states"] + if inputs["output_hidden_states"] is not None + else self.config.output_hidden_states + ) + + if inputs["encoder_outputs"] is None: + inputs["encoder_outputs"] = self.encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + inputs_embeds=inputs["inputs_embeds"], + 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() + + decoder_outputs = self.decoder( + inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=inputs["encoder_outputs"][0], + encoder_attention_mask=inputs["attention_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, + encoder_last_hidden_state=inputs["encoder_outputs"].last_hidden_state, + encoder_hidden_states=inputs["encoder_outputs"].hidden_states, + encoder_attentions=inputs["encoder_outputs"].attentions, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartModel.serving_output + 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 + 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, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + +@add_start_docstrings( + "The BLENDERBOT Model with a language modeling head. Can be used for summarization.", + BLENDERBOT_START_DOCSTRING, +) +class TFBlenderbotForConditionalGeneration(TFBlenderbotPreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + r"model.encoder.embed_tokens.weight", + r"model.decoder.embed_tokens.weight", + ] + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.model = TFBlenderbotModel(config, name="model") + self.use_cache = config.use_cache + # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency. + self.final_logits_bias = self.add_weight( + name="final_logits_bias", shape=[1, config.vocab_size], initializer="zeros", trainable=False + ) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs): + if pretrained_model_name_or_path == "facebook/blenderbot-90M": + warnings.warn( + "The checkpoint `facebook/blenderbot-90M` is deprecated. In the future, please use the identical checkpoint `facebook/small_blenderbot-90M` with `TFBlenderbotSmallForConditionalGeneration.from_pretrained('facebook/small_blenderbot-90M')` instead.", + FutureWarning, + ) + return TFBlenderbotSmallForConditionalGeneration.from_pretrained(pretrained_model_name_or_path) + + return super(TFBlenderbotForConditionalGeneration, cls).from_pretrained( + pretrained_model_name_or_path, *model_args, **kwargs + ) + + @add_start_docstrings_to_model_forward(BLENDERBOT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + @add_end_docstrings(BLENDERBOT_GENERATION_EXAMPLE) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[TFBaseModelOutput] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + labels=None, + training=False, + **kwargs, + ): + r""" + labels (:obj:`torch.LongTensor` of shape :obj:`(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: + + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + labels=labels, + training=training, + kwargs_call=kwargs, + ) + + 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( + inputs["input_ids"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + encoder_outputs=inputs["encoder_outputs"], + decoder_attention_mask=inputs["decoder_attention_mask"], + past_key_values=inputs["past_key_values"], + inputs_embeds=inputs["inputs_embeds"], + 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.model.shared(outputs[0], mode="linear") + lm_logits = lm_logits + self.final_logits_bias + masked_lm_loss = None if inputs["labels"] is None else self.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, # index 1 of d outputs + decoder_hidden_states=outputs.decoder_hidden_states, # index 2 of d outputs + decoder_attentions=outputs.decoder_attentions, # index 3 of d outputs + encoder_last_hidden_state=outputs.last_hidden_state, # index 0 of encoder outputs + encoder_hidden_states=outputs.encoder_hidden_states, # 1 of e out + encoder_attentions=outputs.encoder_attentions, # 2 of e out + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.serving_output + 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 + 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 TFSeq2SeqLMOutput( + logits=output.logits, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.prepare_inputs_for_generation + def prepare_inputs_for_generation(self, decoder_input_ids, past, attention_mask, use_cache, **kwargs) -> Dict: + assert past is not None and len(past) in {1, 2}, f"past has to be an iterable of length 1,2 got {past}" + if len(past) == 1: + assert isinstance(past[0], tf.Tensor), f"`past[0]` has to be of type `tf.Tensor`, but is {type(past[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=past[0]) + past_key_values = None + else: + assert ( + len(past) == 2 + ), "`past` has to be of length 2 with the encoder_outputs at the first position and past_key_values at the second position." + encoder_outputs, past_key_values = past + if isinstance(encoder_outputs, tuple): + assert isinstance( + encoder_outputs[0], tf.Tensor + ), f"`encoder_outputs[0]` has to be of type `tf.Tensor`, but is {type(encoder_outputs[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs[0]) + elif isinstance(encoder_outputs, tf.Tensor): + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs) + assert ( + past_key_values + ), f"decoder cached states must be truthy. got {past_key_values} from the 2nd element of past" + decoder_input_ids = decoder_input_ids[:, -1:] + + assert isinstance( + encoder_outputs, TFBaseModelOutput + ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}." + return { + "input_ids": None, # encoder_outputs is defined. input_ids not needed + "encoder_outputs": encoder_outputs, + "past_key_values": past_key_values, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "use_cache": use_cache, # change this to avoid caching (presumably for debugging) + } + + @staticmethod + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration._reorder_cache + 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) def adjust_logits_during_generation(self, logits, cur_len, max_length): - """Never predict pad_token_id. Predict when max_length is reached.""" - vocab_range = tf.constant(range(self.config.vocab_size)) - logits = tf.where(vocab_range == self.config.pad_token_id, LARGE_NEGATIVE, logits) if cur_len == max_length - 1: - logits = tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits) - return logits + vocab_range = tf.constant(range(self.config.vocab_size)) + return tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits) + else: + return logits + + def get_encoder(self): + return self.model.encoder + + def get_output_embeddings(self): + return self.get_input_embeddings() + + def set_output_embeddings(self, value): + self.set_input_embeddings(value) + + def get_bias(self): + return {"final_logits_bias": self.final_logits_bias} + + def set_bias(self, value): + self.final_logits_bias = value["final_logits_bias"] + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.compute_loss + def compute_loss(self, labels, logits): + """CrossEntropyLoss that ignores pad tokens""" + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True, + reduction=tf.keras.losses.Reduction.NONE, + ) + melted_labels = tf.reshape(labels, (-1,)) + active_loss = tf.not_equal(melted_labels, self.config.pad_token_id) + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(melted_labels, active_loss) + return loss_fn(labels, reduced_logits) diff --git a/src/transformers/models/blenderbot_small/__init__.py b/src/transformers/models/blenderbot_small/__init__.py index 136c17b339..1808979db9 100644 --- a/src/transformers/models/blenderbot_small/__init__.py +++ b/src/transformers/models/blenderbot_small/__init__.py @@ -17,7 +17,7 @@ # limitations under the License. from typing import TYPE_CHECKING -from ...file_utils import _BaseLazyModule, is_torch_available +from ...file_utils import _BaseLazyModule, is_tf_available, is_torch_available _import_structure = { @@ -33,6 +33,11 @@ if is_torch_available(): "BlenderbotSmallPreTrainedModel", ] +if is_tf_available(): + _import_structure["modeling_tf_blenderbot_small"] = [ + "TFBlenderbotSmallForConditionalGeneration", + "TFBlenderbotSmallModel", + ] if TYPE_CHECKING: from .configuration_blenderbot_small import BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP, BlenderbotSmallConfig @@ -46,6 +51,9 @@ if TYPE_CHECKING: BlenderbotSmallPreTrainedModel, ) + if is_tf_available(): + from .modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration, TFBlenderbotModel + else: import importlib import os diff --git a/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py b/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py index c4af450bbe..36cff4c417 100755 --- a/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py +++ b/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py @@ -866,6 +866,7 @@ class BlenderbotSmallDecoder(BlenderbotSmallPreTrainedModel): all_self_attns = () if output_attentions else None all_cross_attentions = () if output_attentions else None next_decoder_cache = () if use_cache else None + for idx, decoder_layer in enumerate(self.layers): # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) if output_hidden_states: diff --git a/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py b/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py new file mode 100644 index 0000000000..aa8a15f174 --- /dev/null +++ b/src/transformers/models/blenderbot_small/modeling_tf_blenderbot_small.py @@ -0,0 +1,1295 @@ +# coding=utf-8 +# Copyright 2021 The Facebook, Inc 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. +""" TF 2.0 BlenderbotSmall model. """ + + +import random +from typing import Dict, Optional, Tuple, Union + +import tensorflow as tf + +from ...activations_tf import get_tf_activation +from ...file_utils import ( + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ...modeling_tf_outputs import ( + TFBaseModelOutput, + TFBaseModelOutputWithPast, + TFSeq2SeqLMOutput, + TFSeq2SeqModelOutput, +) + +# Public API +from ...modeling_tf_utils import ( + DUMMY_INPUTS, + TFPreTrainedModel, + TFSharedEmbeddings, + TFWrappedEmbeddings, + input_processing, + keras_serializable, + shape_list, +) +from ...utils import logging +from .configuration_blenderbot_small import BlenderbotSmallConfig + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "BlenderbotSmallConfig" +_TOKENIZER_FOR_DOC = "BlenderbotSmallTokenizer" + + +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): + shifted_input_ids = tf.cast(input_ids, tf.int32) + shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1) + start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), decoder_start_token_id) + shifted_input_ids = tf.concat([start_tokens, shifted_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 + ) + + # "Verify that `labels` has only positive values and -100" + assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, tf.int32)) + + # 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), dtype=tf.float32) * 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) + mask = tf.cast(mask, tf.float32) + + if past_key_values_length > 0: + mask = tf.concat([tf.zeros((tgt_len, past_key_values_length), dtype=tf.float32), mask], axis=-1) + return tf.broadcast_to(mask[None, None, :, :], (bsz, 1, tgt_len, tgt_len + past_key_values_length)) + + +# 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]`. + """ + bsz, src_len = shape_list(mask) + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = tf.cast(tf.broadcast_to(mask[:, None, None, :], (bsz, 1, tgt_len, src_len)), tf.float32) + + return (1.0 - expanded_mask) * LARGE_NEGATIVE + + +# Copied from transformers.models.blenderbot.modeling_tf_blenderbot.TFBlenderbotLearnedPositionalEmbedding with Blenderbot->BlenderbotSmall +class TFBlenderbotSmallLearnedPositionalEmbedding(TFSharedEmbeddings): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, **kwargs): + assert padding_idx is not None, "padding_idx cannot be None" + super().__init__(num_embeddings, embedding_dim, **kwargs) + + def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): + """Input is expected to be of size [bsz x seqlen].""" + bsz, seq_len = input_shape[:2] + + positions = tf.range( + past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" + ) + return super().call(positions) + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartAttention with Bart->BlenderbotSmall +class TFBlenderbotSmallAttention(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 + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by 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, + 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) + + 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: + 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)}", + ) + 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) + + attn_probs = self.dropout(attn_weights, training=training) + + attn_output = tf.matmul(attn_probs, value_states) + + 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 + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartEncoderLayer with Bart->BlenderbotSmall +class TFBlenderbotSmallEncoderLayer(tf.keras.layers.Layer): + def __init__(self, config: BlenderbotSmallConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFBlenderbotSmallAttention( + 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, training=False): + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + """ + residual = hidden_states + hidden_states, self_attn_weights, _ = self.self_attn( + hidden_states=hidden_states, attention_mask=attention_mask + ) + 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 + hidden_states = self.self_attn_layer_norm(hidden_states) + + residual = 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 + hidden_states = self.final_layer_norm(hidden_states) + + return hidden_states, self_attn_weights + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartDecoderLayer with Bart->BlenderbotSmall +class TFBlenderbotSmallDecoderLayer(tf.keras.layers.Layer): + def __init__(self, config: BlenderbotSmallConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFBlenderbotSmallAttention( + 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 = TFBlenderbotSmallAttention( + 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, + past_key_value: Optional[Tuple[tf.Tensor]] = None, + training=False, + ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`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 (:obj:`tf.Tensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (:obj:`tf.Tensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + past_key_value (:obj:`Tuple(tf.Tensor)`): cached past key and value projection states + """ + residual = 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, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Cross-Attention Block + cross_attn_present_key_value = None + if encoder_hidden_states is not None: + residual = 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_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_value=cross_attn_past_key_value, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + hidden_states = self.encoder_attn_layer_norm(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.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 + hidden_states = self.final_layer_norm(hidden_states) + + return ( + hidden_states, + self_attn_weights, + present_key_value, + ) + + +class TFBlenderbotSmallPreTrainedModel(TFPreTrainedModel): + config_class = BlenderbotSmallConfig + base_model_prefix = "model" + + @property + def dummy_inputs(self): + pad_token = 1 + input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + dummy_inputs = { + "decoder_input_ids": decoder_input_ids, + "attention_mask": tf.math.not_equal(input_ids, pad_token), + "input_ids": input_ids, + } + return dummy_inputs + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.get_input_embeddings + def get_input_embeddings(self): + base_model = getattr(self, self.base_model_prefix, self) + + return base_model.shared + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.set_input_embeddings + def set_input_embeddings(self, value): + base_model = getattr(self, self.base_model_prefix, self) + + try: + base_model.shared.weight = value + except AttributeError: + self(self.dummy_inputs) + base_model.shared.weight = value + + base_model.shared.vocab_size = shape_list(base_model.shared.weight)[0] + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + embed_tokens = TFWrappedEmbeddings(base_model.shared, abs_scope_name=shared_abs_scope_name) + base_model.encoder.set_embed_tokens(embed_tokens) + base_model.decoder.set_embed_tokens(embed_tokens) + + @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"), + } + ] + ) + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.serving + def serving(self, inputs): + output = self.call(inputs) + + return self.serving_output(output) + + +BLENDERBOT_SMALL_START_DOCSTRING = r""" + This model inherits from :class:`~transformers.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 `__ 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. + + .. note:: + + 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 :meth:`tf.keras.Model.fit` method which currently requires having all + the tensors in the first argument of the model call function: :obj:`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 :obj:`input_ids` only and nothing else: :obj:`model(input_ids)` + - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: + :obj:`model([input_ids, attention_mask])` or :obj:`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: + :obj:`model({"input_ids": input_ids, "token_type_ids": token_type_ids})` + + Args: + config (:class:`~transformers.BlenderbotSmallConfig`): 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 :meth:`~transformers.TFPreTrainedModel.from_pretrained` method to load the + model weights. +""" + +BLENDERBOT_SMALL_GENERATION_EXAMPLE = r""" + Conversation example:: + + >>> from transformers import BlenderbotSmallTokenizer, TFBlenderbotSmallForConditionalGeneration + >>> mname = 'facebook/blenderbot_small-90M' + >>> model = BlenderbotSmallForConditionalGeneration.from_pretrained(mname) + >>> tokenizer = TFBlenderbotSmallTokenizer.from_pretrained(mname) + + >>> UTTERANCE = "My friends are cool but they eat too many carbs." + >>> print("Human: ", UTTERANCE) + >>> inputs = tokenizer([UTTERANCE], return_tensors='tf') + >>> inputs.pop("token_type_ids") + + >>> reply_ids = model.generate(**inputs) + >>> print("Bot: ", tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]) + what kind of carbs do they eat? i don't know much about carbs. + + >>> REPLY = "I'm not sure" + >>> print("Human: ", REPLY) + >>> NEXT_UTTERANCE = ( + ... "My friends are cool but they eat too many carbs. " + ... "what kind of carbs do they eat? i don't know much about carbs. " + ... "I'm not sure." + ... ) + + >>> inputs = tokenizer([NEXT_UTTERANCE], return_tensors='tf') + >>> inputs.pop("token_type_ids") + >>> next_reply_ids = model.generate(**inputs) + >>> print("Bot: ", tokenizer.batch_decode(next_reply_ids, skip_special_tokens=True)[0]) +""" + +BLENDERBOT_SMALL_INPUTS_DOCSTRING = r""" + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.BlenderbotSmallTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`({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.html#attention-mask>`__ + decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.BlenderbotSmallTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + BlenderbotSmall uses the :obj:`bos_token_id` as the starting token for :obj:`decoder_input_ids` generation. + If :obj:`past_key_values` is used, optionally only the last :obj:`decoder_input_ids` have to be input (see + :obj:`past_key_values`). + decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. + encoder_outputs (:obj:`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 :obj:`(batch_size, sequence_length, hidden_size)` is a sequence of + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers`) + contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). Set to :obj:`False` during training, :obj:`True` during generation + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for + more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. + training (:obj:`bool`, `optional`, defaults to :obj:`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 TFBlenderbotSmallEncoder(tf.keras.layers.Layer): + config_class = BlenderbotSmallConfig + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + :class:`TFBlenderbotSmallEncoderLayer`. + + Args: + config: BlenderbotSmallConfig + """ + + def __init__(self, config: BlenderbotSmallConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.layerdrop = config.encoder_layerdrop + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_position_embeddings + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + + self.embed_tokens = embed_tokens + self.embed_positions = TFBlenderbotSmallLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.layers = [TFBlenderbotSmallEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)] + self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + """ + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.BlenderbotSmallTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + inputs_embeds=inputs_embeds, + 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 input_ids and 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 input_ids or inputs_embeds") + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + embed_pos = self.embed_positions(input_shape) + hidden_states = inputs["inputs_embeds"] + embed_pos + hidden_states = self.layernorm_embedding(hidden_states) + 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] + attention_mask = _expand_mask(inputs["attention_mask"]) + else: + attention_mask = None + + encoder_states = () if inputs["output_hidden_states"] else None + all_attentions = () if inputs["output_attentions"] else None + + # encoder layers + for encoder_layer in 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, attention_mask) + + if inputs["output_attentions"]: + all_attentions += (attn,) + + 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 TFBlenderbotSmallDecoder(tf.keras.layers.Layer): + config_class = BlenderbotSmallConfig + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a + :class:`TFBlenderbotSmallDecoderLayer` + + Args: + config: BlenderbotSmallConfig + embed_tokens: output embedding + """ + + def __init__(self, config: BlenderbotSmallConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.padding_idx = config.pad_token_id + self.embed_tokens = embed_tokens + self.layerdrop = config.decoder_layerdrop + self.embed_positions = TFBlenderbotSmallLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + self.layers = [TFBlenderbotSmallDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)] + self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") + + self.dropout = tf.keras.layers.Dropout(config.dropout) + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_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 (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.BlenderbotSmallTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + encoder_hidden_states (:obj:`tf.Tensor` of shape :obj:`(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 (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape :obj:`(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 :obj:`past_key_values` are used, the user can optionally input only the last + :obj:`decoder_input_ids` (those that don't have their past key value states given to this model) of + shape :obj:`(batch_size, 1)` instead of all :obj:`decoder_input_ids`` of shape :obj:`(batch_size, + sequence_length)`. + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + 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 = ( + 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["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length) + else: + combined_attention_mask = _expand_mask( + tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] + ) + + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask( + inputs["attention_mask"], tgt_len=input_shape[-1] + ) + + 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(input_shape, past_key_values_length) + + hidden_states = self.layernorm_embedding(inputs["inputs_embeds"]) + positions + hidden_states = self.dropout(hidden_states, training=inputs["training"]) + + # decoder layers + all_hidden_states = () + all_self_attns = () + present_key_values = () + + 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 + + hidden_states, layer_self_attn, present_key_value = decoder_layer( + hidden_states, + attention_mask=combined_attention_mask, + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=inputs["encoder_attention_mask"], + past_key_value=past_key_value, + ) + + if inputs["use_cache"]: + present_key_values += (present_key_value,) + + if inputs["output_attentions"]: + all_self_attns += (layer_self_attn,) + + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + else: + all_hidden_states = None + + all_self_attns = list(all_self_attns) if inputs["output_attentions"] else None + + present_key_values = (encoder_hidden_states, present_key_values) if inputs["use_cache"] else None + + if not inputs["return_dict"]: + return hidden_states, present_key_values, all_hidden_states, all_self_attns + else: + return TFBaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=present_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +@add_start_docstrings( + "The bare BLENDERBOT_SMALL Model outputting raw hidden-states without any specific head on top.", + BLENDERBOT_SMALL_START_DOCSTRING, +) +@keras_serializable +class TFBlenderbotSmallModel(TFBlenderbotSmallPreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config: BlenderbotSmallConfig, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.shared = TFSharedEmbeddings(config.vocab_size, config.d_model, config.pad_token_id, name="model.shared") + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + # Wraps layer to avoid problems with weight restoring and ensuring we're in the correct TF scope. + embed_tokens = TFWrappedEmbeddings(self.shared, abs_scope_name=shared_abs_scope_name) + embed_tokens.vocab_size = self.shared.vocab_size + embed_tokens.hidden_size = self.shared.hidden_size + + self.encoder = TFBlenderbotSmallEncoder(config, embed_tokens, name="encoder") + self.decoder = TFBlenderbotSmallDecoder(config, embed_tokens, name="decoder") + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(BLENDERBOT_SMALL_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + @replace_return_docstrings(output_type=TFSeq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + r""" + Returns: + + Example:: + + >>> from transformers import BlenderbotSmallTokenizer, TFBlenderbotSmallModel + + >>> model = TFBlenderbotSmallModel.from_pretrained("facebook/blenderbot_small-90M") + >>> tokenizer = BlenderbotSmallTokenizer.from_pretrained("facebook/blenderbot_small-90M") + + >>> input_ids = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="tf").input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="tf").input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + + >>> last_hidden_states = outputs.last_hidden_state + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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, + ) + + inputs["output_hidden_states"] = ( + inputs["output_hidden_states"] + if inputs["output_hidden_states"] is not None + else self.config.output_hidden_states + ) + + if inputs["encoder_outputs"] is None: + inputs["encoder_outputs"] = self.encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + inputs_embeds=inputs["inputs_embeds"], + 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() + + decoder_outputs = self.decoder( + inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=inputs["encoder_outputs"][0], + encoder_attention_mask=inputs["attention_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, + encoder_last_hidden_state=inputs["encoder_outputs"].last_hidden_state, + encoder_hidden_states=inputs["encoder_outputs"].hidden_states, + encoder_attentions=inputs["encoder_outputs"].attentions, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartModel.serving_output + 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 + 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, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + +@add_start_docstrings( + "The BLENDERBOT_SMALL Model with a language modeling head. Can be used for summarization.", + BLENDERBOT_SMALL_START_DOCSTRING, +) +class TFBlenderbotSmallForConditionalGeneration(TFBlenderbotSmallPreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + r"model.encoder.embed_tokens.weight", + r"model.decoder.embed_tokens.weight", + ] + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.model = TFBlenderbotSmallModel(config, name="model") + self.use_cache = config.use_cache + # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency. + self.final_logits_bias = self.add_weight( + name="final_logits_bias", shape=[1, config.vocab_size], initializer="zeros", trainable=False + ) + + def get_decoder(self): + return self.model.decoder + + def get_encoder(self): + return self.model.encoder + + def get_output_embeddings(self): + return self.get_input_embeddings() + + def set_output_embeddings(self, value): + self.set_input_embeddings(value) + + def get_bias(self): + return {"final_logits_bias": self.final_logits_bias} + + def set_bias(self, value): + self.final_logits_bias = value["final_logits_bias"] + + @add_start_docstrings_to_model_forward(BLENDERBOT_SMALL_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + @add_end_docstrings(BLENDERBOT_SMALL_GENERATION_EXAMPLE) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[TFBaseModelOutput] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + labels=None, + training=False, + **kwargs, + ): + r""" + labels (:obj:`torch.LongTensor` of shape :obj:`(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: + + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + labels=labels, + training=training, + kwargs_call=kwargs, + ) + + 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( + inputs["input_ids"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + encoder_outputs=inputs["encoder_outputs"], + decoder_attention_mask=inputs["decoder_attention_mask"], + past_key_values=inputs["past_key_values"], + inputs_embeds=inputs["inputs_embeds"], + 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.model.shared(outputs[0], mode="linear") + lm_logits = lm_logits + self.final_logits_bias + masked_lm_loss = None if inputs["labels"] is None else self.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, # index 1 of d outputs + decoder_hidden_states=outputs.decoder_hidden_states, # index 2 of d outputs + decoder_attentions=outputs.decoder_attentions, # index 3 of d outputs + encoder_last_hidden_state=outputs.last_hidden_state, # index 0 of encoder outputs + encoder_hidden_states=outputs.encoder_hidden_states, # 1 of e out + encoder_attentions=outputs.encoder_attentions, # 2 of e out + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.serving_output + 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 + 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 TFSeq2SeqLMOutput( + logits=output.logits, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.prepare_inputs_for_generation + def prepare_inputs_for_generation(self, decoder_input_ids, past, attention_mask, use_cache, **kwargs) -> Dict: + assert past is not None and len(past) in {1, 2}, f"past has to be an iterable of length 1,2 got {past}" + if len(past) == 1: + assert isinstance(past[0], tf.Tensor), f"`past[0]` has to be of type `tf.Tensor`, but is {type(past[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=past[0]) + past_key_values = None + else: + assert ( + len(past) == 2 + ), "`past` has to be of length 2 with the encoder_outputs at the first position and past_key_values at the second position." + encoder_outputs, past_key_values = past + if isinstance(encoder_outputs, tuple): + assert isinstance( + encoder_outputs[0], tf.Tensor + ), f"`encoder_outputs[0]` has to be of type `tf.Tensor`, but is {type(encoder_outputs[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs[0]) + elif isinstance(encoder_outputs, tf.Tensor): + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs) + assert ( + past_key_values + ), f"decoder cached states must be truthy. got {past_key_values} from the 2nd element of past" + decoder_input_ids = decoder_input_ids[:, -1:] + + assert isinstance( + encoder_outputs, TFBaseModelOutput + ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}." + return { + "input_ids": None, # encoder_outputs is defined. input_ids not needed + "encoder_outputs": encoder_outputs, + "past_key_values": past_key_values, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "use_cache": use_cache, # change this to avoid caching (presumably for debugging) + } + + @staticmethod + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration._reorder_cache + 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) + + def adjust_logits_during_generation(self, logits, cur_len, max_length): + if cur_len == max_length - 1: + vocab_range = tf.constant(range(self.config.vocab_size)) + return tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits) + else: + return logits + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.compute_loss + def compute_loss(self, labels, logits): + """CrossEntropyLoss that ignores pad tokens""" + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True, + reduction=tf.keras.losses.Reduction.NONE, + ) + melted_labels = tf.reshape(labels, (-1,)) + active_loss = tf.not_equal(melted_labels, self.config.pad_token_id) + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(melted_labels, active_loss) + return loss_fn(labels, reduced_logits) diff --git a/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py b/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py index 0c8b821ff3..69c60864e4 100644 --- a/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py +++ b/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py @@ -30,7 +30,7 @@ logger = logging.get_logger(__name__) VOCAB_FILES_NAMES = { "vocab_file": "vocab.json", "merges_file": "merges.txt", - # "tokenizer_config_file": "tokenizer_config.json", + "tokenizer_config_file": "tokenizer_config.json", } @@ -75,13 +75,20 @@ class BlenderbotSmallTokenizer(PreTrainedTokenizer): Additional keyword arguments passed along to :class:`~transformers.PreTrainedTokenizer` """ - vocab_files_names = {"vocab_file": "vocab.json", "merges_file": "merges.txt"} + vocab_files_names = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", + "tokenizer_config": "tokenizer_config.json", + } pretrained_vocab_files_map = { "vocab_file": { - "facebook/blenderbot_small-90M": "https://cdn.huggingface.co/facebook/blenderbot_small-90M/vocab.json" + "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/blob/main/vocab.json" }, "merges_file": { - "facebook/blenderbot_small-90M": "https://cdn.huggingface.co/facebook/blenderbot_small-90M/merges.txt" + "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/blob/main/merges.txt" + }, + "tokenizer_config_file": { + "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/blob/main/tokenizer.json" }, } max_model_input_sizes = {"facebook/blenderbot_small-90M": 512} diff --git a/src/transformers/models/led/modeling_tf_led.py b/src/transformers/models/led/modeling_tf_led.py index 373726674f..c93569d5f8 100644 --- a/src/transformers/models/led/modeling_tf_led.py +++ b/src/transformers/models/led/modeling_tf_led.py @@ -1475,7 +1475,7 @@ LED_INPUTS_DOCSTRING = r""" Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for more detail. return_dict (:obj:`bool`, `optional`): - Whether or not to return a :class:`~transformers.file_utils.TFModelOutput` instead of a plain tuple. + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. training (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether or not to use the model in training mode (some modules like dropout modules have different behaviors between training and evaluation). diff --git a/src/transformers/models/marian/__init__.py b/src/transformers/models/marian/__init__.py index 880c5df934..ab1c5e8130 100644 --- a/src/transformers/models/marian/__init__.py +++ b/src/transformers/models/marian/__init__.py @@ -42,7 +42,7 @@ if is_torch_available(): ] if is_tf_available(): - _import_structure["modeling_tf_marian"] = ["TFMarianMTModel"] + _import_structure["modeling_tf_marian"] = ["TFMarianMTModel", "TFMarianModel"] if TYPE_CHECKING: @@ -60,7 +60,7 @@ if TYPE_CHECKING: ) if is_tf_available(): - from .modeling_tf_marian import TFMarianMTModel + from .modeling_tf_marian import TFMarianModel, TFMarianMTModel else: import importlib diff --git a/src/transformers/models/marian/configuration_marian.py b/src/transformers/models/marian/configuration_marian.py index cc54540aba..7e57b6e975 100644 --- a/src/transformers/models/marian/configuration_marian.py +++ b/src/transformers/models/marian/configuration_marian.py @@ -159,17 +159,6 @@ class MarianConfig(PretrainedConfig): self.gradient_checkpointing = gradient_checkpointing self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True - # IMPORTANT - # DELETE ALL OF THE FOLLOWING LINES AS SOON AS TF IS READY - self.extra_pos_embeddings = 0 - self.normalize_before = False - self.add_final_layer_norm = False - self.do_blenderbot_90_layernorm = False - self.normalize_embedding = False - self.static_position_embeddings = True - self.add_bias_logits = False - self.force_bos_token_to_be_generated = False - @property def num_attention_heads(self) -> int: return self.encoder_attention_heads diff --git a/src/transformers/models/marian/modeling_tf_marian.py b/src/transformers/models/marian/modeling_tf_marian.py index f17182306e..43171f9fdb 100644 --- a/src/transformers/models/marian/modeling_tf_marian.py +++ b/src/transformers/models/marian/modeling_tf_marian.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team. +# Copyright 2021 The Marian Team 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. @@ -12,36 +12,1286 @@ # 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. -"""TF Marian model, ported from the fairseq repo.""" +""" TF 2.0 Marian model. """ -from ...file_utils import add_start_docstrings, is_tf_available + +import random +from typing import Dict, Optional, Tuple, Union + +import numpy as np +import tensorflow as tf + +from ...activations_tf import get_tf_activation +from ...file_utils import ( + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ...modeling_tf_outputs import ( + TFBaseModelOutput, + TFBaseModelOutputWithPast, + TFSeq2SeqLMOutput, + TFSeq2SeqModelOutput, +) + +# Public API +from ...modeling_tf_utils import ( + DUMMY_INPUTS, + TFPreTrainedModel, + TFSharedEmbeddings, + TFWrappedEmbeddings, + input_processing, + keras_serializable, + shape_list, +) from ...utils import logging -from ..bart.modeling_tf_bart import BART_START_DOCSTRING, LARGE_NEGATIVE, TFBartForConditionalGeneration from .configuration_marian import MarianConfig -if is_tf_available(): - import tensorflow as tf - - -_CONFIG_FOR_DOC = "MarianConfig" - -START_DOCSTRING = BART_START_DOCSTRING.replace( - "inherits from :class:`~transformers.TFPreTrainedModel`", - "inherits from :class:`~transformers.TFBartForConditionalGeneration`", -).replace("BartConfig", _CONFIG_FOR_DOC) - - logger = logging.get_logger(__name__) +_CONFIG_FOR_DOC = "MarianConfig" +_TOKENIZER_FOR_DOC = "MarianTokenizer" -@add_start_docstrings("Marian model for machine translation", START_DOCSTRING) -class TFMarianMTModel(TFBartForConditionalGeneration): - _keys_to_ignore_on_load_missing = [ - r"model.encoder.embed_positions.weight", - r"model.decoder.embed_positions.weight", - ] + +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): + shifted_input_ids = tf.cast(input_ids, tf.int32) + shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1) + start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), decoder_start_token_id) + shifted_input_ids = tf.concat([start_tokens, shifted_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 + ) + + # "Verify that `labels` has only positive values and -100" + assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, tf.int32)) + + # 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), dtype=tf.float32) * 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) + mask = tf.cast(mask, tf.float32) + + if past_key_values_length > 0: + mask = tf.concat([tf.zeros((tgt_len, past_key_values_length), dtype=tf.float32), mask], axis=-1) + return tf.broadcast_to(mask[None, None, :, :], (bsz, 1, tgt_len, tgt_len + past_key_values_length)) + + +# 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]`. + """ + bsz, src_len = shape_list(mask) + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = tf.cast(tf.broadcast_to(mask[:, None, None, :], (bsz, 1, tgt_len, src_len)), tf.float32) + + return (1.0 - expanded_mask) * LARGE_NEGATIVE + + +class TFMarianSinusoidalPositionalEmbedding(tf.keras.layers.Embedding): + """This module produces sinusoidal positional embeddings of any length.""" + + def __init__(self, num_positions: int, embedding_dim: int, **kwargs): + + if embedding_dim % 2 != 0: + raise NotImplementedError(f"odd embedding_dim {embedding_dim} not supported") + super().__init__( + num_positions, + embedding_dim, + **kwargs, + ) + + 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 + """ + super().build(input_shape) # Instantiates self.weight so it can be loaded + weight: np.ndarray = self._init_weight(self.input_dim, self.output_dim) + self.set_weights([weight]) # overwrite self.weight to correct value + + @staticmethod + def _init_weight(n_pos: int, dim: int): + """ + Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in + the 2nd half of the vector. [dim // 2:] + """ + position_enc = np.array( + [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)] + ) + # index 0 is all zero + position_enc[:, 0 : dim // 2] = np.sin(position_enc[:, 0::2]) + position_enc[:, dim // 2 :] = np.cos(position_enc[:, 1::2]) + # convert to tensor + table = tf.convert_to_tensor(position_enc, dtype=tf.float32) + tf.stop_gradient(table) + return table + + def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): + """Input is expected to be of size [bsz x seqlen].""" + bsz, seq_len = input_shape[:2] + + positions = tf.range( + past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" + ) + return super().call(positions) + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartAttention with Bart->Marian +class TFMarianAttention(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 + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by 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, + 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) + + 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: + 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)}", + ) + 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) + + attn_probs = self.dropout(attn_weights, training=training) + + attn_output = tf.matmul(attn_probs, value_states) + + 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 + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartEncoderLayer with Bart->Marian +class TFMarianEncoderLayer(tf.keras.layers.Layer): + def __init__(self, config: MarianConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFMarianAttention( + 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, training=False): + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + """ + residual = hidden_states + hidden_states, self_attn_weights, _ = self.self_attn( + hidden_states=hidden_states, attention_mask=attention_mask + ) + 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 + hidden_states = self.self_attn_layer_norm(hidden_states) + + residual = 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 + hidden_states = self.final_layer_norm(hidden_states) + + return hidden_states, self_attn_weights + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartDecoderLayer with Bart->Marian +class TFMarianDecoderLayer(tf.keras.layers.Layer): + def __init__(self, config: MarianConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFMarianAttention( + 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 = TFMarianAttention( + 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, + past_key_value: Optional[Tuple[tf.Tensor]] = None, + training=False, + ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`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 (:obj:`tf.Tensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (:obj:`tf.Tensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + past_key_value (:obj:`Tuple(tf.Tensor)`): cached past key and value projection states + """ + residual = 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, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Cross-Attention Block + cross_attn_present_key_value = None + if encoder_hidden_states is not None: + residual = 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_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_value=cross_attn_past_key_value, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + hidden_states = self.encoder_attn_layer_norm(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.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 + hidden_states = self.final_layer_norm(hidden_states) + + return ( + hidden_states, + self_attn_weights, + present_key_value, + ) + + +class TFMarianPreTrainedModel(TFPreTrainedModel): config_class = MarianConfig + base_model_prefix = "model" + + @property + def dummy_inputs(self): + pad_token = 1 + input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + dummy_inputs = { + "decoder_input_ids": decoder_input_ids, + "attention_mask": tf.math.not_equal(input_ids, pad_token), + "input_ids": input_ids, + } + return dummy_inputs + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.get_input_embeddings + def get_input_embeddings(self): + base_model = getattr(self, self.base_model_prefix, self) + + return base_model.shared + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.set_input_embeddings + def set_input_embeddings(self, value): + base_model = getattr(self, self.base_model_prefix, self) + + try: + base_model.shared.weight = value + except AttributeError: + self(self.dummy_inputs) + base_model.shared.weight = value + + base_model.shared.vocab_size = shape_list(base_model.shared.weight)[0] + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + embed_tokens = TFWrappedEmbeddings(base_model.shared, abs_scope_name=shared_abs_scope_name) + base_model.encoder.set_embed_tokens(embed_tokens) + base_model.decoder.set_embed_tokens(embed_tokens) + + @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"), + } + ] + ) + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.serving + def serving(self, inputs): + output = self.call(inputs) + + return self.serving_output(output) + + +MARIAN_START_DOCSTRING = r""" + This model inherits from :class:`~transformers.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 `__ 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. + + .. note:: + + 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 :meth:`tf.keras.Model.fit` method which currently requires having all + the tensors in the first argument of the model call function: :obj:`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 :obj:`input_ids` only and nothing else: :obj:`model(input_ids)` + - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: + :obj:`model([input_ids, attention_mask])` or :obj:`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: + :obj:`model({"input_ids": input_ids, "token_type_ids": token_type_ids})` + + Args: + config (:class:`~transformers.MarianConfig`): 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 :meth:`~transformers.TFPreTrainedModel.from_pretrained` method to load the + model weights. +""" + +MARIAN_GENERATION_EXAMPLE = r""" + TF version of marian-nmt's transformer.h (c++). Designed for the OPUS-NMT translation checkpoints. Available + models are listed `here `__. + + Examples:: + + >>> from transformers import MarianTokenizer, TFMarianMTModel + >>> from typing import List + >>> src = 'fr' # source language + >>> trg = 'en' # target language + >>> sample_text = "où est l'arrêt de bus ?" + >>> mname = f'Helsinki-NLP/opus-mt-{src}-{trg}' + + >>> model = MarianMTModel.from_pretrained(mname) + >>> tok = MarianTokenizer.from_pretrained(mname) + >>> batch = tok.prepare_seq2seq_batch(src_texts=[sample_text], return_tensors="tf") # don't need tgt_text for inference + >>> gen = model.generate(**batch) + >>> words: List[str] = tok.batch_decode(gen, skip_special_tokens=True) # returns "Where is the bus stop ?" +""" + +MARIAN_INPUTS_DOCSTRING = r""" + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.MarianTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`({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.html#attention-mask>`__ + decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.MarianTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + Marian uses the :obj:`pad_token_id` as the starting token for :obj:`decoder_input_ids` generation. If + :obj:`past_key_values` is used, optionally only the last :obj:`decoder_input_ids` have to be input (see + :obj:`past_key_values`). + decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. + encoder_outputs (:obj:`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 :obj:`(batch_size, sequence_length, hidden_size)` is a sequence of + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers`) + contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). Set to :obj:`False` during training, :obj:`True` during generation + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for + more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. + training (:obj:`bool`, `optional`, defaults to :obj:`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 TFMarianEncoder(tf.keras.layers.Layer): + config_class = MarianConfig + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + :class:`TFMarianEncoderLayer`. + + Args: + config: MarianConfig + """ + + def __init__(self, config: MarianConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.layerdrop = config.encoder_layerdrop + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_position_embeddings + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + + self.embed_tokens = embed_tokens + self.embed_positions = TFMarianSinusoidalPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + name="embed_positions", + ) + self.layers = [TFMarianEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)] + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + """ + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.MarianTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + inputs_embeds=inputs_embeds, + 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 input_ids and 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 input_ids or inputs_embeds") + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + embed_pos = self.embed_positions(input_shape) + hidden_states = inputs["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] + attention_mask = _expand_mask(inputs["attention_mask"]) + else: + attention_mask = None + + encoder_states = () if inputs["output_hidden_states"] else None + all_attentions = () if inputs["output_attentions"] else None + + # encoder layers + for encoder_layer in 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, attention_mask) + + if inputs["output_attentions"]: + all_attentions += (attn,) + + 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 TFMarianDecoder(tf.keras.layers.Layer): + config_class = MarianConfig + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a :class:`TFMarianDecoderLayer` + + Args: + config: MarianConfig + embed_tokens: output embedding + """ + + def __init__(self, config: MarianConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.padding_idx = config.pad_token_id + self.embed_tokens = embed_tokens + self.layerdrop = config.decoder_layerdrop + self.embed_positions = TFMarianSinusoidalPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + name="embed_positions", + ) + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + self.layers = [TFMarianDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)] + + self.dropout = tf.keras.layers.Dropout(config.dropout) + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_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 (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.MarianTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + encoder_hidden_states (:obj:`tf.Tensor` of shape :obj:`(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 (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape :obj:`(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 :obj:`past_key_values` are used, the user can optionally input only the last + :obj:`decoder_input_ids` (those that don't have their past key value states given to this model) of + shape :obj:`(batch_size, 1)` instead of all :obj:`decoder_input_ids`` of shape :obj:`(batch_size, + sequence_length)`. + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + 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 = ( + shape_list(inputs["past_key_values"][0][0])[2] if inputs["past_key_values"] is not None else 0 + ) + + # embed positions + positions = self.embed_positions(input_shape, past_key_values_length) + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + hidden_states = inputs["inputs_embeds"] + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length) + else: + combined_attention_mask = _expand_mask( + tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] + ) + + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask( + inputs["attention_mask"], tgt_len=input_shape[-1] + ) + + 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]) + + hidden_states = self.dropout(hidden_states + positions, training=inputs["training"]) + + # decoder layers + all_hidden_states = () + all_self_attns = () + present_key_values = () + 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 + + hidden_states, layer_self_attn, present_key_value = decoder_layer( + hidden_states, + attention_mask=combined_attention_mask, + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=inputs["encoder_attention_mask"], + past_key_value=past_key_value, + ) + + if inputs["use_cache"]: + present_key_values += (present_key_value,) + + if inputs["output_attentions"]: + all_self_attns += (layer_self_attn,) + + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + else: + all_hidden_states = None + + all_self_attns = list(all_self_attns) if inputs["output_attentions"] else None + + present_key_values = (encoder_hidden_states, present_key_values) if inputs["use_cache"] else None + + if not inputs["return_dict"]: + return hidden_states, present_key_values, all_hidden_states, all_self_attns + else: + return TFBaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=present_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +@add_start_docstrings( + "The bare MARIAN Model outputting raw hidden-states without any specific head on top.", + MARIAN_START_DOCSTRING, +) +@keras_serializable +class TFMarianModel(TFMarianPreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config: MarianConfig, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.shared = TFSharedEmbeddings(config.vocab_size, config.d_model, config.pad_token_id, name="model.shared") + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + # Wraps layer to avoid problems with weight restoring and ensuring we're in the correct TF scope. + embed_tokens = TFWrappedEmbeddings(self.shared, abs_scope_name=shared_abs_scope_name) + embed_tokens.vocab_size = self.shared.vocab_size + embed_tokens.hidden_size = self.shared.hidden_size + + self.encoder = TFMarianEncoder(config, embed_tokens, name="encoder") + self.decoder = TFMarianDecoder(config, embed_tokens, name="decoder") + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(MARIAN_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + @replace_return_docstrings(output_type=TFSeq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + r""" + Returns: + + Example:: + + >>> from transformers import MarianTokenizer, TFMarianModel + + >>> tokenizer = MarianTokenizer.from_pretrained('Helsinki-NLP/opus-mt-en-de') + >>> model = TFMarianModel.from_pretrained('Helsinki-NLP/opus-mt-en-de') + + >>> input_ids = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="tf").input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer(" Studien haben gezeigt dass es hilfreich ist einen Hund zu besitzen", + ... return_tensors="tf", add_special_tokens=False).input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + + >>> last_hidden_states = outputs.last_hidden_state + """ + + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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 inputs["decoder_input_ids"] is None and inputs["decoder_inputs_embeds"] is None: + inputs["use_cache"] = False + + inputs["output_hidden_states"] = ( + inputs["output_hidden_states"] + if inputs["output_hidden_states"] is not None + else self.config.output_hidden_states + ) + + if inputs["encoder_outputs"] is None: + inputs["encoder_outputs"] = self.encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + inputs_embeds=inputs["inputs_embeds"], + 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() + + decoder_outputs = self.decoder( + inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=inputs["encoder_outputs"][0], + encoder_attention_mask=inputs["attention_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, + encoder_last_hidden_state=inputs["encoder_outputs"].last_hidden_state, + encoder_hidden_states=inputs["encoder_outputs"].hidden_states, + encoder_attentions=inputs["encoder_outputs"].attentions, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartModel.serving_output + 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 + 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, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + +@add_start_docstrings( + "The MARIAN Model with a language modeling head. Can be used for summarization.", + MARIAN_START_DOCSTRING, +) +class TFMarianMTModel(TFMarianPreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + r"model.encoder.embed_tokens.weight", + r"model.decoder.embed_tokens.weight", + ] + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.model = TFMarianModel(config, name="model") + self.use_cache = config.use_cache + # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency. + self.final_logits_bias = self.add_weight( + name="final_logits_bias", shape=[1, config.vocab_size], initializer="zeros", trainable=False + ) + + def get_decoder(self): + return self.model.decoder + + def get_encoder(self): + return self.model.encoder + + def get_output_embeddings(self): + return self.get_input_embeddings() + + def set_output_embeddings(self, value): + self.set_input_embeddings(value) + + def get_bias(self): + return {"final_logits_bias": self.final_logits_bias} + + def set_bias(self, value): + self.final_logits_bias = value["final_logits_bias"] + + @add_start_docstrings_to_model_forward(MARIAN_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + @add_end_docstrings(MARIAN_GENERATION_EXAMPLE) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[TFBaseModelOutput] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + labels=None, + training=False, + **kwargs, + ): + r""" + labels (:obj:`torch.LongTensor` of shape :obj:`(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: + + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + labels=labels, + training=training, + kwargs_call=kwargs, + ) + + if inputs["labels"] is not None: + inputs["use_cache"] = False + 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( + inputs["input_ids"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + encoder_outputs=inputs["encoder_outputs"], + decoder_attention_mask=inputs["decoder_attention_mask"], + past_key_values=inputs["past_key_values"], + inputs_embeds=inputs["inputs_embeds"], + 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.model.shared(outputs[0], mode="linear") + lm_logits = lm_logits + self.final_logits_bias + masked_lm_loss = None if inputs["labels"] is None else self.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, # index 1 of d outputs + decoder_hidden_states=outputs.decoder_hidden_states, # index 2 of d outputs + decoder_attentions=outputs.decoder_attentions, # index 3 of d outputs + encoder_last_hidden_state=outputs.last_hidden_state, # index 0 of encoder outputs + encoder_hidden_states=outputs.encoder_hidden_states, # 1 of e out + encoder_attentions=outputs.encoder_attentions, # 2 of e out + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.serving_output + 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 + 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 TFSeq2SeqLMOutput( + logits=output.logits, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.prepare_inputs_for_generation + def prepare_inputs_for_generation(self, decoder_input_ids, past, attention_mask, use_cache, **kwargs) -> Dict: + assert past is not None and len(past) in {1, 2}, f"past has to be an iterable of length 1,2 got {past}" + if len(past) == 1: + assert isinstance(past[0], tf.Tensor), f"`past[0]` has to be of type `tf.Tensor`, but is {type(past[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=past[0]) + past_key_values = None + else: + assert ( + len(past) == 2 + ), "`past` has to be of length 2 with the encoder_outputs at the first position and past_key_values at the second position." + encoder_outputs, past_key_values = past + if isinstance(encoder_outputs, tuple): + assert isinstance( + encoder_outputs[0], tf.Tensor + ), f"`encoder_outputs[0]` has to be of type `tf.Tensor`, but is {type(encoder_outputs[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs[0]) + elif isinstance(encoder_outputs, tf.Tensor): + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs) + assert ( + past_key_values + ), f"decoder cached states must be truthy. got {past_key_values} from the 2nd element of past" + decoder_input_ids = decoder_input_ids[:, -1:] + + assert isinstance( + encoder_outputs, TFBaseModelOutput + ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}." + return { + "input_ids": None, # encoder_outputs is defined. input_ids not needed + "encoder_outputs": encoder_outputs, + "past_key_values": past_key_values, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "use_cache": use_cache, # change this to avoid caching (presumably for debugging) + } + + @staticmethod + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration._reorder_cache + 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) def adjust_logits_during_generation(self, logits, cur_len, max_length): """Never predict pad_token_id. Predict when max_length is reached.""" @@ -50,3 +1300,16 @@ class TFMarianMTModel(TFBartForConditionalGeneration): if cur_len == max_length - 1: logits = tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits) return logits + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.compute_loss + def compute_loss(self, labels, logits): + """CrossEntropyLoss that ignores pad tokens""" + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True, + reduction=tf.keras.losses.Reduction.NONE, + ) + melted_labels = tf.reshape(labels, (-1,)) + active_loss = tf.not_equal(melted_labels, self.config.pad_token_id) + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(melted_labels, active_loss) + return loss_fn(labels, reduced_logits) diff --git a/src/transformers/models/mbart/__init__.py b/src/transformers/models/mbart/__init__.py index d14df702f6..db2aca8494 100644 --- a/src/transformers/models/mbart/__init__.py +++ b/src/transformers/models/mbart/__init__.py @@ -47,7 +47,7 @@ if is_torch_available(): ] if is_tf_available(): - _import_structure["modeling_tf_mbart"] = ["TFMBartForConditionalGeneration"] + _import_structure["modeling_tf_mbart"] = ["TFMBartForConditionalGeneration", "TFMBartModel"] if TYPE_CHECKING: @@ -70,7 +70,7 @@ if TYPE_CHECKING: ) if is_tf_available(): - from .modeling_tf_mbart import TFMBartForConditionalGeneration + from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel else: import importlib diff --git a/src/transformers/models/mbart/configuration_mbart.py b/src/transformers/models/mbart/configuration_mbart.py index 4fbacd4c74..b81a713fd9 100644 --- a/src/transformers/models/mbart/configuration_mbart.py +++ b/src/transformers/models/mbart/configuration_mbart.py @@ -159,17 +159,6 @@ class MBartConfig(PretrainedConfig): self.gradient_checkpointing = gradient_checkpointing self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True - # IMPORTANT - # DELETE ALL OF THE FOLLOWING LINES AS SOON AS TF IS READY - self.extra_pos_embeddings = 2 - self.normalize_before = True - self.add_final_layer_norm = True - self.do_blenderbot_90_layernorm = False - self.normalize_embedding = True - self.static_position_embeddings = False - self.add_bias_logits = False - self.force_bos_token_to_be_generated = False - @property def num_attention_heads(self) -> int: return self.encoder_attention_heads diff --git a/src/transformers/models/mbart/modeling_tf_mbart.py b/src/transformers/models/mbart/modeling_tf_mbart.py index 23b30fd4b3..16c5c854b8 100644 --- a/src/transformers/models/mbart/modeling_tf_mbart.py +++ b/src/transformers/models/mbart/modeling_tf_mbart.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team. +# 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. @@ -12,25 +12,1290 @@ # 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. -"""TF mBART model, originally from fairseq.""" -from ...file_utils import add_start_docstrings +""" TF 2.0 MBart model. """ + + +import random +from typing import Dict, Optional, Tuple, Union + +import tensorflow as tf + +from ...activations_tf import get_tf_activation +from ...file_utils import ( + add_code_sample_docstrings, + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ...modeling_tf_outputs import ( + TFBaseModelOutput, + TFBaseModelOutputWithPast, + TFSeq2SeqLMOutput, + TFSeq2SeqModelOutput, +) + +# Public API +from ...modeling_tf_utils import ( + DUMMY_INPUTS, + TFPreTrainedModel, + TFSharedEmbeddings, + TFWrappedEmbeddings, + input_processing, + keras_serializable, + shape_list, +) from ...utils import logging -from ..bart.modeling_tf_bart import BART_START_DOCSTRING, TFBartForConditionalGeneration from .configuration_mbart import MBartConfig -_CONFIG_FOR_DOC = "MBartConfig" - -START_DOCSTRING = BART_START_DOCSTRING.replace( - "inherits from :class:`~transformers.TFPreTrainedModel`", - "inherits from :class:`~transformers.TFBartForConditionalGeneration`", -).replace("BartConfig", _CONFIG_FOR_DOC) - - logger = logging.get_logger(__name__) +_CONFIG_FOR_DOC = "MBartConfig" +_TOKENIZER_FOR_DOC = "MBartTokenizer" -@add_start_docstrings("mBART (multilingual BART) model for machine translation", START_DOCSTRING) -class TFMBartForConditionalGeneration(TFBartForConditionalGeneration): + +LARGE_NEGATIVE = -1e8 + + +def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int): + """ + Shift input ids one token to the right, and wrap the last non pad token (the token) Note that MBart does not + have a single `decoder_start_token_id` in contrast to other Bart-like models. + """ + prev_output_tokens = tf.cast(input_ids, tf.int32) + assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined." + # replace possible -100 values in labels by `pad_token_id` + prev_output_tokens = tf.where( + prev_output_tokens == -100, tf.fill(shape_list(prev_output_tokens), pad_token_id), prev_output_tokens + ) + language_id_index = ( + tf.reduce_sum(tf.cast(tf.math.not_equal(prev_output_tokens, pad_token_id), tf.int32), axis=-1) - 1 + ) + language_id_index = tf.stack([tf.range(shape_list(input_ids)[0]), language_id_index], axis=-1) + languages_ids = tf.gather_nd(prev_output_tokens, language_id_index) + + shifted_input_ids = tf.concat([tf.expand_dims(languages_ids, axis=-1), prev_output_tokens[:, :-1]], axis=-1) + + 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), dtype=tf.float32) * 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) + mask = tf.cast(mask, tf.float32) + + if past_key_values_length > 0: + mask = tf.concat([tf.zeros((tgt_len, past_key_values_length), dtype=tf.float32), mask], axis=-1) + return tf.broadcast_to(mask[None, None, :, :], (bsz, 1, tgt_len, tgt_len + past_key_values_length)) + + +# 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]`. + """ + bsz, src_len = shape_list(mask) + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = tf.cast(tf.broadcast_to(mask[:, None, None, :], (bsz, 1, tgt_len, src_len)), tf.float32) + + return (1.0 - expanded_mask) * LARGE_NEGATIVE + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartLearnedPositionalEmbedding with Bart->MBart +class TFMBartLearnedPositionalEmbedding(TFSharedEmbeddings): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, **kwargs): + assert padding_idx is not None, "padding_idx cannot be None" + # MBart is set up so that if padding_idx is specified then offset the embedding ids by 2 + # and adjust num_embeddings appropriately. Other models dont have this hack + self.offset = 2 + super().__init__(num_embeddings + self.offset, embedding_dim, **kwargs) + + def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): + """Input is expected to be of size [bsz x seqlen].""" + bsz, seq_len = input_shape[:2] + + positions = tf.range( + past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" + ) + return super().call(positions + self.offset) + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartAttention with Bart->MBart +class TFMBartAttention(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 + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by 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, + 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) + + 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: + 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)}", + ) + 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) + + attn_probs = self.dropout(attn_weights, training=training) + + attn_output = tf.matmul(attn_probs, value_states) + + 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 TFMBartEncoderLayer(tf.keras.layers.Layer): + def __init__(self, config: MBartConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFMBartAttention( + 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, training=False): + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + """ + 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 + ) + 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 TFMBartDecoderLayer(tf.keras.layers.Layer): + def __init__(self, config: MBartConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFMBartAttention( + 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 = TFMBartAttention( + 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, + past_key_value: Optional[Tuple[tf.Tensor]] = None, + training=False, + ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`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 (:obj:`tf.Tensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (:obj:`tf.Tensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + past_key_value (:obj:`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, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + # Cross-Attention Block + cross_attn_present_key_value = 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_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_value=cross_attn_past_key_value, + ) + 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, + present_key_value, + ) + + +class TFMBartPreTrainedModel(TFPreTrainedModel): config_class = MBartConfig - # All the code is in src/transformers/models/bart/modeling_tf_bart.py + base_model_prefix = "model" + + @property + def dummy_inputs(self): + pad_token = 1 + input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + dummy_inputs = { + "decoder_input_ids": decoder_input_ids, + "attention_mask": tf.math.not_equal(input_ids, pad_token), + "input_ids": input_ids, + } + return dummy_inputs + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.get_input_embeddings + def get_input_embeddings(self): + base_model = getattr(self, self.base_model_prefix, self) + + return base_model.shared + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.set_input_embeddings + def set_input_embeddings(self, value): + base_model = getattr(self, self.base_model_prefix, self) + + try: + base_model.shared.weight = value + except AttributeError: + self(self.dummy_inputs) + base_model.shared.weight = value + + base_model.shared.vocab_size = shape_list(base_model.shared.weight)[0] + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + embed_tokens = TFWrappedEmbeddings(base_model.shared, abs_scope_name=shared_abs_scope_name) + base_model.encoder.set_embed_tokens(embed_tokens) + base_model.decoder.set_embed_tokens(embed_tokens) + + @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"), + } + ] + ) + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.serving + def serving(self, inputs): + output = self.call(inputs) + + return self.serving_output(output) + + +MBART_START_DOCSTRING = r""" + This model inherits from :class:`~transformers.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 `__ 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. + + .. note:: + + 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 :meth:`tf.keras.Model.fit` method which currently requires having all + the tensors in the first argument of the model call function: :obj:`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 :obj:`input_ids` only and nothing else: :obj:`model(input_ids)` + - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: + :obj:`model([input_ids, attention_mask])` or :obj:`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: + :obj:`model({"input_ids": input_ids, "token_type_ids": token_type_ids})` + + Args: + config (:class:`~transformers.MBartConfig`): 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 :meth:`~transformers.TFPreTrainedModel.from_pretrained` method to load the + model weights. +""" + +MBART_INPUTS_DOCSTRING = r""" + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.MBartTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`({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.html#attention-mask>`__ + decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.MBartTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + MBart uses a specific language id token as the starting token for :obj:`decoder_input_ids` generation that + varies according to source and target language, *e.g.* 25004 for `en_XX`, and 25003 for `de_DE`. If + :obj:`past_key_values` is used, optionally only the last :obj:`decoder_input_ids` have to be input (see + :obj:`past_key_values`). + + For translation and summarization training, :obj:`decoder_input_ids` should be provided. If no + :obj:`decoder_input_ids` is provided, the model will create this tensor by shifting the :obj:`input_ids` to + the right for denoising pre-training following the paper. + decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. + encoder_outputs (:obj:`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 :obj:`(batch_size, sequence_length, hidden_size)` is a sequence of + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers`) + contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). Set to :obj:`False` during training, :obj:`True` during generation + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for + more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. + training (:obj:`bool`, `optional`, defaults to :obj:`False`): + Whether or not to use the model in training mode (some modules like dropout modules have different + behaviors between training and evaluation). +""" + +MBART_GENERATION_EXAMPLE = r""" + Summarization example:: + + >>> from transformers import MBartTokenizer, TFMBartForConditionalGeneration, MBartConfig + + >>> model = MBartForConditionalGeneration.from_pretrained('facebook/mbart-large-cc25') + >>> tokenizer = MBartTokenizer.from_pretrained('facebook/mbart-large-cc25') + + >>> ARTICLE_TO_SUMMARIZE = "Meine Freunde sind cool, aber sie essen zu viel Kuchen." + >>> inputs = tokenizer([ARTICLE_TO_SUMMARIZE], max_length=1024, return_tensors='tf') + + >>> # Generate Summary + >>> summary_ids = model.generate(inputs['input_ids'], num_beams=4, max_length=5, early_stopping=True) + >>> print([tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids]) + + Mask filling example:: + + >>> from transformers import MBartTokenizer, TFMBartForConditionalGeneration + >>> tokenizer = MBartTokenizer.from_pretrained('facebook/mbart-large-cc25') + >>> # de_DE is the language symbol id for German + >>> TXT = " Meine Freunde sind nett aber sie essen zu viel Kuchen. de_DE" + + >>> model = MBartForConditionalGeneration.from_pretrained('facebook/mbart-large-cc25') + >>> input_ids = tokenizer([TXT], add_special_tokens=False, return_tensors='tf')['input_ids'] + >>> logits = model(input_ids).logits + >>> probs = tf.nn.softmax(logits[0]) + >>> # probs[5] is associated with the mask token +""" + + +@keras_serializable +class TFMBartEncoder(tf.keras.layers.Layer): + config_class = MBartConfig + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + :class:`TFMBartEncoderLayer`. + + Args: + config: MBartConfig + """ + + def __init__(self, config: MBartConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.layerdrop = config.encoder_layerdrop + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_position_embeddings + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + + self.embed_tokens = embed_tokens + self.embed_positions = TFMBartLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.layers = [TFMBartEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)] + self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") + self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm") + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + """ + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.MBartTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + inputs_embeds=inputs_embeds, + 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 input_ids and 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 input_ids or inputs_embeds") + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + embed_pos = self.embed_positions(input_shape) + hidden_states = inputs["inputs_embeds"] + embed_pos + hidden_states = self.layernorm_embedding(hidden_states) + 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] + attention_mask = _expand_mask(inputs["attention_mask"]) + else: + attention_mask = None + + encoder_states = () if inputs["output_hidden_states"] else None + all_attentions = () if inputs["output_attentions"] else None + + # encoder layers + for encoder_layer in 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, attention_mask) + + 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 TFMBartDecoder(tf.keras.layers.Layer): + config_class = MBartConfig + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a :class:`TFMBartDecoderLayer` + + Args: + config: MBartConfig + embed_tokens: output embedding + """ + + def __init__(self, config: MBartConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.padding_idx = config.pad_token_id + self.embed_tokens = embed_tokens + self.layerdrop = config.decoder_layerdrop + self.embed_positions = TFMBartLearnedPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + self.padding_idx, + name="embed_positions", + ) + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + self.layers = [TFMBartDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)] + self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") + self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm") + + self.dropout = tf.keras.layers.Dropout(config.dropout) + + def set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_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 (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.MBartTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + encoder_hidden_states (:obj:`tf.Tensor` of shape :obj:`(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 (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape :obj:`(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 :obj:`past_key_values` are used, the user can optionally input only the last + :obj:`decoder_input_ids` (those that don't have their past key value states given to this model) of + shape :obj:`(batch_size, 1)` instead of all :obj:`decoder_input_ids`` of shape :obj:`(batch_size, + sequence_length)`. + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + 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 = ( + shape_list(inputs["past_key_values"][0][0])[2] if inputs["past_key_values"] is not None else 0 + ) + + # embed positions + positions = self.embed_positions(input_shape, past_key_values_length) + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + hidden_states = inputs["inputs_embeds"] + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length) + else: + combined_attention_mask = _expand_mask( + tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] + ) + + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask( + inputs["attention_mask"], tgt_len=input_shape[-1] + ) + + 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]) + + hidden_states = self.layernorm_embedding(hidden_states + positions) + hidden_states = self.dropout(hidden_states, training=inputs["training"]) + + # decoder layers + all_hidden_states = () + all_self_attns = () + present_key_values = () + 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 + + hidden_states, layer_self_attn, present_key_value = decoder_layer( + hidden_states, + attention_mask=combined_attention_mask, + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=inputs["encoder_attention_mask"], + past_key_value=past_key_value, + ) + + if inputs["use_cache"]: + present_key_values += (present_key_value,) + + if inputs["output_attentions"]: + all_self_attns += (layer_self_attn,) + + hidden_states = self.layer_norm(hidden_states) + + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + else: + all_hidden_states = None + + all_self_attns = list(all_self_attns) if inputs["output_attentions"] else None + + present_key_values = (encoder_hidden_states, present_key_values) if inputs["use_cache"] else None + + if not inputs["return_dict"]: + return hidden_states, present_key_values, all_hidden_states, all_self_attns + else: + return TFBaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=present_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +@add_start_docstrings( + "The bare MBART Model outputting raw hidden-states without any specific head on top.", + MBART_START_DOCSTRING, +) +@keras_serializable +class TFMBartModel(TFMBartPreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config: MBartConfig, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.shared = TFSharedEmbeddings(config.vocab_size, config.d_model, config.pad_token_id, name="model.shared") + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + # Wraps layer to avoid problems with weight restoring and ensuring we're in the correct TF scope. + embed_tokens = TFWrappedEmbeddings(self.shared, abs_scope_name=shared_abs_scope_name) + embed_tokens.vocab_size = self.shared.vocab_size + embed_tokens.hidden_size = self.shared.hidden_size + + self.encoder = TFMBartEncoder(config, embed_tokens, name="encoder") + self.decoder = TFMBartDecoder(config, embed_tokens, name="decoder") + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(MBART_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + @add_code_sample_docstrings( + tokenizer_class=_TOKENIZER_FOR_DOC, + checkpoint="facebook/mbart-large-cc25", + output_type=TFSeq2SeqModelOutput, + config_class=_CONFIG_FOR_DOC, + ) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]] = None, + past_key_values=None, + inputs_embeds=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_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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 inputs["decoder_input_ids"] is None and inputs["decoder_inputs_embeds"] is None: + inputs["use_cache"] = False + + inputs["output_hidden_states"] = ( + inputs["output_hidden_states"] + if inputs["output_hidden_states"] is not None + else self.config.output_hidden_states + ) + + if inputs["decoder_input_ids"] is None and inputs["input_ids"] is not None: + inputs["decoder_input_ids"] = shift_tokens_right(inputs["input_ids"], self.config.pad_token_id) + + if inputs["encoder_outputs"] is None: + inputs["encoder_outputs"] = self.encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + inputs_embeds=inputs["inputs_embeds"], + 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() + + decoder_outputs = self.decoder( + inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=inputs["encoder_outputs"][0], + encoder_attention_mask=inputs["attention_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, + encoder_last_hidden_state=inputs["encoder_outputs"].last_hidden_state, + encoder_hidden_states=inputs["encoder_outputs"].hidden_states, + encoder_attentions=inputs["encoder_outputs"].attentions, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartModel.serving_output + 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 + 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, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + +@add_start_docstrings( + "The MBART Model with a language modeling head. Can be used for summarization.", + MBART_START_DOCSTRING, +) +class TFMBartForConditionalGeneration(TFMBartPreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + r"model.encoder.embed_tokens.weight", + r"model.decoder.embed_tokens.weight", + ] + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.model = TFMBartModel(config, name="model") + self.use_cache = config.use_cache + # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency. + self.final_logits_bias = self.add_weight( + name="final_logits_bias", shape=[1, config.vocab_size], initializer="zeros", trainable=False + ) + + def get_decoder(self): + return self.model.decoder + + def get_encoder(self): + return self.model.encoder + + def get_output_embeddings(self): + return self.get_input_embeddings() + + def set_output_embeddings(self, value): + self.set_input_embeddings(value) + + def get_bias(self): + return {"final_logits_bias": self.final_logits_bias} + + def set_bias(self, value): + self.final_logits_bias = value["final_logits_bias"] + + @add_start_docstrings_to_model_forward(MBART_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + @add_end_docstrings(MBART_GENERATION_EXAMPLE) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[TFBaseModelOutput] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + labels=None, + training=False, + **kwargs, + ): + """ + labels (:obj:`tf.Tensor` of shape :obj:`(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: + + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + labels=labels, + training=training, + kwargs_call=kwargs, + ) + + if inputs["labels"] is not None: + inputs["use_cache"] = False + if inputs["decoder_input_ids"] is None: + inputs["decoder_input_ids"] = shift_tokens_right(inputs["labels"], self.config.pad_token_id) + + outputs = self.model( + inputs["input_ids"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + encoder_outputs=inputs["encoder_outputs"], + decoder_attention_mask=inputs["decoder_attention_mask"], + past_key_values=inputs["past_key_values"], + inputs_embeds=inputs["inputs_embeds"], + 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.model.shared(outputs[0], mode="linear") + lm_logits = lm_logits + self.final_logits_bias + masked_lm_loss = None if inputs["labels"] is None else self.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, # index 1 of d outputs + decoder_hidden_states=outputs.decoder_hidden_states, # index 2 of d outputs + decoder_attentions=outputs.decoder_attentions, # index 3 of d outputs + encoder_last_hidden_state=outputs.last_hidden_state, # index 0 of encoder outputs + encoder_hidden_states=outputs.encoder_hidden_states, # 1 of e out + encoder_attentions=outputs.encoder_attentions, # 2 of e out + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.serving_output + 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 + 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 TFSeq2SeqLMOutput( + logits=output.logits, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.prepare_inputs_for_generation + def prepare_inputs_for_generation(self, decoder_input_ids, past, attention_mask, use_cache, **kwargs) -> Dict: + assert past is not None and len(past) in {1, 2}, f"past has to be an iterable of length 1,2 got {past}" + if len(past) == 1: + assert isinstance(past[0], tf.Tensor), f"`past[0]` has to be of type `tf.Tensor`, but is {type(past[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=past[0]) + past_key_values = None + else: + assert ( + len(past) == 2 + ), "`past` has to be of length 2 with the encoder_outputs at the first position and past_key_values at the second position." + encoder_outputs, past_key_values = past + if isinstance(encoder_outputs, tuple): + assert isinstance( + encoder_outputs[0], tf.Tensor + ), f"`encoder_outputs[0]` has to be of type `tf.Tensor`, but is {type(encoder_outputs[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs[0]) + elif isinstance(encoder_outputs, tf.Tensor): + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs) + assert ( + past_key_values + ), f"decoder cached states must be truthy. got {past_key_values} from the 2nd element of past" + decoder_input_ids = decoder_input_ids[:, -1:] + + assert isinstance( + encoder_outputs, TFBaseModelOutput + ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}." + return { + "input_ids": None, # encoder_outputs is defined. input_ids not needed + "encoder_outputs": encoder_outputs, + "past_key_values": past_key_values, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "use_cache": use_cache, # change this to avoid caching (presumably for debugging) + } + + @staticmethod + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration._reorder_cache + 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) + + def adjust_logits_during_generation(self, logits, cur_len, max_length): + if cur_len == max_length - 1: + vocab_range = tf.constant(range(self.config.vocab_size)) + return tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits) + else: + return logits + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.compute_loss + def compute_loss(self, labels, logits): + """CrossEntropyLoss that ignores pad tokens""" + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True, + reduction=tf.keras.losses.Reduction.NONE, + ) + melted_labels = tf.reshape(labels, (-1,)) + active_loss = tf.not_equal(melted_labels, self.config.pad_token_id) + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(melted_labels, active_loss) + return loss_fn(labels, reduced_logits) diff --git a/src/transformers/models/pegasus/__init__.py b/src/transformers/models/pegasus/__init__.py index 4041e0db7d..e263c1b86c 100644 --- a/src/transformers/models/pegasus/__init__.py +++ b/src/transformers/models/pegasus/__init__.py @@ -45,7 +45,7 @@ if is_torch_available(): ] if is_tf_available(): - _import_structure["modeling_tf_pegasus"] = ["TFPegasusForConditionalGeneration"] + _import_structure["modeling_tf_pegasus"] = ["TFPegasusForConditionalGeneration", "TFPegasusModel"] if TYPE_CHECKING: @@ -66,7 +66,7 @@ if TYPE_CHECKING: ) if is_tf_available(): - from .modeling_tf_pegasus import TFPegasusForConditionalGeneration + from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel else: import importlib diff --git a/src/transformers/models/pegasus/configuration_pegasus.py b/src/transformers/models/pegasus/configuration_pegasus.py index 0dee48e7bb..0ed78b25fa 100644 --- a/src/transformers/models/pegasus/configuration_pegasus.py +++ b/src/transformers/models/pegasus/configuration_pegasus.py @@ -159,17 +159,6 @@ class PegasusConfig(PretrainedConfig): self.gradient_checkpointing = gradient_checkpointing self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True - # IMPORTANT - # DELETE ALL OF THE FOLLOWING LINES AS SOON AS TF IS READY - self.extra_pos_embeddings = 0 - self.normalize_before = True - self.add_final_layer_norm = True - self.do_blenderbot_90_layernorm = False - self.normalize_embedding = False - self.static_position_embeddings = True - self.add_bias_logits = False - self.force_bos_token_to_be_generated = False - @property def num_attention_heads(self) -> int: return self.encoder_attention_heads diff --git a/src/transformers/models/pegasus/modeling_tf_pegasus.py b/src/transformers/models/pegasus/modeling_tf_pegasus.py index bec856575d..1f41d02790 100644 --- a/src/transformers/models/pegasus/modeling_tf_pegasus.py +++ b/src/transformers/models/pegasus/modeling_tf_pegasus.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team. +# Copyright 2021, Google Inc. 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. @@ -12,30 +12,1312 @@ # 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. -"""TF Pegasus model, ported from the fairseq repo.""" -from ...file_utils import add_start_docstrings +""" TF 2.0 Pegasus model. """ + + +import random +from typing import Dict, Optional, Tuple, Union + +import numpy as np +import tensorflow as tf + +from ...activations_tf import get_tf_activation +from ...file_utils import ( + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + replace_return_docstrings, +) +from ...modeling_tf_outputs import ( + TFBaseModelOutput, + TFBaseModelOutputWithPast, + TFSeq2SeqLMOutput, + TFSeq2SeqModelOutput, +) + +# Public API +from ...modeling_tf_utils import ( + DUMMY_INPUTS, + TFPreTrainedModel, + TFSharedEmbeddings, + TFWrappedEmbeddings, + input_processing, + keras_serializable, + shape_list, +) from ...utils import logging -from ..bart.modeling_tf_bart import BART_START_DOCSTRING, TFBartForConditionalGeneration from .configuration_pegasus import PegasusConfig -_CONFIG_FOR_DOC = "PegasusConfig" - -START_DOCSTRING = BART_START_DOCSTRING.replace( - "inherits from :class:`~transformers.TFPreTrainedModel`", - "inherits from :class:`~transformers.TFBartForConditionalGeneration`", -).replace("BartConfig", _CONFIG_FOR_DOC) - - logger = logging.get_logger(__name__) +_CONFIG_FOR_DOC = "PegasusConfig" +_TOKENIZER_FOR_DOC = "PegasusTokenizer" -@add_start_docstrings("Pegasus model for summarization", START_DOCSTRING) -class TFPegasusForConditionalGeneration(TFBartForConditionalGeneration): - _keys_to_ignore_on_load_missing = [ - r"final_logits_bias", - r"model.encoder.embed_positions.weight", - r"model.decoder.embed_positions.weight", - ] + +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): + shifted_input_ids = tf.cast(input_ids, tf.int32) + shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1) + start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), decoder_start_token_id) + shifted_input_ids = tf.concat([start_tokens, shifted_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 + ) + + # "Verify that `labels` has only positive values and -100" + assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, tf.int32)) + + # 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), dtype=tf.float32) * 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) + mask = tf.cast(mask, tf.float32) + + if past_key_values_length > 0: + mask = tf.concat([tf.zeros((tgt_len, past_key_values_length), dtype=tf.float32), mask], axis=-1) + return tf.broadcast_to(mask[None, None, :, :], (bsz, 1, tgt_len, tgt_len + past_key_values_length)) + + +# 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]`. + """ + bsz, src_len = shape_list(mask) + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = tf.cast(tf.broadcast_to(mask[:, None, None, :], (bsz, 1, tgt_len, src_len)), tf.float32) + + return (1.0 - expanded_mask) * LARGE_NEGATIVE + + +# Copied from transformers.models.marian.modeling_tf_marian.TFMarianSinusoidalPositionalEmbedding with Marian->Pegasus +class TFPegasusSinusoidalPositionalEmbedding(tf.keras.layers.Embedding): + """This module produces sinusoidal positional embeddings of any length.""" + + def __init__(self, num_positions: int, embedding_dim: int, **kwargs): + + if embedding_dim % 2 != 0: + raise NotImplementedError(f"odd embedding_dim {embedding_dim} not supported") + super().__init__( + num_positions, + embedding_dim, + **kwargs, + ) + + 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 + """ + super().build(input_shape) # Instantiates self.weight so it can be loaded + weight: np.ndarray = self._init_weight(self.input_dim, self.output_dim) + self.set_weights([weight]) # overwrite self.weight to correct value + + @staticmethod + def _init_weight(n_pos: int, dim: int): + """ + Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in + the 2nd half of the vector. [dim // 2:] + """ + position_enc = np.array( + [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)] + ) + # index 0 is all zero + position_enc[:, 0 : dim // 2] = np.sin(position_enc[:, 0::2]) + position_enc[:, dim // 2 :] = np.cos(position_enc[:, 1::2]) + # convert to tensor + table = tf.convert_to_tensor(position_enc, dtype=tf.float32) + tf.stop_gradient(table) + return table + + def call(self, input_shape: tf.TensorShape, past_key_values_length: int = 0): + """Input is expected to be of size [bsz x seqlen].""" + bsz, seq_len = input_shape[:2] + + positions = tf.range( + past_key_values_length, seq_len + past_key_values_length, delta=1, dtype=tf.int32, name="range" + ) + return super().call(positions) + + +# Copied from transformers.models.bart.modeling_tf_bart.TFBartAttention with Bart->Pegasus +class TFPegasusAttention(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 + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by 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, + 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) + + 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: + 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)}", + ) + 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) + + attn_probs = self.dropout(attn_weights, training=training) + + attn_output = tf.matmul(attn_probs, value_states) + + 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 + + +# Copied from transformers.models.mbart.modeling_tf_mbart.TFMBartEncoderLayer with MBart->Pegasus +class TFPegasusEncoderLayer(tf.keras.layers.Layer): + def __init__(self, config: PegasusConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFPegasusAttention( + 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, training=False): + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`tf.Tensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + """ + 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 + ) + 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 + + +# Copied from transformers.models.mbart.modeling_tf_mbart.TFMBartDecoderLayer with MBart->Pegasus +class TFPegasusDecoderLayer(tf.keras.layers.Layer): + def __init__(self, config: PegasusConfig, **kwargs): + super().__init__(**kwargs) + self.embed_dim = config.d_model + self.self_attn = TFPegasusAttention( + 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 = TFPegasusAttention( + 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, + past_key_value: Optional[Tuple[tf.Tensor]] = None, + training=False, + ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]: + """ + Args: + hidden_states (:obj:`tf.Tensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (:obj:`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 (:obj:`tf.Tensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (:obj:`tf.Tensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + past_key_value (:obj:`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, + ) + hidden_states = self.dropout(hidden_states, training=training) + hidden_states = residual + hidden_states + + # Cross-Attention Block + cross_attn_present_key_value = 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_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_value=cross_attn_past_key_value, + ) + 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, + present_key_value, + ) + + +class TFPegasusPreTrainedModel(TFPreTrainedModel): config_class = PegasusConfig - # All the code is in src/transformers/models/bart/modeling_tf_bart.py + base_model_prefix = "model" + + @property + def dummy_inputs(self): + pad_token = 1 + input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + decoder_input_ids = tf.cast(tf.convert_to_tensor(DUMMY_INPUTS), tf.int32) + dummy_inputs = { + "decoder_input_ids": decoder_input_ids, + "attention_mask": tf.math.not_equal(input_ids, pad_token), + "input_ids": input_ids, + } + return dummy_inputs + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.get_input_embeddings + def get_input_embeddings(self): + base_model = getattr(self, self.base_model_prefix, self) + + return base_model.shared + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.set_input_embeddings + def set_input_embeddings(self, value): + base_model = getattr(self, self.base_model_prefix, self) + + try: + base_model.shared.weight = value + except AttributeError: + self(self.dummy_inputs) + base_model.shared.weight = value + + base_model.shared.vocab_size = shape_list(base_model.shared.weight)[0] + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + embed_tokens = TFWrappedEmbeddings(base_model.shared, abs_scope_name=shared_abs_scope_name) + base_model.encoder.set_embed_tokens(embed_tokens) + base_model.decoder.set_embed_tokens(embed_tokens) + + @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"), + } + ] + ) + # Copied from transformers.models.bart.modeling_tf_bart.TFBartPretrainedModel.serving + def serving(self, inputs): + output = self.call(inputs) + + return self.serving_output(output) + + +PEGASUS_START_DOCSTRING = r""" + This model inherits from :class:`~transformers.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 `__ 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. + + .. note:: + + 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 :meth:`tf.keras.Model.fit` method which currently requires having all + the tensors in the first argument of the model call function: :obj:`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 :obj:`input_ids` only and nothing else: :obj:`model(input_ids)` + - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: + :obj:`model([input_ids, attention_mask])` or :obj:`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: + :obj:`model({"input_ids": input_ids, "token_type_ids": token_type_ids})` + + Args: + config (:class:`~transformers.PegasusConfig`): 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 :meth:`~transformers.TFPreTrainedModel.from_pretrained` method to load the + model weights. +""" + +PEGASUS_GENERATION_EXAMPLE = r""" + Summarization example:: + + >>> from transformers import PegasusTokenizer, TFPegasusForConditionalGeneration + + >>> model = TFPegasusForConditionalGeneration.from_pretrained('google/pegasus-xsum') + >>> tokenizer = PegasusTokenizer.from_pretrained('google/pegasus-xsum') + + >>> ARTICLE_TO_SUMMARIZE = ( + ... "PG&E stated it scheduled the blackouts in response to forecasts for high winds " + ... "amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were " + ... "scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow." + ... ) + >>> inputs = tokenizer([ARTICLE_TO_SUMMARIZE], max_length=1024, return_tensors='tf') + + >>> # Generate Summary + >>> summary_ids = model.generate(inputs['input_ids']) + >>> print([tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids]) +""" + +PEGASUS_INPUTS_DOCSTRING = r""" + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.PegasusTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`({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.html#attention-mask>`__ + decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.PegasusTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + Pegasus uses the :obj:`pad_token_id` as the starting token for :obj:`decoder_input_ids` generation. If + :obj:`past_key_values` is used, optionally only the last :obj:`decoder_input_ids` have to be input (see + :obj:`past_key_values`). + + For translation and summarization training, :obj:`decoder_input_ids` should be provided. If no + :obj:`decoder_input_ids` is provided, the model will create this tensor by shifting the :obj:`input_ids` to + the right for denoising pre-training following the paper. + decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. + encoder_outputs (:obj:`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 :obj:`(batch_size, sequence_length, hidden_size)` is a sequence of + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers`) + contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids` + (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)` + instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up + decoding (see :obj:`past_key_values`). Set to :obj:`False` during training, :obj:`True` during generation + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for + more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. + training (:obj:`bool`, `optional`, defaults to :obj:`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 TFPegasusEncoder(tf.keras.layers.Layer): + config_class = PegasusConfig + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + :class:`TFPegasusEncoderLayer`. + + Args: + config: PegasusConfig + """ + + def __init__(self, config: PegasusConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.dropout = tf.keras.layers.Dropout(config.dropout) + self.layerdrop = config.encoder_layerdrop + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_position_embeddings + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + + self.embed_tokens = embed_tokens + self.embed_positions = TFPegasusSinusoidalPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + name="embed_positions", + ) + self.layers = [TFPegasusEncoderLayer(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 set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs, + ): + """ + Args: + input_ids (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.PegasusTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + inputs_embeds=inputs_embeds, + 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 input_ids and 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 input_ids or inputs_embeds") + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + embed_pos = self.embed_positions(input_shape) + hidden_states = inputs["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] + attention_mask = _expand_mask(inputs["attention_mask"]) + else: + attention_mask = None + + encoder_states = () if inputs["output_hidden_states"] else None + all_attentions = () if inputs["output_attentions"] else None + + # encoder layers + for encoder_layer in 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, attention_mask) + + 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 TFPegasusDecoder(tf.keras.layers.Layer): + config_class = PegasusConfig + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a :class:`TFPegasusDecoderLayer` + + Args: + config: PegasusConfig + embed_tokens: output embedding + """ + + def __init__(self, config: PegasusConfig, embed_tokens: Optional[TFSharedEmbeddings] = None, **kwargs): + super().__init__(**kwargs) + self.config = config + self.padding_idx = config.pad_token_id + self.embed_tokens = embed_tokens + self.layerdrop = config.decoder_layerdrop + self.embed_positions = TFPegasusSinusoidalPositionalEmbedding( + config.max_position_embeddings, + config.d_model, + name="embed_positions", + ) + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 + self.layers = [TFPegasusDecoderLayer(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 set_embed_tokens(self, embed_tokens): + self.embed_tokens = embed_tokens + + def call( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_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 (:obj:`tf.Tensor` of shape :obj:`(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 :class:`~transformers.PegasusTokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` + for details. + + `What are input IDs? <../glossary.html#input-ids>`__ + attention_mask (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + encoder_hidden_states (:obj:`tf.Tensor` of shape :obj:`(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 (:obj:`tf.Tensor` of shape :obj:`(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.html#attention-mask>`__ + past_key_values (:obj:`Tuple[Tuple[tf.Tensor]]` of length :obj:`config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape :obj:`(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 :obj:`past_key_values` are used, the user can optionally input only the last + :obj:`decoder_input_ids` (those that don't have their past key value states given to this model) of + shape :obj:`(batch_size, 1)` instead of all :obj:`decoder_input_ids`` of shape :obj:`(batch_size, + sequence_length)`. + inputs_embeds (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): + Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded + representation. This is useful if you want more control over how to convert :obj:`input_ids` indices + into associated vectors than the model's internal embedding lookup matrix. + output_attentions (:obj:`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 (:obj:`bool`, `optional`): + Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors + for more detail. + return_dict (:obj:`bool`, `optional`): + Whether or not to return a :class:`~transformers.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, + 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 = ( + shape_list(inputs["past_key_values"][0][0])[2] if inputs["past_key_values"] is not None else 0 + ) + + # embed positions + positions = self.embed_positions(input_shape, past_key_values_length) + + if inputs["inputs_embeds"] is None: + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale + + hidden_states = inputs["inputs_embeds"] + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length) + else: + combined_attention_mask = _expand_mask( + tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] + ) + + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask( + inputs["attention_mask"], tgt_len=input_shape[-1] + ) + + 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]) + + hidden_states = self.dropout(hidden_states + positions, training=inputs["training"]) + + # decoder layers + all_hidden_states = () + all_self_attns = () + present_key_values = () + 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 + + hidden_states, layer_self_attn, present_key_value = decoder_layer( + hidden_states, + attention_mask=combined_attention_mask, + encoder_hidden_states=inputs["encoder_hidden_states"], + encoder_attention_mask=inputs["encoder_attention_mask"], + past_key_value=past_key_value, + ) + + if inputs["use_cache"]: + present_key_values += (present_key_value,) + + if inputs["output_attentions"]: + all_self_attns += (layer_self_attn,) + + hidden_states = self.layer_norm(hidden_states) + + if inputs["output_hidden_states"]: + all_hidden_states += (hidden_states,) + else: + all_hidden_states = None + + all_self_attns = list(all_self_attns) if inputs["output_attentions"] else None + + present_key_values = (encoder_hidden_states, present_key_values) if inputs["use_cache"] else None + + if not inputs["return_dict"]: + return hidden_states, present_key_values, all_hidden_states, all_self_attns + else: + return TFBaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=present_key_values, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +@add_start_docstrings( + "The bare PEGASUS Model outputting raw hidden-states without any specific head on top.", + PEGASUS_START_DOCSTRING, +) +@keras_serializable +class TFPegasusModel(TFPegasusPreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config: PegasusConfig, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.shared = TFSharedEmbeddings(config.vocab_size, config.d_model, config.pad_token_id, name="model.shared") + + with tf.compat.v1.variable_scope("model.shared") as shared_abs_scope_name: + pass + + # Wraps layer to avoid problems with weight restoring and ensuring we're in the correct TF scope. + embed_tokens = TFWrappedEmbeddings(self.shared, abs_scope_name=shared_abs_scope_name) + embed_tokens.vocab_size = self.shared.vocab_size + embed_tokens.hidden_size = self.shared.hidden_size + + self.encoder = TFPegasusEncoder(config, embed_tokens, name="encoder") + self.decoder = TFPegasusDecoder(config, embed_tokens, name="decoder") + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(PEGASUS_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + @replace_return_docstrings(output_type=TFSeq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + training=False, + **kwargs + ): + r""" + Returns: + + Example:: + + >>> from transformers import PegasusTokenizer, TFPegasusModel + + >>> tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-large") + >>> model = TFPegasusModel.from_pretrained("google/pegasus-large") + + >>> input_ids = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="tf").input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="tf").input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + + >>> last_hidden_states = outputs.last_hidden_state + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + 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 inputs["decoder_input_ids"] is None and inputs["decoder_inputs_embeds"] is None: + inputs["use_cache"] = False + + inputs["output_hidden_states"] = ( + inputs["output_hidden_states"] + if inputs["output_hidden_states"] is not None + else self.config.output_hidden_states + ) + + if inputs["encoder_outputs"] is None: + inputs["encoder_outputs"] = self.encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + inputs_embeds=inputs["inputs_embeds"], + 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() + + decoder_outputs = self.decoder( + inputs["decoder_input_ids"], + attention_mask=inputs["decoder_attention_mask"], + encoder_hidden_states=inputs["encoder_outputs"][0], + encoder_attention_mask=inputs["attention_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, + encoder_last_hidden_state=inputs["encoder_outputs"].last_hidden_state, + encoder_hidden_states=inputs["encoder_outputs"].hidden_states, + encoder_attentions=inputs["encoder_outputs"].attentions, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartModel.serving_output + 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 + 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, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + +@add_start_docstrings( + "The PEGASUS Model with a language modeling head. Can be used for summarization.", + PEGASUS_START_DOCSTRING, +) +class TFPegasusForConditionalGeneration(TFPegasusPreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + r"model.encoder.embed_tokens.weight", + r"model.decoder.embed_tokens.weight", + ] + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + self.model = TFPegasusModel(config, name="model") + self.use_cache = config.use_cache + # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency. + self.final_logits_bias = self.add_weight( + name="final_logits_bias", shape=[1, config.vocab_size], initializer="zeros", trainable=False + ) + + def get_decoder(self): + return self.model.decoder + + def get_encoder(self): + return self.model.encoder + + def get_output_embeddings(self): + return self.get_input_embeddings() + + def set_output_embeddings(self, value): + self.set_input_embeddings(value) + + def get_bias(self): + return {"final_logits_bias": self.final_logits_bias} + + def set_bias(self, value): + self.final_logits_bias = value["final_logits_bias"] + + @add_start_docstrings_to_model_forward(PEGASUS_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) + @add_end_docstrings(PEGASUS_GENERATION_EXAMPLE) + def call( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + encoder_outputs: Optional[TFBaseModelOutput] = None, + past_key_values=None, + inputs_embeds=None, + decoder_inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + labels=None, + training=False, + **kwargs, + ): + """ + labels (:obj:`torch.LongTensor` of shape :obj:`(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: + + """ + inputs = input_processing( + func=self.call, + config=self.config, + input_ids=input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + labels=labels, + training=training, + kwargs_call=kwargs, + ) + + if inputs["labels"] is not None: + inputs["use_cache"] = False + 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( + inputs["input_ids"], + attention_mask=inputs["attention_mask"], + decoder_input_ids=inputs["decoder_input_ids"], + encoder_outputs=inputs["encoder_outputs"], + decoder_attention_mask=inputs["decoder_attention_mask"], + past_key_values=inputs["past_key_values"], + inputs_embeds=inputs["inputs_embeds"], + 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.model.shared(outputs[0], mode="linear") + lm_logits = lm_logits + self.final_logits_bias + masked_lm_loss = None if inputs["labels"] is None else self.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, # index 1 of d outputs + decoder_hidden_states=outputs.decoder_hidden_states, # index 2 of d outputs + decoder_attentions=outputs.decoder_attentions, # index 3 of d outputs + encoder_last_hidden_state=outputs.last_hidden_state, # index 0 of encoder outputs + encoder_hidden_states=outputs.encoder_hidden_states, # 1 of e out + encoder_attentions=outputs.encoder_attentions, # 2 of e out + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.serving_output + 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 + 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 TFSeq2SeqLMOutput( + logits=output.logits, + past_key_values=pkv, + decoder_hidden_states=dec_hs, + decoder_attentions=dec_attns, + encoder_last_hidden_state=output.encoder_last_hidden_state, + encoder_hidden_states=enc_hs, + encoder_attentions=enc_attns, + ) + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.prepare_inputs_for_generation + def prepare_inputs_for_generation(self, decoder_input_ids, past, attention_mask, use_cache, **kwargs) -> Dict: + assert past is not None and len(past) in {1, 2}, f"past has to be an iterable of length 1,2 got {past}" + if len(past) == 1: + assert isinstance(past[0], tf.Tensor), f"`past[0]` has to be of type `tf.Tensor`, but is {type(past[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=past[0]) + past_key_values = None + else: + assert ( + len(past) == 2 + ), "`past` has to be of length 2 with the encoder_outputs at the first position and past_key_values at the second position." + encoder_outputs, past_key_values = past + if isinstance(encoder_outputs, tuple): + assert isinstance( + encoder_outputs[0], tf.Tensor + ), f"`encoder_outputs[0]` has to be of type `tf.Tensor`, but is {type(encoder_outputs[0])}" + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs[0]) + elif isinstance(encoder_outputs, tf.Tensor): + encoder_outputs = TFBaseModelOutput(last_hidden_state=encoder_outputs) + assert ( + past_key_values + ), f"decoder cached states must be truthy. got {past_key_values} from the 2nd element of past" + decoder_input_ids = decoder_input_ids[:, -1:] + + assert isinstance( + encoder_outputs, TFBaseModelOutput + ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}." + return { + "input_ids": None, # encoder_outputs is defined. input_ids not needed + "encoder_outputs": encoder_outputs, + "past_key_values": past_key_values, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "use_cache": use_cache, # change this to avoid caching (presumably for debugging) + } + + @staticmethod + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration._reorder_cache + 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) + + def adjust_logits_during_generation(self, logits, cur_len, max_length): + if cur_len == max_length - 1: + vocab_range = tf.constant(range(self.config.vocab_size)) + return tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits) + else: + return logits + + # Copied from transformers.models.bart.modeling_tf_bart.TFBartForConditionalGeneration.compute_loss + def compute_loss(self, labels, logits): + """CrossEntropyLoss that ignores pad tokens""" + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True, + reduction=tf.keras.losses.Reduction.NONE, + ) + melted_labels = tf.reshape(labels, (-1,)) + active_loss = tf.not_equal(melted_labels, self.config.pad_token_id) + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(melted_labels, active_loss) + return loss_fn(labels, reduced_logits) diff --git a/src/transformers/utils/dummy_tf_objects.py b/src/transformers/utils/dummy_tf_objects.py index 2a2b8f27cf..1413ae2cce 100644 --- a/src/transformers/utils/dummy_tf_objects.py +++ b/src/transformers/utils/dummy_tf_objects.py @@ -369,6 +369,33 @@ class TFBlenderbotForConditionalGeneration: requires_tf(self) +class TFBlenderbotModel: + def __init__(self, *args, **kwargs): + requires_tf(self) + + @classmethod + def from_pretrained(self, *args, **kwargs): + requires_tf(self) + + +class TFBlenderbotSmallForConditionalGeneration: + def __init__(self, *args, **kwargs): + requires_tf(self) + + @classmethod + def from_pretrained(self, *args, **kwargs): + requires_tf(self) + + +class TFBlenderbotSmallModel: + def __init__(self, *args, **kwargs): + requires_tf(self) + + @classmethod + def from_pretrained(self, *args, **kwargs): + requires_tf(self) + + TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST = None @@ -952,6 +979,11 @@ class TFLxmertVisualFeatureEncoder: requires_tf(self) +class TFMarian: + def __init__(self, *args, **kwargs): + requires_tf(self) + + class TFMarianMTModel: def __init__(self, *args, **kwargs): requires_tf(self) @@ -970,6 +1002,15 @@ class TFMBartForConditionalGeneration: requires_tf(self) +class TFMBartModel: + def __init__(self, *args, **kwargs): + requires_tf(self) + + @classmethod + def from_pretrained(self, *args, **kwargs): + requires_tf(self) + + TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST = None @@ -1211,6 +1252,15 @@ class TFPegasusForConditionalGeneration: requires_tf(self) +class TFPegasusModel: + def __init__(self, *args, **kwargs): + requires_tf(self) + + @classmethod + def from_pretrained(self, *args, **kwargs): + requires_tf(self) + + TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = None diff --git a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_tf_{{cookiecutter.lowercase_modelname}}.py b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_tf_{{cookiecutter.lowercase_modelname}}.py index 332fc79880..3526d1bfde 100644 --- a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_tf_{{cookiecutter.lowercase_modelname}}.py +++ b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_tf_{{cookiecutter.lowercase_modelname}}.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright {{cookiecutter.authors}} and The HuggingFace Inc. team. All rights reserved. +# Copyright 2021 {{cookiecutter.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. @@ -1464,7 +1464,6 @@ class TF{{cookiecutter.camelcase_modelname}}ForQuestionAnswering(TF{{cookiecutte ) {% else %} -import math import random from typing import Dict, Optional, Tuple, Union @@ -1936,7 +1935,7 @@ class TF{{cookiecutter.camelcase_modelname}}PreTrainedModel(TFPreTrainedModel): input_ids (:obj:`tf.Tensor` of shape :obj:`({0})`): Indices of input sequence tokens in the vocabulary. - Indices can be obtained using :class:`~transformers.BertTokenizer`. See + Indices can be obtained using :class:`~transformers.{{cookiecutter.camelcase_modelname}}Tokenizer`. See :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for details. @@ -1949,8 +1948,21 @@ class TF{{cookiecutter.camelcase_modelname}}PreTrainedModel(TFPreTrainedModel): `What are attention masks? <../glossary.html#attention-mask>`__ decoder_input_ids (:obj:`tf.Tensor` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): - Provide for translation and summarization training. By default, the model will create this tensor by - shifting the input_ids right, following the paper. + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using :class:`~transformers.{{cookiecutter.camelcase_modelname}}Tokenizer`. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + + `What are input IDs? <../glossary.html#input-ids>`__ + + {{cookiecutter.camelcase_modelname}} uses the :obj:`eos_token_id` as the starting token for + :obj:`decoder_input_ids` generation. If :obj:`past_key_values` is used, optionally only the last + :obj:`decoder_input_ids` have to be input (see :obj:`past_key_values`). + + For translation and summarization training, :obj:`decoder_input_ids` should be provided. If no + :obj:`decoder_input_ids` is provided, the model will create this tensor by shifting the :obj:`input_ids` to + the right for denoising pre-training following the paper. decoder_attention_mask (:obj:`tf.Tensor` of shape :obj:`(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. encoder_outputs (:obj:`tf.FloatTensor`, `optional`): @@ -1971,7 +1983,7 @@ class TF{{cookiecutter.camelcase_modelname}}PreTrainedModel(TFPreTrainedModel): Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for more detail. return_dict (:obj:`bool`, `optional`): - Whether or not to return a :class:`~transformers.file_utils.TFModelOutput` instead of a plain tuple. + Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. training (:obj:`bool`, `optional`, defaults to :obj:`False`): Whether or not to use the model in training mode (some modules like dropout modules have different behaviors between training and evaluation). @@ -1996,7 +2008,7 @@ class TF{{cookiecutter.camelcase_modelname}}Encoder(tf.keras.layers.Layer): self.layerdrop = config.encoder_layerdrop self.padding_idx = config.pad_token_id self.max_source_positions = config.max_position_embeddings - self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 self.embed_tokens = embed_tokens @@ -2077,14 +2089,10 @@ class TF{{cookiecutter.camelcase_modelname}}Encoder(tf.keras.layers.Layer): raise ValueError("You have to specify either input_ids or inputs_embeds") if inputs["inputs_embeds"] is None: - inputs_embeds = self.embed_tokens(inputs["input_ids"]) * self.embed_scale - else: - inputs_embeds = inputs["inputs_embeds"] - - inputs_embeds = inputs_embeds + inputs["inputs_embeds"] = self.embed_tokens(inputs["input_ids"]) * self.embed_scale embed_pos = self.embed_positions(input_shape) - hidden_states = inputs_embeds + embed_pos + hidden_states = inputs["inputs_embeds"] + embed_pos hidden_states = self.layernorm_embedding(hidden_states) hidden_states = self.dropout(hidden_states, training=inputs["training"]) @@ -2146,7 +2154,7 @@ class TF{{cookiecutter.camelcase_modelname}}Decoder(tf.keras.layers.Layer): self.padding_idx, name="embed_positions", ) - self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 + self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0 self.layers = [TF{{cookiecutter.camelcase_modelname}}DecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)] self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding") @@ -2259,7 +2267,6 @@ class TF{{cookiecutter.camelcase_modelname}}Decoder(tf.keras.layers.Layer): hidden_states = inputs["inputs_embeds"] # [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) else: @@ -2267,21 +2274,8 @@ class TF{{cookiecutter.camelcase_modelname}}Decoder(tf.keras.layers.Layer): tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1] ) - if inputs["attention_mask"] is None and inputs["input_ids"] is not None and input_shape[-1] > 1: - inputs["attention_mask"] = tf.cast( - tf.math.not_equal(inputs["input_ids"], self.config.pad_token_id), inputs["input_ids"].dtype - ) - inputs["attention_mask"] = tf.concat( - [ - tf.ones((input_shape[0], past_key_values_length), dtype=inputs["attention_mask"].dtype), - inputs["attention_mask"], - ], - axis=-1, - ) - else: - inputs["attention_mask"] = tf.ones( - (input_shape[0], input_shape[1] + past_key_values_length), dtype=tf.int32 - ) + if inputs["attention_mask"] is not None and input_shape[-1] > 1: + combined_attention_mask = combined_attention_mask + _expand_mask(inputs["attention_mask"], tgt_len=input_shape[-1]) 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] @@ -2683,7 +2677,7 @@ class TF{{cookiecutter.camelcase_modelname}}ForConditionalGeneration(TF{{cookiec 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), + 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/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_{{cookiecutter.lowercase_modelname}}.py b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_{{cookiecutter.lowercase_modelname}}.py index c3b02e905d..5a12c9b795 100755 --- a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_{{cookiecutter.lowercase_modelname}}.py +++ b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_{{cookiecutter.lowercase_modelname}}.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright {{cookiecutter.authors}} and The HuggingFace Inc. team. All rights reserved. +# Copyright 2021 {{cookiecutter.authors}} 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. diff --git a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_tf_{{cookiecutter.lowercase_modelname}}.py b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_tf_{{cookiecutter.lowercase_modelname}}.py index 7ccd691ca6..c9637cd607 100644 --- a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_tf_{{cookiecutter.lowercase_modelname}}.py +++ b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_tf_{{cookiecutter.lowercase_modelname}}.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright {{cookiecutter.authors}} and The HuggingFace Inc. team. All rights reserved. +# 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. @@ -17,7 +17,7 @@ import unittest -from transformers import {{cookiecutter.camelcase_modelname}}Config, is_tf_available +from transformers import is_tf_available, {{cookiecutter.camelcase_modelname}}Config from transformers.testing_utils import require_tf, slow from .test_configuration_common import ConfigTester @@ -28,12 +28,12 @@ if is_tf_available(): import tensorflow as tf from transformers import ( + TF{{cookiecutter.camelcase_modelname}}ForCausalLM, TF{{cookiecutter.camelcase_modelname}}ForMaskedLM, TF{{cookiecutter.camelcase_modelname}}ForMultipleChoice, TF{{cookiecutter.camelcase_modelname}}ForQuestionAnswering, TF{{cookiecutter.camelcase_modelname}}ForSequenceClassification, TF{{cookiecutter.camelcase_modelname}}ForTokenClassification, - TF{{cookiecutter.camelcase_modelname}}ForCausalLM, TF{{cookiecutter.camelcase_modelname}}Model, ) @@ -323,8 +323,12 @@ class TF{{cookiecutter.camelcase_modelname}}ModelIntegrationTest(unittest.TestCa {% else %} import unittest -from transformers import {{cookiecutter.camelcase_modelname}}Config, {{cookiecutter.camelcase_modelname}}Tokenizer, is_tf_available -from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_tf, slow +from transformers import ( + is_tf_available, + {{cookiecutter.camelcase_modelname}}Config, + {{cookiecutter.camelcase_modelname}}Tokenizer, +) +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, ids_tensor @@ -333,7 +337,10 @@ from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor if is_tf_available(): import tensorflow as tf - from transformers import TF{{cookiecutter.camelcase_modelname}}ForConditionalGeneration, TF{{cookiecutter.camelcase_modelname}}Model + from transformers import ( + TF{{cookiecutter.camelcase_modelname}}ForConditionalGeneration, + TF{{cookiecutter.camelcase_modelname}}Model, + ) @require_tf @@ -453,7 +460,7 @@ def prepare_{{cookiecutter.lowercase_modelname}}_inputs_dict( if attention_mask is None: attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) if decoder_attention_mask is None: - decoder_attention_mask = tf.cast(tf.math.not_equal(decoder_input_ids, config.pad_token_id), tf.int8) + decoder_attention_mask = tf.concat([tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8)], axis=-1) return { "input_ids": input_ids, "decoder_input_ids": decoder_input_ids, diff --git a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_{{cookiecutter.lowercase_modelname}}.py b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_{{cookiecutter.lowercase_modelname}}.py index d5ee422b12..9d7d303df7 100644 --- a/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_{{cookiecutter.lowercase_modelname}}.py +++ b/templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/test_modeling_{{cookiecutter.lowercase_modelname}}.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright {{cookiecutter.authors}} and The HuggingFace Inc. team. All rights reserved. +# 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. @@ -21,8 +21,8 @@ import unittest from tests.test_modeling_common import floats_tensor from transformers import is_torch_available from transformers.testing_utils import require_torch, slow, torch_device -from .test_configuration_common import ConfigTester +from .test_configuration_common import ConfigTester from .test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask @@ -31,15 +31,17 @@ if is_torch_available(): from transformers import ( {{cookiecutter.camelcase_modelname}}Config, - {{cookiecutter.camelcase_modelname}}ForMaskedLM, {{cookiecutter.camelcase_modelname}}ForCausalLM, + {{cookiecutter.camelcase_modelname}}ForMaskedLM, {{cookiecutter.camelcase_modelname}}ForMultipleChoice, {{cookiecutter.camelcase_modelname}}ForQuestionAnswering, {{cookiecutter.camelcase_modelname}}ForSequenceClassification, {{cookiecutter.camelcase_modelname}}ForTokenClassification, {{cookiecutter.camelcase_modelname}}Model, ) - from transformers.models.{{cookiecutter.lowercase_modelname}}.modeling_{{cookiecutter.lowercase_modelname}} import {{cookiecutter.uppercase_modelname}}_PRETRAINED_MODEL_ARCHIVE_LIST + from transformers.models.{{cookiecutter.lowercase_modelname}}.modeling_{{cookiecutter.lowercase_modelname}} import ( + {{cookiecutter.uppercase_modelname}}_PRETRAINED_MODEL_ARCHIVE_LIST, + ) class {{cookiecutter.camelcase_modelname}}ModelTester: diff --git a/tests/test_modeling_tf_bart.py b/tests/test_modeling_tf_bart.py index 2765eb7a41..9520aea57e 100644 --- a/tests/test_modeling_tf_bart.py +++ b/tests/test_modeling_tf_bart.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 The Huggingface Inc. team +# 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. @@ -13,14 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. - import unittest import numpy as np from transformers import BartConfig, BartTokenizer, is_tf_available from transformers.file_utils import cached_property -from transformers.testing_utils import is_pt_tf_cross_test, require_tf, slow +from transformers.testing_utils import require_tf, slow from .test_configuration_common import ConfigTester from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor @@ -30,7 +29,6 @@ if is_tf_available(): import tensorflow as tf from transformers import TFBartForConditionalGeneration, TFBartModel - from transformers.models.bart.modeling_tf_bart import TFBartSinusoidalPositionalEmbedding @require_tf @@ -39,30 +37,49 @@ class TFBartModelTester: config_updates = {} hidden_act = "gelu" - def __init__(self, parent): + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=20, + eos_token_id=2, + pad_token_id=1, + bos_token_id=0, + ): self.parent = parent - self.batch_size = 13 - self.seq_length = 7 - self.is_training = True - self.use_labels = False - self.vocab_size = 99 - self.hidden_size = 32 - self.num_hidden_layers = 5 - self.num_attention_heads = 4 - self.intermediate_size = 37 + 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.hidden_dropout_prob = 0.1 - self.attention_probs_dropout_prob = 0.1 - self.max_position_embeddings = 20 - self.eos_token_ids = [2] - self.pad_token_id = 1 - self.bos_token_id = 0 + self.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + 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_for_common(self): input_ids = ids_tensor([self.batch_size, self.seq_length - 1], self.vocab_size) - eos_tensor = tf.expand_dims(tf.constant([2] * self.batch_size), 1) + eos_tensor = tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size), 1) input_ids = tf.concat([input_ids, eos_tensor], axis=1) - input_ids = tf.clip_by_value(input_ids, 3, self.vocab_size + 1) + + decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) config = self.config_cls( vocab_size=self.vocab_size, @@ -82,7 +99,7 @@ class TFBartModelTester: decoder_start_token_id=self.pad_token_id, **self.config_updates, ) - inputs_dict = prepare_bart_inputs_dict(config, input_ids) + inputs_dict = prepare_bart_inputs_dict(config, input_ids, decoder_input_ids) return config, inputs_dict def check_decoder_model_past_large_inputs(self, config, inputs_dict): @@ -124,14 +141,25 @@ class TFBartModelTester: def prepare_bart_inputs_dict( config, input_ids, + decoder_input_ids, attention_mask=None, + decoder_attention_mask=None, ): if attention_mask is None: attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) + if decoder_attention_mask is None: + decoder_attention_mask = tf.concat( + [ + tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), + tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8), + ], + axis=-1, + ) return { "input_ids": input_ids, - "decoder_input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, } @@ -236,14 +264,46 @@ class TFBartModelTest(TFModelTesterMixin, unittest.TestCase): models_equal = False self.assertTrue(models_equal) - def test_saved_model_creation(self): - # This test is too long (>30sec) and makes fail the CI + @slow + def test_saved_model_with_hidden_states_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR pass + @slow + def test_saved_model_with_attentions_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation_extended(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + +def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): + """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" + if a is None and b is None: + return True + try: + if tf.debugging.assert_near(a, b, atol=atol): + return True + raise + except Exception: + msg = "{} != {}".format(a, b) + if prefix: + msg = prefix + ": " + msg + raise AssertionError(msg) + + +def _long_tensor(tok_lst): + return tf.constant(tok_lst, dtype=tf.int32) + @require_tf class TFBartHeadTests(unittest.TestCase): - vocab_size = 99 def _get_config_and_data(self): @@ -295,33 +355,10 @@ class TFBartHeadTests(unittest.TestCase): self.assertEqual(outputs.logits.shape, expected_shape) -def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): - """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" - if a is None and b is None: - return True - try: - if tf.debugging.assert_near(a, b, atol=atol): - return True - raise - except Exception: - msg = "{} != {}".format(a, b) - if prefix: - msg = prefix + ": " + msg - raise AssertionError(msg) - - -def _long_tensor(tok_lst): - return tf.constant(tok_lst, dtype=tf.int32) - - -TOLERANCE = 1e-4 - - -@is_pt_tf_cross_test @slow class TFBartModelIntegrationTest(unittest.TestCase): def test_inference_no_head(self): - model = TFBartModel.from_pretrained("facebook/bart-large", from_pt=True) + model = TFBartForConditionalGeneration.from_pretrained("facebook/bart-large").model input_ids = _long_tensor([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]) attention_mask = tf.cast(tf.math.not_equal(input_ids, model.config.pad_token_id), tf.int8) @@ -331,10 +368,10 @@ class TFBartModelIntegrationTest(unittest.TestCase): expected_slice = tf.convert_to_tensor( [[0.7144, 0.8143, -1.2813], [0.7144, 0.8143, -1.2813], [-0.0467, 2.5911, -2.1845]], ) - tf.debugging.assert_near(output[:, :3, :3], expected_slice, atol=TOLERANCE) + tf.debugging.assert_near(output[:, :3, :3], expected_slice, atol=1e-3) def test_cnn_summarization_same_as_fairseq_hard(self): - hf = TFBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn", from_pt=True) + hf = TFBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn") tok = self.tok FRANCE_ARTICLE = ' Marseille, France (CNN)The French prosecutor leading an investigation into the crash of Germanwings Flight 9525 insisted Wednesday that he was not aware of any video footage from on board the plane. Marseille prosecutor Brice Robin told CNN that "so far no videos were used in the crash investigation." He added, "A person who has such a video needs to immediately give it to the investigators." Robin\'s comments follow claims by two magazines, German daily Bild and French Paris Match, of a cell phone video showing the harrowing final seconds from on board Germanwings Flight 9525 as it crashed into the French Alps. All 150 on board were killed. Paris Match and Bild reported that the video was recovered from a phone at the wreckage site. The two publications described the supposed video, but did not post it on their websites. The publications said that they watched the video, which was found by a source close to the investigation. "One can hear cries of \'My God\' in several languages," Paris Match reported. "Metallic banging can also be heard more than three times, perhaps of the pilot trying to open the cockpit door with a heavy object. Towards the end, after a heavy shake, stronger than the others, the screaming intensifies. Then nothing." "It is a very disturbing scene," said Julian Reichelt, editor-in-chief of Bild online. An official with France\'s accident investigation agency, the BEA, said the agency is not aware of any such video. Lt. Col. Jean-Marc Menichini, a French Gendarmerie spokesman in charge of communications on rescue efforts around the Germanwings crash site, told CNN that the reports were "completely wrong" and "unwarranted." Cell phones have been collected at the site, he said, but that they "hadn\'t been exploited yet." Menichini said he believed the cell phones would need to be sent to the Criminal Research Institute in Rosny sous-Bois, near Paris, in order to be analyzed by specialized technicians working hand-in-hand with investigators. But none of the cell phones found so far have been sent to the institute, Menichini said. Asked whether staff involved in the search could have leaked a memory card to the media, Menichini answered with a categorical "no." Reichelt told "Erin Burnett: Outfront" that he had watched the video and stood by the report, saying Bild and Paris Match are "very confident" that the clip is real. He noted that investigators only revealed they\'d recovered cell phones from the crash site after Bild and Paris Match published their reports. "That is something we did not know before. ... Overall we can say many things of the investigation weren\'t revealed by the investigation at the beginning," he said. What was mental state of Germanwings co-pilot? German airline Lufthansa confirmed Tuesday that co-pilot Andreas Lubitz had battled depression years before he took the controls of Germanwings Flight 9525, which he\'s accused of deliberately crashing last week in the French Alps. Lubitz told his Lufthansa flight training school in 2009 that he had a "previous episode of severe depression," the airline said Tuesday. Email correspondence between Lubitz and the school discovered in an internal investigation, Lufthansa said, included medical documents he submitted in connection with resuming his flight training. The announcement indicates that Lufthansa, the parent company of Germanwings, knew of Lubitz\'s battle with depression, allowed him to continue training and ultimately put him in the cockpit. Lufthansa, whose CEO Carsten Spohr previously said Lubitz was 100% fit to fly, described its statement Tuesday as a "swift and seamless clarification" and said it was sharing the information and documents -- including training and medical records -- with public prosecutors. Spohr traveled to the crash site Wednesday, where recovery teams have been working for the past week to recover human remains and plane debris scattered across a steep mountainside. He saw the crisis center set up in Seyne-les-Alpes, laid a wreath in the village of Le Vernet, closer to the crash site, where grieving families have left flowers at a simple stone memorial. Menichini told CNN late Tuesday that no visible human remains were left at the site but recovery teams would keep searching. French President Francois Hollande, speaking Tuesday, said that it should be possible to identify all the victims using DNA analysis by the end of the week, sooner than authorities had previously suggested. In the meantime, the recovery of the victims\' personal belongings will start Wednesday, Menichini said. Among those personal belongings could be more cell phones belonging to the 144 passengers and six crew on board. Check out the latest from our correspondents . The details about Lubitz\'s correspondence with the flight school during his training were among several developments as investigators continued to delve into what caused the crash and Lubitz\'s possible motive for downing the jet. A Lufthansa spokesperson told CNN on Tuesday that Lubitz had a valid medical certificate, had passed all his examinations and "held all the licenses required." Earlier, a spokesman for the prosecutor\'s office in Dusseldorf, Christoph Kumpa, said medical records reveal Lubitz suffered from suicidal tendencies at some point before his aviation career and underwent psychotherapy before he got his pilot\'s license. Kumpa emphasized there\'s no evidence suggesting Lubitz was suicidal or acting aggressively before the crash. Investigators are looking into whether Lubitz feared his medical condition would cause him to lose his pilot\'s license, a European government official briefed on the investigation told CNN on Tuesday. While flying was "a big part of his life," the source said, it\'s only one theory being considered. Another source, a law enforcement official briefed on the investigation, also told CNN that authorities believe the primary motive for Lubitz to bring down the plane was that he feared he would not be allowed to fly because of his medical problems. Lubitz\'s girlfriend told investigators he had seen an eye doctor and a neuropsychologist, both of whom deemed him unfit to work recently and concluded he had psychological issues, the European government official said. But no matter what details emerge about his previous mental health struggles, there\'s more to the story, said Brian Russell, a forensic psychologist. "Psychology can explain why somebody would turn rage inward on themselves about the fact that maybe they weren\'t going to keep doing their job and they\'re upset about that and so they\'re suicidal," he said. "But there is no mental illness that explains why somebody then feels entitled to also take that rage and turn it outward on 149 other people who had nothing to do with the person\'s problems." Germanwings crash compensation: What we know . Who was the captain of Germanwings Flight 9525? CNN\'s Margot Haddad reported from Marseille and Pamela Brown from Dusseldorf, while Laura Smith-Spark wrote from London. CNN\'s Frederik Pleitgen, Pamela Boykoff, Antonia Mortensen, Sandrine Amiel and Anna-Maja Rappard contributed to this report.' # @noqa @@ -437,27 +474,3 @@ class FasterTFBartModelIntegrationTests(unittest.TestCase): expected = np.array([[-0.0828, -0.0251, -0.0674], [0.1277, 0.3311, -0.0255], [0.2613, -0.0840, -0.2763]]) assert np.allclose(features[0, :3, :3].numpy(), expected, atol=1e-3) - - -@require_tf -class TestTFSinusoidalPositionalEmbeddings(unittest.TestCase): - desired_weights = [ - [0, 0, 0, 0, 0], - [0.84147096, 0.82177866, 0.80180490, 0.78165019, 0.76140374], - [0.90929741, 0.93651021, 0.95829457, 0.97505713, 0.98720258], - ] - - def test_positional_emb_cache_logic(self): - emb1 = TFBartSinusoidalPositionalEmbedding(num_positions=32, embedding_dim=6) - no_cache = emb1((4, 10), past_key_values_length=0) - yes_cache = emb1((4, 10), past_key_values_length=2) - self.assertTrue(no_cache.shape == yes_cache.shape == (10, 6)) - self.assertListEqual(no_cache[2:].numpy().tolist(), yes_cache[:-2].numpy().tolist()) - - def test_positional_emb_weights_against_marian(self): - emb1 = TFBartSinusoidalPositionalEmbedding(num_positions=512, embedding_dim=512) - emb1.build(None) - weights = emb1.embeddings.numpy() - for i, (expected_weight, actual_weight) in enumerate(zip(self.desired_weights, weights)): - for j in range(5): - self.assertAlmostEqual(expected_weight[j], actual_weight[j], places=3) diff --git a/tests/test_modeling_tf_blenderbot.py b/tests/test_modeling_tf_blenderbot.py index 7cb63f47d6..cf71e75ba3 100644 --- a/tests/test_modeling_tf_blenderbot.py +++ b/tests/test_modeling_tf_blenderbot.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 HuggingFace Inc. team. +# 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. @@ -13,35 +13,158 @@ # See the License for the specific language governing permissions and # limitations under the License. + import unittest -from tests.test_configuration_common import ConfigTester -from tests.test_modeling_tf_bart import TFBartModelTester -from tests.test_modeling_tf_common import TFModelTesterMixin -from transformers import BlenderbotConfig, BlenderbotSmallTokenizer, is_tf_available +from transformers import BlenderbotConfig, BlenderbotTokenizer, is_tf_available from transformers.file_utils import cached_property -from transformers.testing_utils import is_pt_tf_cross_test, require_tf, require_tokenizers, slow +from transformers.testing_utils import require_tf, require_tokenizers, slow + +from .test_configuration_common import ConfigTester +from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor if is_tf_available(): import tensorflow as tf - from transformers import TFAutoModelForSeq2SeqLM, TFBlenderbotForConditionalGeneration + from transformers import TFAutoModelForSeq2SeqLM, TFBlenderbotForConditionalGeneration, TFBlenderbotModel -class TFBlenderbotModelTester(TFBartModelTester): - config_updates = dict( - normalize_before=True, - static_position_embeddings=True, - do_blenderbot_90_layernorm=True, - normalize_embeddings=True, - ) +@require_tf +class TFBlenderbotModelTester: config_cls = BlenderbotConfig + config_updates = {} + hidden_act = "gelu" + + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=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.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + 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_for_common(self): + input_ids = ids_tensor([self.batch_size, self.seq_length - 1], self.vocab_size) + eos_tensor = tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size), 1) + input_ids = tf.concat([input_ids, eos_tensor], axis=1) + + decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + config = self.config_cls( + 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, + dropout=self.hidden_dropout_prob, + attention_dropout=self.attention_probs_dropout_prob, + max_position_embeddings=self.max_position_embeddings, + eos_token_ids=[2], + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + decoder_start_token_id=self.pad_token_id, + **self.config_updates, + ) + inputs_dict = prepare_blenderbot_inputs_dict(config, input_ids, decoder_input_ids) + return config, inputs_dict + + def check_decoder_model_past_large_inputs(self, config, inputs_dict): + model = TFBlenderbotModel(config=config).get_decoder() + input_ids = inputs_dict["input_ids"] + + input_ids = input_ids[:1, :] + attention_mask = inputs_dict["attention_mask"][:1, :] + self.batch_size = 1 + + # first forward pass + outputs = model(input_ids, attention_mask=attention_mask, use_cache=True) + + output, past_key_values = outputs.to_tuple() + past_key_values = past_key_values[1] + + # create hypothetical next token and extent to next_input_ids + next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size) + next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8) + + # 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)[0] + output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0] + + self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1]) + + # 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] + + # test that outputs are equal for slice + tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3) + + +def prepare_blenderbot_inputs_dict( + config, + input_ids, + decoder_input_ids, + attention_mask=None, + decoder_attention_mask=None, +): + if attention_mask is None: + attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) + if decoder_attention_mask is None: + decoder_attention_mask = tf.concat( + [ + tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), + tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8), + ], + axis=-1, + ) + return { + "input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + } @require_tf class TFBlenderbotModelTest(TFModelTesterMixin, unittest.TestCase): - all_model_classes = (TFBlenderbotForConditionalGeneration,) if is_tf_available() else () + all_model_classes = (TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else () all_generative_model_classes = (TFBlenderbotForConditionalGeneration,) if is_tf_available() else () is_encoder_decoder = True test_pruning = False @@ -53,9 +176,9 @@ class TFBlenderbotModelTest(TFModelTesterMixin, unittest.TestCase): def test_config(self): self.config_tester.run_common_tests() - def test_inputs_embeds(self): - # inputs_embeds not supported - pass + def test_decoder_model_past_large_inputs(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common() + self.model_tester.check_decoder_model_past_large_inputs(*config_and_inputs) def test_model_common_attributes(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -77,8 +200,22 @@ class TFBlenderbotModelTest(TFModelTesterMixin, unittest.TestCase): name = model.get_bias() assert name is None + @slow + def test_saved_model_with_hidden_states_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + @slow + def test_saved_model_with_attentions_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + def test_saved_model_creation(self): - # This test is too long (>30sec) and makes fail the CI + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation_extended(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR pass def test_resize_token_embeddings(self): @@ -145,17 +282,33 @@ class TFBlenderbotModelTest(TFModelTesterMixin, unittest.TestCase): self.assertTrue(models_equal) -@is_pt_tf_cross_test +def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): + """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" + if a is None and b is None: + return True + try: + if tf.debugging.assert_near(a, b, atol=atol): + return True + raise + except Exception: + msg = "{} != {}".format(a, b) + if prefix: + msg = prefix + ": " + msg + raise AssertionError(msg) + + +def _long_tensor(tok_lst): + return tf.constant(tok_lst, dtype=tf.int32) + + @require_tokenizers -class TFBlenderbot90MIntegrationTests(unittest.TestCase): - src_text = [ - "Social anxiety\nWow, I am never shy. Do you have anxiety?\nYes. I end up sweating and blushing and feel like i'm going to throw up.\nand why is that?" - ] - model_name = "facebook/blenderbot-90M" +class TFBlenderbot400MIntegrationTests(unittest.TestCase): + src_text = ["My friends are cool but they eat too many carbs."] + model_name = "facebook/blenderbot-400M-distill" @cached_property def tokenizer(self): - return BlenderbotSmallTokenizer.from_pretrained(self.model_name) + return BlenderbotTokenizer.from_pretrained(self.model_name) @cached_property def model(self): @@ -163,17 +316,13 @@ class TFBlenderbot90MIntegrationTests(unittest.TestCase): return model @slow - def test_90_generation_from_long_input(self): + def test_generation_from_long_input(self): model_inputs = self.tokenizer(self.src_text, return_tensors="tf") generated_ids = self.model.generate( model_inputs.input_ids, - attention_mask=model_inputs.attention_mask, - num_beams=2, - use_cache=True, ) generated_words = self.tokenizer.batch_decode(generated_ids.numpy(), skip_special_tokens=True)[0] - assert generated_words in ( - "i don't know. i just feel like i'm going to throw up. it's not fun.", - "i'm not sure. i just feel like i've been feeling like i have to be in a certain place", - "i'm not sure. i just feel like i've been in a bad situation.", + assert ( + generated_words + == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?" ) diff --git a/tests/test_modeling_tf_blenderbot_small.py b/tests/test_modeling_tf_blenderbot_small.py new file mode 100644 index 0000000000..dfba5f40a0 --- /dev/null +++ b/tests/test_modeling_tf_blenderbot_small.py @@ -0,0 +1,328 @@ +# 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. + + +import unittest + +from transformers import BlenderbotSmallConfig, BlenderbotSmallTokenizer, is_tf_available +from transformers.file_utils import cached_property +from transformers.testing_utils import require_tf, require_tokenizers, slow + +from .test_configuration_common import ConfigTester +from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor + + +if is_tf_available(): + import tensorflow as tf + + from transformers import TFAutoModelForSeq2SeqLM, TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel + + +@require_tf +class TFBlenderbotSmallModelTester: + config_cls = BlenderbotSmallConfig + config_updates = {} + hidden_act = "gelu" + + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=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.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + 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_for_common(self): + input_ids = ids_tensor([self.batch_size, self.seq_length - 1], self.vocab_size) + eos_tensor = tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size), 1) + input_ids = tf.concat([input_ids, eos_tensor], axis=1) + + decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + config = self.config_cls( + 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, + dropout=self.hidden_dropout_prob, + attention_dropout=self.attention_probs_dropout_prob, + max_position_embeddings=self.max_position_embeddings, + eos_token_ids=[2], + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + decoder_start_token_id=self.pad_token_id, + **self.config_updates, + ) + inputs_dict = prepare_blenderbot_small_inputs_dict(config, input_ids, decoder_input_ids) + return config, inputs_dict + + def check_decoder_model_past_large_inputs(self, config, inputs_dict): + model = TFBlenderbotSmallModel(config=config).get_decoder() + input_ids = inputs_dict["input_ids"] + + input_ids = input_ids[:1, :] + attention_mask = inputs_dict["attention_mask"][:1, :] + self.batch_size = 1 + + # first forward pass + outputs = model(input_ids, attention_mask=attention_mask, use_cache=True) + + output, past_key_values = outputs.to_tuple() + past_key_values = past_key_values[1] + + # create hypothetical next token and extent to next_input_ids + next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size) + next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8) + + # 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)[0] + output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0] + + self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1]) + + # 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] + + # test that outputs are equal for slice + tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3) + + +def prepare_blenderbot_small_inputs_dict( + config, + input_ids, + decoder_input_ids, + attention_mask=None, + decoder_attention_mask=None, +): + if attention_mask is None: + attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) + if decoder_attention_mask is None: + decoder_attention_mask = tf.concat( + [ + tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), + tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8), + ], + axis=-1, + ) + return { + "input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + } + + +@require_tf +class TFBlenderbotSmallModelTest(TFModelTesterMixin, unittest.TestCase): + all_model_classes = ( + (TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel) if is_tf_available() else () + ) + all_generative_model_classes = (TFBlenderbotSmallForConditionalGeneration,) if is_tf_available() else () + is_encoder_decoder = True + test_pruning = False + + def setUp(self): + self.model_tester = TFBlenderbotSmallModelTester(self) + self.config_tester = ConfigTester(self, config_class=BlenderbotSmallConfig) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_decoder_model_past_large_inputs(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common() + self.model_tester.check_decoder_model_past_large_inputs(*config_and_inputs) + + def test_model_common_attributes(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + for model_class in self.all_model_classes: + model = model_class(config) + assert isinstance(model.get_input_embeddings(), tf.keras.layers.Layer) + x = model.get_output_layer_with_bias() + assert x is None + name = model.get_prefix_bias_name() + assert name is None + + @slow + def test_saved_model_with_hidden_states_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + @slow + def test_saved_model_with_attentions_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation_extended(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_resize_token_embeddings(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + def _get_word_embedding_weight(model, embedding_layer): + if hasattr(embedding_layer, "weight"): + return embedding_layer.weight + else: + # Here we build the word embeddings weights if not exists. + # And then we retry to get the attribute once built. + model(model.dummy_inputs) + if hasattr(embedding_layer, "weight"): + return embedding_layer.weight + else: + return None + + for model_class in self.all_model_classes: + for size in [config.vocab_size - 10, config.vocab_size + 10, None]: + # build the embeddings + model = model_class(config=config) + old_input_embeddings = _get_word_embedding_weight(model, model.get_input_embeddings()) + old_output_embeddings = _get_word_embedding_weight(model, model.get_output_embeddings()) + old_final_logits_bias = model.get_bias() + + # reshape the embeddings + model.resize_token_embeddings(size) + new_input_embeddings = _get_word_embedding_weight(model, model.get_input_embeddings()) + new_output_embeddings = _get_word_embedding_weight(model, model.get_output_embeddings()) + new_final_logits_bias = model.get_bias() + + # check that the resized embeddings size matches the desired size. + assert_size = size if size is not None else config.vocab_size + + self.assertEqual(new_input_embeddings.shape[0], assert_size) + + # check that weights remain the same after resizing + models_equal = True + for p1, p2 in zip(old_input_embeddings.value(), new_input_embeddings.value()): + if tf.math.reduce_sum(tf.math.abs(p1 - p2)) > 0: + models_equal = False + self.assertTrue(models_equal) + + if old_output_embeddings is not None and new_output_embeddings is not None: + self.assertEqual(new_output_embeddings.shape[0], assert_size) + + models_equal = True + for p1, p2 in zip(old_output_embeddings.value(), new_output_embeddings.value()): + if tf.math.reduce_sum(tf.math.abs(p1 - p2)) > 0: + models_equal = False + self.assertTrue(models_equal) + + if old_final_logits_bias is not None and new_final_logits_bias is not None: + old_final_logits_bias = old_final_logits_bias["final_logits_bias"] + new_final_logits_bias = new_final_logits_bias["final_logits_bias"] + self.assertEqual(new_final_logits_bias.shape[0], 1) + self.assertEqual(new_final_logits_bias.shape[1], assert_size) + + models_equal = True + for old, new in zip(old_final_logits_bias.value(), new_final_logits_bias.value()): + for p1, p2 in zip(old, new): + if tf.math.reduce_sum(tf.math.abs(p1 - p2)) > 0: + models_equal = False + self.assertTrue(models_equal) + + +def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): + """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" + if a is None and b is None: + return True + try: + if tf.debugging.assert_near(a, b, atol=atol): + return True + raise + except Exception: + msg = "{} != {}".format(a, b) + if prefix: + msg = prefix + ": " + msg + raise AssertionError(msg) + + +def _long_tensor(tok_lst): + return tf.constant(tok_lst, dtype=tf.int32) + + +@require_tokenizers +class TFBlenderbot90MIntegrationTests(unittest.TestCase): + src_text = [ + "Social anxiety\nWow, I am never shy. Do you have anxiety?\nYes. I end up sweating and blushing and feel like i'm going to throw up.\nand why is that?" + ] + model_name = "facebook/blenderbot_small-90M" + + @cached_property + def tokenizer(self): + # use "old" tokenizer here because of bug when downloading new tokenizer + return BlenderbotSmallTokenizer.from_pretrained("facebook/blenderbot-90M") + + @cached_property + def model(self): + model = TFAutoModelForSeq2SeqLM.from_pretrained(self.model_name) + return model + + @slow + def test_90_generation_from_long_input(self): + model_inputs = self.tokenizer(self.src_text, return_tensors="tf") + generated_ids = self.model.generate( + model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + num_beams=2, + use_cache=True, + ) + generated_words = self.tokenizer.batch_decode(generated_ids.numpy(), skip_special_tokens=True)[0] + assert generated_words in ( + "i don't know. i just feel like i'm going to throw up. it's not fun.", + "i'm not sure. i just feel like i've been feeling like i have to be in a certain place", + "i'm not sure. i just feel like i've been in a bad situation.", + ) diff --git a/tests/test_modeling_tf_marian.py b/tests/test_modeling_tf_marian.py index 7c3c87e0d3..dec14450a9 100644 --- a/tests/test_modeling_tf_marian.py +++ b/tests/test_modeling_tf_marian.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 HuggingFace Inc. team. +# 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. @@ -13,48 +13,174 @@ # See the License for the specific language governing permissions and # limitations under the License. + import tempfile import unittest import warnings from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available from transformers.file_utils import cached_property -from transformers.testing_utils import is_pt_tf_cross_test, require_sentencepiece, require_tf, require_tokenizers, slow +from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow from .test_configuration_common import ConfigTester -from .test_modeling_tf_bart import TFBartModelTester -from .test_modeling_tf_common import TFModelTesterMixin +from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor if is_tf_available(): import tensorflow as tf - from transformers import TFAutoModelForSeq2SeqLM, TFMarianMTModel - - -class ModelTester(TFBartModelTester): - config_updates = dict(static_position_embeddings=True, add_bias_logits=True) - config_cls = MarianConfig + from transformers import TFAutoModelForSeq2SeqLM, TFMarianModel, TFMarianMTModel @require_tf -class TFMarianMTModelTest(TFModelTesterMixin, unittest.TestCase): - all_model_classes = (TFMarianMTModel,) if is_tf_available() else () +class TFMarianModelTester: + config_cls = MarianConfig + config_updates = {} + hidden_act = "gelu" + + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=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.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + 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_for_common(self): + input_ids = ids_tensor([self.batch_size, self.seq_length - 1], self.vocab_size) + eos_tensor = tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size), 1) + input_ids = tf.concat([input_ids, eos_tensor], axis=1) + + decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + config = self.config_cls( + 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, + dropout=self.hidden_dropout_prob, + attention_dropout=self.attention_probs_dropout_prob, + max_position_embeddings=self.max_position_embeddings, + eos_token_ids=[2], + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + decoder_start_token_id=self.pad_token_id, + **self.config_updates, + ) + inputs_dict = prepare_marian_inputs_dict(config, input_ids, decoder_input_ids) + return config, inputs_dict + + def check_decoder_model_past_large_inputs(self, config, inputs_dict): + model = TFMarianModel(config=config).get_decoder() + input_ids = inputs_dict["input_ids"] + + input_ids = input_ids[:1, :] + attention_mask = inputs_dict["attention_mask"][:1, :] + self.batch_size = 1 + + # first forward pass + outputs = model(input_ids, attention_mask=attention_mask, use_cache=True) + + output, past_key_values = outputs.to_tuple() + past_key_values = past_key_values[1] + + # create hypothetical next token and extent to next_input_ids + next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size) + next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8) + + # 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)[0] + output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0] + + self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1]) + + # 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] + + # test that outputs are equal for slice + tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3) + + +def prepare_marian_inputs_dict( + config, + input_ids, + decoder_input_ids, + attention_mask=None, + decoder_attention_mask=None, +): + if attention_mask is None: + attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) + if decoder_attention_mask is None: + decoder_attention_mask = tf.concat( + [ + tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), + tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8), + ], + axis=-1, + ) + return { + "input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + } + + +@require_tf +class TFMarianModelTest(TFModelTesterMixin, unittest.TestCase): + all_model_classes = (TFMarianMTModel, TFMarianModel) if is_tf_available() else () all_generative_model_classes = (TFMarianMTModel,) if is_tf_available() else () - model_tester_cls = ModelTester is_encoder_decoder = True test_pruning = False def setUp(self): - self.model_tester = self.model_tester_cls(self) + self.model_tester = TFMarianModelTester(self) self.config_tester = ConfigTester(self, config_class=MarianConfig) def test_config(self): self.config_tester.run_common_tests() - def test_inputs_embeds(self): - # inputs_embeds not supported - pass + def test_decoder_model_past_large_inputs(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common() + self.model_tester.check_decoder_model_past_large_inputs(*config_and_inputs) def test_compile_tf_model(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -107,8 +233,22 @@ class TFMarianMTModelTest(TFModelTesterMixin, unittest.TestCase): name = model.get_bias() assert name is None + @slow + def test_saved_model_with_hidden_states_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + @slow + def test_saved_model_with_attentions_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + def test_saved_model_creation(self): - # This test is too long (>30sec) and makes fail the CI + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation_extended(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR pass def test_resize_token_embeddings(self): @@ -175,6 +315,25 @@ class TFMarianMTModelTest(TFModelTesterMixin, unittest.TestCase): self.assertTrue(models_equal) +def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): + """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" + if a is None and b is None: + return True + try: + if tf.debugging.assert_near(a, b, atol=atol): + return True + raise + except Exception: + msg = "{} != {}".format(a, b) + if prefix: + msg = prefix + ": " + msg + raise AssertionError(msg) + + +def _long_tensor(tok_lst): + return tf.constant(tok_lst, dtype=tf.int32) + + class AbstractMarianIntegrationTest(unittest.TestCase): maxDiff = 1000 # show more chars for failing integration tests @@ -219,7 +378,6 @@ class AbstractMarianIntegrationTest(unittest.TestCase): @require_sentencepiece @require_tokenizers -@is_pt_tf_cross_test class TestMarian_MT_EN(AbstractMarianIntegrationTest): """Cover low resource/high perplexity setting. This breaks if pad_token_id logits not set to LARGE_NEGATIVE.""" @@ -233,7 +391,6 @@ class TestMarian_MT_EN(AbstractMarianIntegrationTest): self._assert_generated_batch_equal_expected() -@is_pt_tf_cross_test @require_sentencepiece @require_tokenizers class TestMarian_en_zh(AbstractMarianIntegrationTest): @@ -247,7 +404,6 @@ class TestMarian_en_zh(AbstractMarianIntegrationTest): self._assert_generated_batch_equal_expected() -@is_pt_tf_cross_test @require_sentencepiece @require_tokenizers class TestMarian_en_ROMANCE(AbstractMarianIntegrationTest): diff --git a/tests/test_modeling_tf_mbart.py b/tests/test_modeling_tf_mbart.py index 0916f19b11..6208622f2d 100644 --- a/tests/test_modeling_tf_mbart.py +++ b/tests/test_modeling_tf_mbart.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 HuggingFace Inc. team. +# 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. @@ -12,47 +12,107 @@ # 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. + import tempfile import unittest -from tests.test_configuration_common import ConfigTester -from tests.test_modeling_tf_bart import TFBartModelTester -from tests.test_modeling_tf_common import TFModelTesterMixin from transformers import AutoTokenizer, MBartConfig, is_tf_available from transformers.file_utils import cached_property -from transformers.testing_utils import is_pt_tf_cross_test, require_sentencepiece, require_tf, require_tokenizers, slow +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, ids_tensor if is_tf_available(): - import tensorflow as tf - from transformers import TFAutoModelForSeq2SeqLM, TFMBartForConditionalGeneration - - -class ModelTester(TFBartModelTester): - config_updates = dict(normalize_before=True, add_final_layer_norm=True) - config_cls = MBartConfig + from transformers import TFAutoModelForSeq2SeqLM, TFMBartForConditionalGeneration, TFMBartModel @require_tf -class TFMBartModelTest(TFModelTesterMixin, unittest.TestCase): - all_model_classes = (TFMBartForConditionalGeneration,) if is_tf_available() else () - all_generative_model_classes = (TFMBartForConditionalGeneration,) if is_tf_available() else () - model_tester_cls = ModelTester - is_encoder_decoder = True - test_pruning = False +class TFMBartModelTester: + config_cls = MBartConfig + config_updates = {} + hidden_act = "gelu" - def setUp(self): - self.model_tester = self.model_tester_cls(self) - self.config_tester = ConfigTester(self, config_class=MBartConfig) + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=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.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id - def test_config(self): - self.config_tester.run_common_tests() + def prepare_config_and_inputs_for_common(self): + input_ids = ids_tensor([self.batch_size, self.seq_length - 1], self.vocab_size) + eos_tensor = tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size), 1) + input_ids = tf.concat([input_ids, eos_tensor], axis=1) - def test_inputs_embeds(self): - # inputs_embeds not supported - pass + decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + config = self.config_cls( + 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, + dropout=self.hidden_dropout_prob, + attention_dropout=self.attention_probs_dropout_prob, + max_position_embeddings=self.max_position_embeddings, + eos_token_ids=[2], + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + decoder_start_token_id=self.pad_token_id, + **self.config_updates, + ) + inputs_dict = prepare_mbart_inputs_dict(config, input_ids, decoder_input_ids) + return config, inputs_dict + + def check_decoder_model_past_large_inputs(self, config, inputs_dict): + model = TFMBartModel(config=config).get_decoder() + input_ids = inputs_dict["input_ids"] + + input_ids = input_ids[:1, :] + attention_mask = inputs_dict["attention_mask"][:1, :] + self.batch_size = 1 + + # first forward pass + outputs = model(input_ids, attention_mask=attention_mask, use_cache=True) + + output, past_key_values = outputs.to_tuple() + past_key_values = past_key_values[1] def test_compile_tf_model(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -60,13 +120,11 @@ class TFMBartModelTest(TFModelTesterMixin, unittest.TestCase): optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0) loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) metric = tf.keras.metrics.SparseCategoricalAccuracy("accuracy") - model_class = self.all_generative_model_classes[0] input_ids = { "decoder_input_ids": tf.keras.Input(batch_shape=(2, 2000), name="decoder_input_ids", dtype="int32"), "input_ids": tf.keras.Input(batch_shape=(2, 2000), name="input_ids", dtype="int32"), } - # Prepare our model model = model_class(config) model(self._prepare_for_class(inputs_dict, model_class)) # Model must be called before saving. @@ -74,17 +132,58 @@ class TFMBartModelTest(TFModelTesterMixin, unittest.TestCase): with tempfile.TemporaryDirectory() as tmpdirname: model.save_pretrained(tmpdirname) model = model_class.from_pretrained(tmpdirname) - outputs_dict = model(input_ids) hidden_states = outputs_dict[0] - # Add a dense layer on top to test integration with other keras modules outputs = tf.keras.layers.Dense(2, activation="softmax", name="outputs")(hidden_states) - # Compile extended model extended_model = tf.keras.Model(inputs=[input_ids], outputs=[outputs]) extended_model.compile(optimizer=optimizer, loss=loss, metrics=[metric]) + +def prepare_mbart_inputs_dict( + config, + input_ids, + decoder_input_ids, + attention_mask=None, + decoder_attention_mask=None, +): + if attention_mask is None: + attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) + if decoder_attention_mask is None: + decoder_attention_mask = tf.concat( + [ + tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), + tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8), + ], + axis=-1, + ) + return { + "input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + } + + +@require_tf +class TFMBartModelTest(TFModelTesterMixin, unittest.TestCase): + all_model_classes = (TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else () + all_generative_model_classes = (TFMBartForConditionalGeneration,) if is_tf_available() else () + is_encoder_decoder = True + test_pruning = False + + def setUp(self): + self.model_tester = TFMBartModelTester(self) + self.config_tester = ConfigTester(self, config_class=MBartConfig) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_decoder_model_past_large_inputs(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common() + self.model_tester.check_decoder_model_past_large_inputs(*config_and_inputs) + def test_model_common_attributes(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -105,8 +204,22 @@ class TFMBartModelTest(TFModelTesterMixin, unittest.TestCase): name = model.get_bias() assert name is None + @slow + def test_saved_model_with_hidden_states_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + @slow + def test_saved_model_with_attentions_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + def test_saved_model_creation(self): - # This test is too long (>30sec) and makes fail the CI + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation_extended(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR pass def test_resize_token_embeddings(self): @@ -173,10 +286,31 @@ class TFMBartModelTest(TFModelTesterMixin, unittest.TestCase): self.assertTrue(models_equal) -@is_pt_tf_cross_test +def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): + """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" + if a is None and b is None: + return True + try: + if tf.debugging.assert_near(a, b, atol=atol): + return True + raise + except Exception: + msg = "{} != {}".format(a, b) + if prefix: + msg = prefix + ": " + msg + raise AssertionError(msg) + + +def _long_tensor(tok_lst): + return tf.constant(tok_lst, dtype=tf.int32) + + +TOLERANCE = 1e-4 + + @require_sentencepiece @require_tokenizers -class TestMBartEnRO(unittest.TestCase): +class TFMBartModelIntegrationTest(unittest.TestCase): src_text = [ " UN Chief Says There Is No Military Solution in Syria", ] @@ -191,7 +325,7 @@ class TestMBartEnRO(unittest.TestCase): @cached_property def model(self): - model = TFAutoModelForSeq2SeqLM.from_pretrained(self.model_name, from_pt=True) + model = TFAutoModelForSeq2SeqLM.from_pretrained(self.model_name) return model def _assert_generated_batch_equal_expected(self, **tokenizer_kwargs): diff --git a/tests/test_modeling_tf_pegasus.py b/tests/test_modeling_tf_pegasus.py index 5c91f35af2..559c644fbb 100644 --- a/tests/test_modeling_tf_pegasus.py +++ b/tests/test_modeling_tf_pegasus.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2020 HuggingFace Inc. team. +# 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. @@ -18,46 +18,167 @@ import unittest from transformers import AutoTokenizer, PegasusConfig, is_tf_available from transformers.file_utils import cached_property -from transformers.testing_utils import is_pt_tf_cross_test, require_sentencepiece, require_tf, require_tokenizers, slow +from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow from .test_configuration_common import ConfigTester -from .test_modeling_tf_bart import TFBartModelTester -from .test_modeling_tf_common import TFModelTesterMixin +from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor if is_tf_available(): import tensorflow as tf - from transformers import TFAutoModelForSeq2SeqLM, TFPegasusForConditionalGeneration + from transformers import TFAutoModelForSeq2SeqLM, TFPegasusForConditionalGeneration, TFPegasusModel -class ModelTester(TFBartModelTester): - config_updates = dict( - normalize_before=True, - static_position_embeddings=True, - ) - hidden_act = "relu" +@require_tf +class TFPegasusModelTester: config_cls = PegasusConfig + config_updates = {} + hidden_act = "gelu" + + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_labels=False, + vocab_size=99, + hidden_size=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=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.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + 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_for_common(self): + input_ids = ids_tensor([self.batch_size, self.seq_length - 1], self.vocab_size) + eos_tensor = tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size), 1) + input_ids = tf.concat([input_ids, eos_tensor], axis=1) + + decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + config = self.config_cls( + 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, + dropout=self.hidden_dropout_prob, + attention_dropout=self.attention_probs_dropout_prob, + max_position_embeddings=self.max_position_embeddings, + eos_token_ids=[2], + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + decoder_start_token_id=self.pad_token_id, + **self.config_updates, + ) + inputs_dict = prepare_pegasus_inputs_dict(config, input_ids, decoder_input_ids) + return config, inputs_dict + + def check_decoder_model_past_large_inputs(self, config, inputs_dict): + model = TFPegasusModel(config=config).get_decoder() + input_ids = inputs_dict["input_ids"] + + input_ids = input_ids[:1, :] + attention_mask = inputs_dict["attention_mask"][:1, :] + self.batch_size = 1 + + # first forward pass + outputs = model(input_ids, attention_mask=attention_mask, use_cache=True) + + output, past_key_values = outputs.to_tuple() + past_key_values = past_key_values[1] + + # create hypothetical next token and extent to next_input_ids + next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size) + next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8) + + # 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)[0] + output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0] + + self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1]) + + # 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] + + # test that outputs are equal for slice + tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3) + + +def prepare_pegasus_inputs_dict( + config, + input_ids, + decoder_input_ids, + attention_mask=None, + decoder_attention_mask=None, +): + if attention_mask is None: + attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8) + if decoder_attention_mask is None: + decoder_attention_mask = tf.concat( + [ + tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8), + tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8), + ], + axis=-1, + ) + return { + "input_ids": input_ids, + "decoder_input_ids": decoder_input_ids, + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + } @require_tf class TFPegasusModelTest(TFModelTesterMixin, unittest.TestCase): - all_model_classes = (TFPegasusForConditionalGeneration,) if is_tf_available() else () + all_model_classes = (TFPegasusForConditionalGeneration, TFPegasusModel) if is_tf_available() else () all_generative_model_classes = (TFPegasusForConditionalGeneration,) if is_tf_available() else () - model_tester_cls = ModelTester is_encoder_decoder = True test_pruning = False def setUp(self): - self.model_tester = self.model_tester_cls(self) + self.model_tester = TFPegasusModelTester(self) self.config_tester = ConfigTester(self, config_class=PegasusConfig) def test_config(self): self.config_tester.run_common_tests() - def test_inputs_embeds(self): - # inputs_embeds not supported - pass + def test_decoder_model_past_large_inputs(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common() + self.model_tester.check_decoder_model_past_large_inputs(*config_and_inputs) def test_compile_tf_model(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -110,8 +231,22 @@ class TFPegasusModelTest(TFModelTesterMixin, unittest.TestCase): name = model.get_bias() assert name is None + @slow + def test_saved_model_with_hidden_states_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + @slow + def test_saved_model_with_attentions_output(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + def test_saved_model_creation(self): - # This test is too long (>30sec) and makes fail the CI + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR + pass + + def test_saved_model_creation_extended(self): + # TODO(JPLU, PVP) - fix this with s2s tf-serving PR pass def test_resize_token_embeddings(self): @@ -178,7 +313,25 @@ class TFPegasusModelTest(TFModelTesterMixin, unittest.TestCase): self.assertTrue(models_equal) -@is_pt_tf_cross_test +def _assert_tensors_equal(a, b, atol=1e-12, prefix=""): + """If tensors not close, or a and b arent both tensors, raise a nice Assertion error.""" + if a is None and b is None: + return True + try: + if tf.debugging.assert_near(a, b, atol=atol): + return True + raise + except Exception: + msg = "{} != {}".format(a, b) + if prefix: + msg = prefix + ": " + msg + raise AssertionError(msg) + + +def _long_tensor(tok_lst): + return tf.constant(tok_lst, dtype=tf.int32) + + @require_sentencepiece @require_tokenizers class TFPegasusIntegrationTests(unittest.TestCase): @@ -198,7 +351,7 @@ class TFPegasusIntegrationTests(unittest.TestCase): @cached_property def model(self): - model = TFAutoModelForSeq2SeqLM.from_pretrained(self.model_name, from_pt=True) + model = TFAutoModelForSeq2SeqLM.from_pretrained(self.model_name) return model def _assert_generated_batch_equal_expected(self, **tokenizer_kwargs):