From b256f3518d470ba53be519992c3b9d97d174db48 Mon Sep 17 00:00:00 2001 From: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com> Date: Wed, 9 Mar 2022 19:53:01 +0100 Subject: [PATCH] Add FlaxBartForCausalLM (#15995) * add causal lm * add CausalLM tests * Add FlaxBartForCausalLM * Add EncoderDecoder model tests * change docstring * make repo-consistency * suggested changes * remove jax ops * correction * rename pre-trained decoder model --- docs/source/model_doc/bart.mdx | 5 + src/transformers/__init__.py | 4 + .../models/auto/modeling_flax_auto.py | 1 + src/transformers/models/bart/__init__.py | 4 + .../models/bart/modeling_flax_bart.py | 261 ++++++++++++++++++ src/transformers/utils/dummy_flax_objects.py | 14 + tests/bart/test_modeling_flax_bart.py | 98 ++++++- .../test_modeling_flax_encoder_decoder.py | 40 +++ ...st_modeling_flax_speech_encoder_decoder.py | 119 ++++++++ utils/check_repo.py | 1 + 10 files changed, 543 insertions(+), 4 deletions(-) diff --git a/docs/source/model_doc/bart.mdx b/docs/source/model_doc/bart.mdx index 38d6b6ea95..54edb509d9 100644 --- a/docs/source/model_doc/bart.mdx +++ b/docs/source/model_doc/bart.mdx @@ -152,3 +152,8 @@ assert tok.batch_decode(generated_ids, skip_special_tokens=True) == [ - __call__ - encode - decode + +## FlaxBartForCausalLM + +[[autodoc]] FlaxBartForCausalLM + - __call__ \ No newline at end of file diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 69f21f0120..c2c3e7a2c4 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -2198,6 +2198,8 @@ if is_flax_available(): _import_structure["models.bart"].extend( [ + "FlaxBartDecoderPreTrainedModel", + "FlaxBartForCausalLM", "FlaxBartForConditionalGeneration", "FlaxBartForQuestionAnswering", "FlaxBartForSequenceClassification", @@ -4170,6 +4172,8 @@ if TYPE_CHECKING: FlaxAutoModelForVision2Seq, ) from .models.bart import ( + FlaxBartDecoderPreTrainedModel, + FlaxBartForCausalLM, FlaxBartForConditionalGeneration, FlaxBartForQuestionAnswering, FlaxBartForSequenceClassification, diff --git a/src/transformers/models/auto/modeling_flax_auto.py b/src/transformers/models/auto/modeling_flax_auto.py index 3956d823e9..4475766bdf 100644 --- a/src/transformers/models/auto/modeling_flax_auto.py +++ b/src/transformers/models/auto/modeling_flax_auto.py @@ -126,6 +126,7 @@ FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict( ("gpt_neo", "FlaxGPTNeoForCausalLM"), ("gptj", "FlaxGPTJForCausalLM"), ("xglm", "FlaxXGLMForCausalLM"), + ("bart", "FlaxBartForCausalLM"), ] ) diff --git a/src/transformers/models/bart/__init__.py b/src/transformers/models/bart/__init__.py index ffcf517416..db499e3ce0 100644 --- a/src/transformers/models/bart/__init__.py +++ b/src/transformers/models/bart/__init__.py @@ -45,6 +45,8 @@ if is_tf_available(): if is_flax_available(): _import_structure["modeling_flax_bart"] = [ + "FlaxBartDecoderPreTrainedModel", + "FlaxBartForCausalLM", "FlaxBartForConditionalGeneration", "FlaxBartForQuestionAnswering", "FlaxBartForSequenceClassification", @@ -76,6 +78,8 @@ if TYPE_CHECKING: if is_flax_available(): from .modeling_flax_bart import ( + FlaxBartDecoderPreTrainedModel, + FlaxBartForCausalLM, FlaxBartForConditionalGeneration, FlaxBartForQuestionAnswering, FlaxBartForSequenceClassification, diff --git a/src/transformers/models/bart/modeling_flax_bart.py b/src/transformers/models/bart/modeling_flax_bart.py index cdec52f6e1..386bddbb26 100644 --- a/src/transformers/models/bart/modeling_flax_bart.py +++ b/src/transformers/models/bart/modeling_flax_bart.py @@ -1725,3 +1725,264 @@ append_call_sample_docstring( FlaxSeq2SeqQuestionAnsweringModelOutput, _CONFIG_FOR_DOC, ) + + +class FlaxBartDecoderPreTrainedModel(FlaxPreTrainedModel): + config_class = BartConfig + base_model_prefix: str = "model" + module_class: nn.Module = None + + def __init__( + self, + config: BartConfig, + input_shape: Tuple[int] = (1, 1), + seed: int = 0, + dtype: jnp.dtype = jnp.float32, + **kwargs + ): + config.is_decoder = True + config.is_encoder_decoder = False + module = self.module_class(config=config, dtype=dtype, **kwargs) + super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype) + + def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict: + # init input tensors + input_ids = jnp.zeros(input_shape, dtype="i4") + attention_mask = jnp.ones_like(input_ids) + + batch_size, sequence_length = input_ids.shape + position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)) + + params_rng, dropout_rng = jax.random.split(rng) + rngs = {"params": params_rng, "dropout": dropout_rng} + encoder_hidden_states = jnp.zeros(input_shape + (self.config.d_model,)) + encoder_attention_mask = attention_mask + module_init_outputs = self.module.init( + rngs, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states, + encoder_attention_mask, + return_dict=False, + ) + return module_init_outputs["params"] + + def init_cache(self, batch_size, max_length): + r""" + Args: + batch_size (`int`): + batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache. + max_length (`int`): + maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized + cache. + """ + # init input variables to retrieve cache + input_ids = jnp.ones((batch_size, max_length), dtype="i4") + attention_mask = jnp.ones_like(input_ids, dtype="i4") + position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape) + + init_variables = self.module.init( + jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True + ) + return unfreeze(init_variables["cache"]) + + @add_start_docstrings_to_model_forward(BART_DECODE_INPUTS_DOCSTRING) + def __call__( + self, + input_ids: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + position_ids: Optional[jnp.ndarray] = None, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + train: bool = False, + params: dict = None, + past_key_values: dict = None, + dropout_rng: PRNGKey = None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if encoder_hidden_states is not None and encoder_attention_mask is None: + batch_size, sequence_length = encoder_hidden_states.shape[:2] + encoder_attention_mask = jnp.ones((batch_size, sequence_length)) + + # prepare decoder inputs + if attention_mask is None: + attention_mask = jnp.ones_like(input_ids) + if position_ids is None: + batch_size, sequence_length = input_ids.shape + position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)) + + # Handle any PRNG if needed + rngs = {"dropout": dropout_rng} if dropout_rng is not None else {} + + inputs = {"params": params or self.params} + + # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed + # down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be + # changed by FlaxBartAttention module + if past_key_values: + inputs["cache"] = past_key_values + mutable = ["cache"] + else: + mutable = False + + outputs = self.module.apply( + inputs, + input_ids=jnp.array(input_ids, dtype="i4"), + attention_mask=jnp.array(attention_mask, dtype="i4"), + position_ids=jnp.array(position_ids, dtype="i4"), + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + deterministic=not train, + rngs=rngs, + mutable=mutable, + ) + + # add updated cache to model output + if past_key_values is not None and return_dict: + outputs, past_key_values = outputs + outputs["past_key_values"] = unfreeze(past_key_values["cache"]) + return outputs + elif past_key_values is not None and not return_dict: + outputs, past_key_values = outputs + outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:] + + return outputs + + +class FlaxBartDecoderWrapper(nn.Module): + """ + This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is + used in combination with the [`EncoderDecoderModel`] framework. + """ + + config: BartConfig + dtype: jnp.dtype = jnp.float32 + + def setup(self): + embed_dim = self.config.d_model + embed_tokens = nn.Embed( + self.config.vocab_size, + embed_dim, + embedding_init=jax.nn.initializers.normal(self.config.init_std), + ) + self.decoder = FlaxBartDecoder(config=self.config, embed_tokens=embed_tokens, dtype=self.dtype) + + def __call__(self, *args, **kwargs): + return self.decoder(*args, **kwargs) + + +class FlaxBartForCausalLMModule(nn.Module): + config: BartConfig + dtype: jnp.dtype = jnp.float32 + + def setup(self): + self.model = FlaxBartDecoderWrapper(config=self.config, dtype=self.dtype) + self.lm_head = nn.Dense( + self.config.vocab_size, + use_bias=False, + dtype=self.dtype, + kernel_init=jax.nn.initializers.normal(self.config.init_std), + ) + + def __call__( + self, + input_ids, + attention_mask, + position_ids, + encoder_hidden_states: Optional[jnp.ndarray] = None, + encoder_attention_mask: Optional[jnp.ndarray] = None, + init_cache: bool = False, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + deterministic: bool = True, + ): + + outputs = self.model( + input_ids, + attention_mask, + position_ids, + encoder_hidden_states, + encoder_attention_mask, + deterministic=deterministic, + init_cache=init_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + + if self.config.tie_word_embeddings: + shared_embedding = self.model.variables["params"]["decoder"]["embed_tokens"]["embedding"] + lm_logits = self.lm_head.apply({"params": {"kernel": shared_embedding.T}}, hidden_states) + else: + lm_logits = self.lm_head(hidden_states) + + if not return_dict: + return (lm_logits,) + outputs[1:] + + return FlaxCausalLMOutputWithCrossAttentions( + logits=lm_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + + +@add_start_docstrings( + """ + Bart Decoder Model with a language modeling head on top (linear layer with weights tied to the input embeddings) + e.g for autoregressive tasks. + """, + BART_START_DOCSTRING, +) +class FlaxBartForCausalLM(FlaxBartDecoderPreTrainedModel): + module_class = FlaxBartForCausalLMModule + + def prepare_inputs_for_generation(self, input_ids, max_length, attention_mask: Optional[jnp.DeviceArray] = None): + # initializing the cache + batch_size, seq_length = input_ids.shape + + past_key_values = self.init_cache(batch_size, max_length) + # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length. + # But since the decoder uses a causal mask, those positions are masked anyway. + # Thus, we can create a single static attention_mask here, which is more efficient for compilation + extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4") + if attention_mask is not None: + position_ids = attention_mask.cumsum(axis=-1) - 1 + extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, attention_mask, (0, 0)) + else: + position_ids = jnp.broadcast_to(jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)) + + return { + "past_key_values": past_key_values, + "attention_mask": extended_attention_mask, + "position_ids": position_ids, + } + + def update_inputs_for_generation(self, model_outputs, model_kwargs): + model_kwargs["past_key_values"] = model_outputs.past_key_values + model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1 + return model_kwargs + + +append_call_sample_docstring( + FlaxBartForCausalLM, + _TOKENIZER_FOR_DOC, + _CHECKPOINT_FOR_DOC, + FlaxCausalLMOutputWithCrossAttentions, + _CONFIG_FOR_DOC, +) diff --git a/src/transformers/utils/dummy_flax_objects.py b/src/transformers/utils/dummy_flax_objects.py index 3962cdfb52..166cecaeba 100644 --- a/src/transformers/utils/dummy_flax_objects.py +++ b/src/transformers/utils/dummy_flax_objects.py @@ -249,6 +249,20 @@ class FlaxAutoModelForVision2Seq(metaclass=DummyObject): requires_backends(self, ["flax"]) +class FlaxBartDecoderPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + class FlaxBartForConditionalGeneration(metaclass=DummyObject): _backends = ["flax"] diff --git a/tests/bart/test_modeling_flax_bart.py b/tests/bart/test_modeling_flax_bart.py index dce757e884..219d41cae2 100644 --- a/tests/bart/test_modeling_flax_bart.py +++ b/tests/bart/test_modeling_flax_bart.py @@ -11,17 +11,16 @@ # 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 import numpy as np import timeout_decorator # noqa -from transformers import BartConfig, is_flax_available +from transformers import BartConfig, BartTokenizer, is_flax_available from transformers.testing_utils import require_flax, slow from ..generation.test_generation_flax_utils import FlaxGenerationTesterMixin -from ..test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor +from ..test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask if is_flax_available(): @@ -34,7 +33,6 @@ if is_flax_available(): import jax import jax.numpy as jnp - from transformers import BartTokenizer from transformers.models.bart.modeling_flax_bart import ( FlaxBartForConditionalGeneration, FlaxBartForQuestionAnswering, @@ -475,3 +473,95 @@ class FlaxBartModelTest(FlaxModelTesterMixin, unittest.TestCase, FlaxGenerationT hypotheses_batch.tolist(), clean_up_tokenization_spaces=True, skip_special_tokens=True ) assert generated_summaries == EXPECTED + + +class FlaxBartStandaloneDecoderModelTester: + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_attention_mask=True, + use_labels=False, + vocab_size=99, + hidden_size=16, + num_hidden_layers=2, + num_attention_heads=4, + intermediate_size=4, + hidden_act="gelu", + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=32, + eos_token_id=2, + pad_token_id=1, + bos_token_id=0, + initializer_range=0.02, + ): + self.parent = parent + self.batch_size = batch_size + self.seq_length = seq_length + self.is_training = is_training + self.use_attention_mask = use_attention_mask + 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_act = hidden_act + self.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.initializer_range = initializer_range + + def prepare_config_and_inputs(self): + input_ids = jnp.clip(ids_tensor([self.batch_size, self.seq_length], self.vocab_size), 3, self.vocab_size) + + attention_mask = None + if self.use_attention_mask: + attention_mask = random_attention_mask([self.batch_size, self.seq_length]) + + config = BartConfig( + 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_id=self.eos_token_id, + bos_token_id=self.bos_token_id, + pad_token_id=self.pad_token_id, + initializer_range=self.initializer_range, + use_cache=False, + ) + + return config, input_ids, attention_mask + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + config, input_ids, attention_mask = config_and_inputs + inputs_dict = {"input_ids": input_ids, "attention_mask": attention_mask} + return config, inputs_dict + + def prepare_config_and_inputs_for_decoder(self): + config, input_ids, attention_mask = self.prepare_config_and_inputs() + + encoder_hidden_states = floats_tensor([self.batch_size, self.seq_length, self.hidden_size]) + encoder_attention_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2) + + return ( + config, + input_ids, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + ) diff --git a/tests/encoder_decoder/test_modeling_flax_encoder_decoder.py b/tests/encoder_decoder/test_modeling_flax_encoder_decoder.py index 60be9f420c..e6f0a49c16 100644 --- a/tests/encoder_decoder/test_modeling_flax_encoder_decoder.py +++ b/tests/encoder_decoder/test_modeling_flax_encoder_decoder.py @@ -22,6 +22,7 @@ import numpy as np from transformers import is_flax_available, is_torch_available from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow, torch_device +from ..bart.test_modeling_flax_bart import FlaxBartStandaloneDecoderModelTester from ..bert.test_modeling_flax_bert import FlaxBertModelTester from ..gpt2.test_modeling_flax_gpt2 import FlaxGPT2ModelTester from ..test_modeling_flax_common import ids_tensor @@ -31,6 +32,7 @@ if is_flax_available(): from transformers import ( AutoTokenizer, EncoderDecoderConfig, + FlaxBartForCausalLM, FlaxBertModel, FlaxEncoderDecoderModel, FlaxGPT2LMHeadModel, @@ -360,6 +362,7 @@ class FlaxEncoderDecoderMixin: self.assertTrue(decoder_config.cross_attention_hidden_size is None) # check without `enc_to_dec_proj` projection + decoder_config.hidden_size = config.hidden_size self.assertTrue(config.hidden_size == decoder_config.hidden_size) self.check_equivalence_pt_to_flax(config, decoder_config, inputs_dict) self.check_equivalence_flax_to_pt(config, decoder_config, inputs_dict) @@ -456,6 +459,43 @@ class FlaxGPT2EncoderDecoderModelTest(FlaxEncoderDecoderMixin, unittest.TestCase self.assertEqual(summary, [EXPECTED_SUMMARY_STUDENTS]) +@require_flax +class FlaxBartEncoderDecoderModelTest(FlaxEncoderDecoderMixin, unittest.TestCase): + def get_encoder_decoder_model(self, config, decoder_config): + encoder_model = FlaxBertModel(config) + decoder_model = FlaxBartForCausalLM(decoder_config) + return encoder_model, decoder_model + + def prepare_config_and_inputs(self): + model_tester_encoder = FlaxBertModelTester(self, batch_size=13) + model_tester_decoder = FlaxBartStandaloneDecoderModelTester(self, batch_size=13) + encoder_config_and_inputs = model_tester_encoder.prepare_config_and_inputs() + decoder_config_and_inputs = model_tester_decoder.prepare_config_and_inputs_for_decoder() + (config, input_ids, token_type_ids, attention_mask) = encoder_config_and_inputs + ( + decoder_config, + decoder_input_ids, + decoder_attention_mask, + encoder_hidden_states, + encoder_attention_mask, + ) = decoder_config_and_inputs + + # make sure that cross attention layers are added + decoder_config.add_cross_attention = True + return { + "config": config, + "input_ids": input_ids, + "attention_mask": attention_mask, + "decoder_config": decoder_config, + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + "encoder_hidden_states": encoder_hidden_states, + } + + def get_pretrained_model(self): + return FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("bert-base-cased", "facebook/bart-base") + + @require_flax class FlaxEncoderDecoderModelTest(unittest.TestCase): def get_from_encoderdecoder_pretrained_model(self): diff --git a/tests/speech_encoder_decoder/test_modeling_flax_speech_encoder_decoder.py b/tests/speech_encoder_decoder/test_modeling_flax_speech_encoder_decoder.py index 51868a851d..f204dae530 100644 --- a/tests/speech_encoder_decoder/test_modeling_flax_speech_encoder_decoder.py +++ b/tests/speech_encoder_decoder/test_modeling_flax_speech_encoder_decoder.py @@ -21,6 +21,7 @@ import numpy as np from transformers import is_flax_available, is_torch_available from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow, torch_device +from ..bart.test_modeling_flax_bart import FlaxBartStandaloneDecoderModelTester from ..gpt2.test_modeling_flax_gpt2 import FlaxGPT2ModelTester from ..test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask from ..wav2vec2.test_modeling_flax_wav2vec2 import FlaxWav2Vec2ModelTester @@ -28,6 +29,7 @@ from ..wav2vec2.test_modeling_flax_wav2vec2 import FlaxWav2Vec2ModelTester if is_flax_available(): from transformers import ( + FlaxBartForCausalLM, FlaxGPT2LMHeadModel, FlaxSpeechEncoderDecoderModel, FlaxWav2Vec2Model, @@ -553,3 +555,120 @@ class FlaxWav2Vec2GPT2ModelTest(FlaxEncoderDecoderMixin, unittest.TestCase): self.assertEqual(len(fx_outputs), len(pt_outputs_loaded), "Output lengths differ between Flax and PyTorch") self.assert_almost_equals(fx_logits, pt_logits_loaded.numpy(), 4e-2) + + +@require_flax +class FlaxWav2Vec2BartModelTest(FlaxEncoderDecoderMixin, unittest.TestCase): + def get_pretrained_model_and_inputs(self): + model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained( + "facebook/wav2vec2-large-lv60", "bart-large" + ) + batch_size = 13 + input_values = floats_tensor([batch_size, 512], model.config.encoder.vocab_size) + attention_mask = random_attention_mask([batch_size, 512]) + decoder_input_ids = ids_tensor([batch_size, 4], model.config.decoder.vocab_size) + decoder_attention_mask = random_attention_mask([batch_size, 4]) + inputs = { + "inputs": input_values, + "attention_mask": attention_mask, + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + } + + return model, inputs + + def get_encoder_decoder_model(self, config, decoder_config): + encoder_model = FlaxWav2Vec2Model(config) + decoder_model = FlaxBartForCausalLM(decoder_config) + return encoder_model, decoder_model + + def prepare_config_and_inputs(self): + model_tester_encoder = FlaxWav2Vec2ModelTester(self, batch_size=13) + model_tester_decoder = FlaxBartStandaloneDecoderModelTester(self, batch_size=13) + encoder_config_and_inputs = model_tester_encoder.prepare_config_and_inputs() + decoder_config_and_inputs = model_tester_decoder.prepare_config_and_inputs_for_decoder() + (config, inputs, attention_mask) = encoder_config_and_inputs + ( + decoder_config, + decoder_input_ids, + decoder_attention_mask, + encoder_hidden_states, + encoder_attention_mask, + ) = decoder_config_and_inputs + + # make sure that cross attention layers are added + decoder_config.add_cross_attention = True + return { + "config": config, + "inputs": inputs, + "attention_mask": attention_mask, + "decoder_config": decoder_config, + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + "encoder_hidden_states": encoder_hidden_states, + } + + @slow + def test_flaxwav2vec2bart_pt_flax_equivalence(self): + pt_model = SpeechEncoderDecoderModel.from_pretrained("patrickvonplaten/wav2vec2-2-bart-large") + fx_model = FlaxSpeechEncoderDecoderModel.from_pretrained( + "patrickvonplaten/wav2vec2-2-bart-large", from_pt=True + ) + + pt_model.to(torch_device) + pt_model.eval() + + # prepare inputs + batch_size = 13 + input_values = floats_tensor([batch_size, 512], fx_model.config.encoder.vocab_size) + attention_mask = random_attention_mask([batch_size, 512]) + decoder_input_ids = ids_tensor([batch_size, 4], fx_model.config.decoder.vocab_size) + decoder_attention_mask = random_attention_mask([batch_size, 4]) + inputs_dict = { + "inputs": input_values, + "attention_mask": attention_mask, + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + } + + flax_inputs = inputs_dict + pt_inputs = {k: torch.tensor(v.tolist()) for k, v in flax_inputs.items()} + + with torch.no_grad(): + pt_outputs = pt_model(**pt_inputs) + pt_logits = pt_outputs.logits + pt_outputs = pt_outputs.to_tuple() + + fx_outputs = fx_model(**inputs_dict) + fx_logits = fx_outputs.logits + fx_outputs = fx_outputs.to_tuple() + + self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch") + self.assert_almost_equals(fx_logits, pt_logits.numpy(), 4e-2) + + # PT -> Flax + with tempfile.TemporaryDirectory() as tmpdirname: + pt_model.save_pretrained(tmpdirname) + fx_model_loaded = FlaxSpeechEncoderDecoderModel.from_pretrained(tmpdirname, from_pt=True) + + fx_outputs_loaded = fx_model_loaded(**inputs_dict) + fx_logits_loaded = fx_outputs_loaded.logits + fx_outputs_loaded = fx_outputs_loaded.to_tuple() + self.assertEqual(len(fx_outputs_loaded), len(pt_outputs), "Output lengths differ between Flax and PyTorch") + self.assert_almost_equals(fx_logits_loaded, pt_logits.numpy(), 4e-2) + + # Flax -> PT + with tempfile.TemporaryDirectory() as tmpdirname: + fx_model.save_pretrained(tmpdirname) + pt_model_loaded = SpeechEncoderDecoderModel.from_pretrained(tmpdirname, from_flax=True) + + pt_model_loaded.to(torch_device) + pt_model_loaded.eval() + + with torch.no_grad(): + pt_outputs_loaded = pt_model_loaded(**pt_inputs) + pt_logits_loaded = pt_outputs_loaded.logits + pt_outputs_loaded = pt_outputs_loaded.to_tuple() + + self.assertEqual(len(fx_outputs), len(pt_outputs_loaded), "Output lengths differ between Flax and PyTorch") + self.assert_almost_equals(fx_logits, pt_logits_loaded.numpy(), 4e-2) diff --git a/utils/check_repo.py b/utils/check_repo.py index 46fe871ef0..308a853113 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -89,6 +89,7 @@ IGNORE_NON_TESTED = PRIVATE_MODELS.copy() + [ "TFRobertaForMultipleChoice", # TODO: fix "TrOCRDecoderWrapper", # Building part of bigger (tested) model. "SeparableConv1D", # Building part of bigger (tested) model. + "FlaxBartForCausalLM", # Building part of bigger (tested) model. ] # Update this list with test files that don't have a tester with a `all_model_classes` variable and which don't