From b8462b5b2ac84f63293900ae168dbde039443a22 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Wed, 6 Jan 2021 19:37:02 +0100 Subject: [PATCH] [GenerationOutputs] Fix GenerationOutputs Tests (#9443) * fix generation models * fix led * fix docs * add is_decoder * fix last docstrings * make style * fix t5 cross attentions * correct t5 --- src/transformers/generation_utils.py | 72 +++++++++---------- src/transformers/models/bert/modeling_bert.py | 2 +- .../modeling_bert_generation.py | 2 +- .../modeling_encoder_decoder.py | 2 +- .../models/roberta/modeling_roberta.py | 2 +- src/transformers/models/t5/modeling_t5.py | 2 + ...ng_{{cookiecutter.lowercase_modelname}}.py | 2 +- tests/test_generation_utils.py | 32 ++++++--- tests/test_modeling_led.py | 26 +++++++ 9 files changed, 89 insertions(+), 53 deletions(-) diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py index fb55b97859..8147c103dd 100644 --- a/src/transformers/generation_utils.py +++ b/src/transformers/generation_utils.py @@ -113,21 +113,20 @@ class SampleDecoderOnlyOutput(ModelOutput): Args: - sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`): + sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size*num_return_sequences, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all batches finished early due to the :obj:`eos_token_id`. scores (:obj:`tuple(torch.FloatTensor)` `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) at each generation step. :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of - shape :obj:`(batch_size * num_return_sequences, config.vocab_size)`). + shape :obj:`(batch_size*num_return_sequences, config.vocab_size)`). attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(num_return_sequences * batch_size, num_heads, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(num_return_sequences*batch_size, num_heads, generated_length, sequence_length)`. hidden_states (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(num_return_sequences * batch_size, generated_length, - hidden_size)`. + :obj:`torch.FloatTensor` of shape :obj:`(num_return_sequences*batch_size, generated_length, hidden_size)`. """ sequences: torch.LongTensor = None @@ -145,27 +144,26 @@ class SampleEncoderDecoderOutput(ModelOutput): Args: - sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`): + sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size*num_return_sequences, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all batches finished early due to the :obj:`eos_token_id`. scores (:obj:`tuple(torch.FloatTensor)` `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) at each generation step. :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of - shape :obj:`(batch_size * num_return_sequences, config.vocab_size)`). + shape :obj:`(batch_size*num_return_sequences, config.vocab_size)`). encoder_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): - Tuple of :obj:`torch.FloatTensor` (one for each layer of the decoder) of shape :obj:`(batch_size * - num_return_sequences, num_heads, sequence_length, sequence_length)`. + Tuple of :obj:`torch.FloatTensor` (one for each layer of the decoder) of shape + :obj:`(batch_size*num_return_sequences, num_heads, sequence_length, sequence_length)`. encoder_hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) - of shape :obj:`(batch_size * num_return_sequences, sequence_length, hidden_size)`. + of shape :obj:`(batch_size*num_return_sequences, sequence_length, hidden_size)`. decoder_attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_return_sequences, num_heads, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_return_sequences, num_heads, generated_length, sequence_length)`. decoder_hidden_states (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_return_sequences, generated_length, - hidden_size)`. + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_return_sequences, generated_length, hidden_size)`. """ sequences: torch.LongTensor = None @@ -182,23 +180,23 @@ class BeamSearchDecoderOnlyOutput(ModelOutput): Base class for outputs of decoder-only generation models using beam search. Args: - sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`): + sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size*num_return_sequences, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all batches finished early due to the :obj:`eos_token_id`. - sequences_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_return_sequences)`, `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): + sequences_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_return_sequences)`, `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Final beam scores of the generated ``sequences``. scores (:obj:`tuple(torch.FloatTensor)` `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this beam - . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape :obj:`(batch_size - * num_beams * num_return_sequences, config.vocab_size)`). + . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape + :obj:`(batch_size*num_beams*num_return_sequences, config.vocab_size)`). attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams, num_heads, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams, num_heads, generated_length, sequence_length)`. hidden_states (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams * num_return_sequences, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`. """ @@ -217,30 +215,30 @@ class BeamSearchEncoderDecoderOutput(ModelOutput): attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) Args: - sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`): + sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size*num_return_sequences, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all batches finished early due to the :obj:`eos_token_id`. - sequences_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_return_sequences)`, `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): + sequences_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_return_sequences)`, `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Final beam scores of the generated ``sequences``. scores (:obj:`tuple(torch.FloatTensor)` `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this beam - . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape :obj:`(batch_size - * num_beams, config.vocab_size)`). + . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape + :obj:`(batch_size*num_beams, config.vocab_size)`). attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): encoder_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple of :obj:`torch.FloatTensor` (one for each layer of the decoder) of shape :obj:`(batch_size, num_heads, sequence_length, sequence_length)`. encoder_hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) - of shape :obj:`(batch_size * num_beams * num_return_sequences, sequence_length, hidden_size)`. + of shape :obj:`(batch_size*num_beams*num_return_sequences, sequence_length, hidden_size)`. decoder_attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams * num_return_sequences, num_heads, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams*num_return_sequences, num_heads, generated_length, sequence_length)`. decoder_hidden_states (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams * num_return_sequences, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`. """ @@ -259,7 +257,7 @@ class BeamSampleDecoderOnlyOutput(ModelOutput): Base class for outputs of decoder-only generation models using beam sample. Args: - sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`): + sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size*num_return_sequences, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all batches finished early due to the :obj:`eos_token_id`. sequences_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_return_sequence)`, `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): @@ -267,15 +265,15 @@ class BeamSampleDecoderOnlyOutput(ModelOutput): scores (:obj:`tuple(torch.FloatTensor)` `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this beam - . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape :obj:`(batch_size - * num_beams * num_return_sequences, config.vocab_size)`). + . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape + :obj:`(batch_size*num_beams*num_return_sequences, config.vocab_size)`). attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams, num_heads, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams, num_heads, generated_length, sequence_length)`. hidden_states (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams, generated_length, hidden_size)`. + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams, generated_length, hidden_size)`. """ sequences: torch.LongTensor = None @@ -293,7 +291,7 @@ class BeamSampleEncoderDecoderOutput(ModelOutput): encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) Args: - sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size * num_beams, sequence_length)`): + sequences (:obj:`torch.LongTensor` of shape :obj:`(batch_size*num_beams, sequence_length)`): The generated sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all batches finished early due to the :obj:`eos_token_id`. sequences_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_return_sequence)`, `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): @@ -301,21 +299,21 @@ class BeamSampleEncoderDecoderOutput(ModelOutput): scores (:obj:`tuple(torch.FloatTensor)` `optional`, returned when ``output_scores=True`` is passed or when ``config.output_scores=True``): Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this beam - . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape :obj:`(batch_size - * num_beams, config.vocab_size)`). + . :obj:`(max_length,)`-shaped tuple of :obj:`torch.FloatTensor` with each tensor of shape + :obj:`(batch_size*num_beams, config.vocab_size)`). encoder_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple of :obj:`torch.FloatTensor` (one for each layer of the decoder) of shape :obj:`(batch_size, num_heads, sequence_length, sequence_length)`. encoder_hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) - of shape :obj:`(batch_size * num_beams, sequence_length, hidden_size)`. + of shape :obj:`(batch_size*num_beams, sequence_length, hidden_size)`. decoder_attentions (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_attentions=True`` is passed or ``config.output_attentions=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams, num_heads, generated_length, + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams, num_heads, generated_length, sequence_length)`. decoder_hidden_states (:obj:`tuple(tuple(torch.FloatTensor))`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - :obj:`torch.FloatTensor` of shape :obj:`(batch_size * num_beams, generated_length, hidden_size)`. + :obj:`torch.FloatTensor` of shape :obj:`(batch_size*num_beams, generated_length, hidden_size)`. """ sequences: torch.LongTensor = None diff --git a/src/transformers/models/bert/modeling_bert.py b/src/transformers/models/bert/modeling_bert.py index b4d11fda43..590363a124 100755 --- a/src/transformers/models/bert/modeling_bert.py +++ b/src/transformers/models/bert/modeling_bert.py @@ -1227,7 +1227,7 @@ class BertLMHeadModel(BertPreTrainedModel): if past is not None: input_ids = input_ids[:, -1:] - return {"input_ids": input_ids, "attention_mask": attention_mask} + return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past} def _reorder_cache(self, past, beam_idx): reordered_past = () diff --git a/src/transformers/models/bert_generation/modeling_bert_generation.py b/src/transformers/models/bert_generation/modeling_bert_generation.py index 35e6d15f93..0f89efb385 100755 --- a/src/transformers/models/bert_generation/modeling_bert_generation.py +++ b/src/transformers/models/bert_generation/modeling_bert_generation.py @@ -570,7 +570,7 @@ class BertGenerationDecoder(BertGenerationPreTrainedModel): if past is not None: input_ids = input_ids[:, -1:] - return {"input_ids": input_ids, "attention_mask": attention_mask} + return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past} def _reorder_cache(self, past, beam_idx): reordered_past = () diff --git a/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py b/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py index e6b19842ba..2486521847 100644 --- a/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py +++ b/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py @@ -455,7 +455,7 @@ class EncoderDecoderModel(PreTrainedModel): "decoder_attention_mask": decoder_attention_mask, "decoder_input_ids": decoder_inputs["input_ids"], "encoder_outputs": encoder_outputs, - "past_key_values": past, + "past_key_values": decoder_inputs["past_key_values"], "use_cache": use_cache, } return input_dict diff --git a/src/transformers/models/roberta/modeling_roberta.py b/src/transformers/models/roberta/modeling_roberta.py index 3e9eca1767..2f648d3dc0 100644 --- a/src/transformers/models/roberta/modeling_roberta.py +++ b/src/transformers/models/roberta/modeling_roberta.py @@ -962,7 +962,7 @@ class RobertaForCausalLM(RobertaPreTrainedModel): if past is not None: input_ids = input_ids[:, -1:] - return {"input_ids": input_ids, "attention_mask": attention_mask} + return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past} def _reorder_cache(self, past, beam_idx): reordered_past = () diff --git a/src/transformers/models/t5/modeling_t5.py b/src/transformers/models/t5/modeling_t5.py index 7e292a6882..3b096c36db 100644 --- a/src/transformers/models/t5/modeling_t5.py +++ b/src/transformers/models/t5/modeling_t5.py @@ -1148,6 +1148,7 @@ class T5Model(T5PreTrainedModel): self.shared = nn.Embedding(config.vocab_size, config.d_model) encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False encoder_config.use_cache = False encoder_config.is_encoder_decoder = False self.encoder = T5Stack(encoder_config, self.shared) @@ -1325,6 +1326,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel): self.shared = nn.Embedding(config.vocab_size, config.d_model) encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False encoder_config.use_cache = False encoder_config.is_encoder_decoder = False self.encoder = T5Stack(encoder_config, self.shared) 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 77c8847a39..c3b02e905d 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 @@ -1132,7 +1132,7 @@ class {{cookiecutter.camelcase_modelname}}ForCausalLM({{cookiecutter.camelcase_m if past is not None: input_ids = input_ids[:, -1:] - return {"input_ids": input_ids, "attention_mask": attention_mask} + return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past} def _reorder_cache(self, past, beam_idx): reordered_past = () diff --git a/tests/test_generation_utils.py b/tests/test_generation_utils.py index 5359f348f8..dac8337d52 100644 --- a/tests/test_generation_utils.py +++ b/tests/test_generation_utils.py @@ -522,6 +522,7 @@ class GenerationTesterMixin: return config.use_cache = True + config.is_decoder = True model = model_class(config).to(torch_device).eval() output_greedy, output_generate = self._greedy_generate( model=model, @@ -730,6 +731,7 @@ class GenerationTesterMixin: beam_kwargs, beam_scorer = self._get_beam_scorer_and_kwargs(input_ids.shape[0], max_length) config.use_cache = True + config.is_decoder = True model = model_class(config).to(torch_device).eval() output_beam, output_generate = self._beam_search_generate( model=model, @@ -962,12 +964,7 @@ class GenerationTesterMixin: # Attentions if config.is_encoder_decoder: # encoder - encoder_expected_shape = (batch_size, config.num_attention_heads, seq_length, seq_length) - self.assertIsInstance(output.encoder_attentions, tuple) - self.assertListEqual( - [layer_attentions.shape for layer_attentions in output.encoder_attentions], - [encoder_expected_shape] * len(output.encoder_attentions), - ) + self._check_encoder_attention_for_generate(output.encoder_attentions, batch_size, config, seq_length) # decoder self._check_attentions_for_generate( num_sequences_in_output, @@ -993,11 +990,8 @@ class GenerationTesterMixin: # Hidden States if config.is_encoder_decoder: # encoder - encoder_expected_shape = (batch_size, seq_length, config.hidden_size) - self.assertIsInstance(output.encoder_hidden_states, tuple) - self.assertListEqual( - [layer_hidden_states.shape for layer_hidden_states in output.encoder_hidden_states], - [encoder_expected_shape] * len(output.encoder_hidden_states), + self._check_encoder_hidden_states_for_generate( + output.encoder_hidden_states, batch_size, config, seq_length ) # decoder @@ -1052,6 +1046,14 @@ class GenerationTesterMixin: [layer_attention.shape for layer_attention in iter_attentions], [expected_shape] * len(iter_attentions) ) + def _check_encoder_attention_for_generate(self, attentions, batch_size, config, seq_length): + encoder_expected_shape = (batch_size, config.num_attention_heads, seq_length, seq_length) + self.assertIsInstance(attentions, tuple) + self.assertListEqual( + [layer_attentions.shape for layer_attentions in attentions], + [encoder_expected_shape] * len(attentions), + ) + def _check_hidden_states_for_generate( self, batch_size, hidden_states, min_length, max_length, config, use_cache=False, num_beam_groups=1 ): @@ -1071,6 +1073,14 @@ class GenerationTesterMixin: [expected_shape] * len(iter_hidden_states), ) + def _check_encoder_hidden_states_for_generate(self, hidden_states, batch_size, config, seq_length): + encoder_expected_shape = (batch_size, seq_length, config.hidden_size) + self.assertIsInstance(hidden_states, tuple) + self.assertListEqual( + [layer_hidden_states.shape for layer_hidden_states in hidden_states], + [encoder_expected_shape] * len(hidden_states), + ) + @require_torch class UtilsFunctionsTest(unittest.TestCase): diff --git a/tests/test_modeling_led.py b/tests/test_modeling_led.py index 915e42f394..55b63dd24b 100644 --- a/tests/test_modeling_led.py +++ b/tests/test_modeling_led.py @@ -327,6 +327,32 @@ class LEDModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase): # longformer cannot keep gradients in attentions or hidden states return + def _check_encoder_attention_for_generate(self, attentions, batch_size, config, seq_length): + # make sure tgt_length is padded + tgt_length = ( + seq_length // config.attention_window[0] + (seq_length % config.attention_window[0] != 0) + ) * config.attention_window[0] + + encoder_expected_shape = (batch_size, config.num_attention_heads, tgt_length, seq_length) + self.assertIsInstance(attentions, tuple) + self.assertListEqual( + [layer_attentions.shape for layer_attentions in attentions], + [encoder_expected_shape] * len(attentions), + ) + + def _check_encoder_hidden_states_for_generate(self, hidden_states, batch_size, config, seq_length): + # make sure seq_length is padded + seq_length = ( + seq_length // config.attention_window[0] + (seq_length % config.attention_window[0] != 0) + ) * config.attention_window[0] + + encoder_expected_shape = (batch_size, seq_length, config.hidden_size) + self.assertIsInstance(hidden_states, tuple) + self.assertListEqual( + [layer_hidden_states.shape for layer_hidden_states in hidden_states], + [encoder_expected_shape] * len(hidden_states), + ) + def test_attention_outputs(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() config.return_dict = True