def check_encoder_decoder_model_from_pretrained_configs(
self, config, input_ids, attention_mask, encoder_hidden_states,
decoder_config, decoder_input_ids, decoder_attention_mask,
**kwargs):
encoder_decoder_config = EncoderDecoderConfig.from_encoder_decoder_configs(
config, decoder_config)
self.assertTrue(encoder_decoder_config.decoder.is_decoder)
enc_dec_model = FlaxEncoderDecoderModel(encoder_decoder_config)
self.assertTrue(enc_dec_model.config.is_encoder_decoder)
outputs_encoder_decoder = enc_dec_model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
self.assertEqual(outputs_encoder_decoder["logits"].shape,
(decoder_input_ids.shape +
(decoder_config.vocab_size, )))
self.assertEqual(
outputs_encoder_decoder["encoder_last_hidden_state"].shape,
(input_ids.shape + (config.hidden_size, )))
def check_equivalence_flax_to_pt(self, config, decoder_config, inputs_dict):
encoder_decoder_config = EncoderDecoderConfig.from_encoder_decoder_configs(config, decoder_config)
pt_model = EncoderDecoderModel(encoder_decoder_config)
fx_model = FlaxEncoderDecoderModel(encoder_decoder_config)
pt_model = load_flax_weights_in_pytorch_model(pt_model, fx_model.params)
self.check_pt_flax_equivalence(pt_model, fx_model, inputs_dict)
def check_equivalence_pt_to_flax(self, config, decoder_config, inputs_dict):
encoder_decoder_config = EncoderDecoderConfig.from_encoder_decoder_configs(config, decoder_config)
pt_model = EncoderDecoderModel(encoder_decoder_config)
fx_model = FlaxEncoderDecoderModel(encoder_decoder_config)
fx_state = convert_pytorch_state_dict_to_flax(pt_model.state_dict(), fx_model)
fx_model.params = fx_state
self.check_pt_flax_equivalence(pt_model, fx_model, inputs_dict)