Fix bug in x-attentions output for roberta and harden test to catch it (huggingface#8660)

Yossi Synett · web-flow · commit 18c8cf000bed · 2020-11-23T13:28:29.000+01:00
diff --git a/src/transformers/models/roberta/modeling_roberta.py b/src/transformers/models/roberta/modeling_roberta.py
@@ -814,7 +814,7 @@ def forward(
             logits=prediction_scores,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-            cross_attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
         )
 
     def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
diff --git a/tests/test_modeling_encoder_decoder.py b/tests/test_modeling_encoder_decoder.py
@@ -300,6 +300,9 @@ def check_encoder_decoder_model_output_attentions(
         labels,
         **kwargs
     ):
+        # make the decoder inputs a different shape from the encoder inputs to harden the test
+        decoder_input_ids = decoder_input_ids[:, :-1]
+        decoder_attention_mask = decoder_attention_mask[:, :-1]
         encoder_model, decoder_model = self.get_encoder_decoder_model(config, decoder_config)
         enc_dec_model = EncoderDecoderModel(encoder=encoder_model, decoder=decoder_model)
         enc_dec_model.to(torch_device)
@@ -314,9 +317,8 @@ def check_encoder_decoder_model_output_attentions(
         encoder_attentions = outputs_encoder_decoder["encoder_attentions"]
         self.assertEqual(len(encoder_attentions), config.num_hidden_layers)
 
-        self.assertListEqual(
-            list(encoder_attentions[0].shape[-3:]),
-            [config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1]],
+        self.assertEqual(
+            encoder_attentions[0].shape[-3:], (config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])
         )
 
         decoder_attentions = outputs_encoder_decoder["decoder_attentions"]
@@ -327,20 +329,20 @@ def check_encoder_decoder_model_output_attentions(
         )
         self.assertEqual(len(decoder_attentions), num_decoder_layers)
 
-        self.assertListEqual(
-            list(decoder_attentions[0].shape[-3:]),
-            [decoder_config.num_attention_heads, decoder_input_ids.shape[-1], decoder_input_ids.shape[-1]],
+        self.assertEqual(
+            decoder_attentions[0].shape[-3:],
+            (decoder_config.num_attention_heads, decoder_input_ids.shape[-1], decoder_input_ids.shape[-1]),
         )
 
         cross_attentions = outputs_encoder_decoder["cross_attentions"]
         self.assertEqual(len(cross_attentions), num_decoder_layers)
 
-        cross_attention_input_seq_len = input_ids.shape[-1] * (
+        cross_attention_input_seq_len = decoder_input_ids.shape[-1] * (
             1 + (decoder_config.ngram if hasattr(decoder_config, "ngram") else 0)
         )
-        self.assertListEqual(
-            list(cross_attentions[0].shape[-3:]),
-            [decoder_config.num_attention_heads, cross_attention_input_seq_len, decoder_input_ids.shape[-1]],
+        self.assertEqual(
+            cross_attentions[0].shape[-3:],
+            (decoder_config.num_attention_heads, cross_attention_input_seq_len, input_ids.shape[-1]),
         )
 
     def check_encoder_decoder_model_generate(self, input_ids, config, decoder_config, **kwargs):

Original file line number	Diff line number	Diff line change
`@@ -814,7 +814,7 @@ def forward(`
`814`	`814`	`logits=prediction_scores,`
`815`	`815`	`hidden_states=outputs.hidden_states,`
`816`	`816`	`attentions=outputs.attentions,`
`817`		`- cross_attentions=outputs.attentions,`
	`817`	`+ cross_attentions=outputs.cross_attentions,`
`818`	`818`	`)`
`819`	`819`
`820`	`820`	`def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):`