Add flags to return scores, hidden states and / or attention weights in GenerationMixin (#9150)

* Define new output dataclasses for greedy generation

* Add output_[...] flags in greedy generation methods

Added output_attentions, output_hidden_states, output_scores flags in
generate and greedy_search methods in GenerationMixin.

* [WIP] Implement logic and tests for output flags in generation

* Update GreedySearchOutput classes & docstring

* Implement greedy search output accumulation logic

Update greedy_search unittests

Fix generate method return value docstring

Properly init flags with the default config

* Update configuration to add output_scores flag

* Fix test_generation_utils

Sort imports and fix isinstance tests for GreedySearchOutputs

* Fix typo in generation_utils

* Add return_dict_in_generate for backwards compatibility

* Add return_dict_in_generate flag in config

* Fix tyPo in configuration

* Fix handling of attentions and hidden_states flags

* Make style & quality

* first attempt attentions

* some corrections

* improve tests

* special models requires special test

* disable xlm test for now

* clean tests

* fix for tf

* isort

* Add output dataclasses for other generation methods

* Add logic to return dict in sample generation

* Complete test for sample generation

- Pass output_attentions and output_hidden_states flags to encoder in
encoder-decoder models
- Fix import satements order in test_generation_utils file

* Add logic to return dict in sample generation

- Refactor tests to avoid using self.assertTrue, which provides
scarce information when the test fails
- Add tests for the three beam_search methods: vanilla, sample and
grouped

* Style doc

* Fix copy-paste error in generation tests

* Rename logits to scores and refactor

* Refactor group_beam_search for consistency

* make style

* add sequences_scores

* fix all tests

* add docs

* fix beam search finalize test

* correct docstring

* clean some files

* Made suggested changes to the documentation

* Style doc ?

* Style doc using the Python util

* Update src/transformers/generation_utils.py

* fix empty lines

* fix all test

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

tests/test_generation_beam_search.py

@@ -190,7 +190,7 @@ class BeamSearchTester:
         input_ids = torch.cat([input_ids[output_indices, :], output_tokens.unsqueeze(-1)], dim=-1)
 
         # finalize
-        decoded = beam_scorer.finalize(
+        sequence_output = beam_scorer.finalize(
             input_ids,
             output_scores,
             output_tokens,
@@ -198,19 +198,27 @@ class BeamSearchTester:
             pad_token_id=self.pad_token_id,
             eos_token_id=self.eos_token_id,
         )
+
+        sequences = sequence_output["sequences"]
+        sequence_scores = sequence_output["sequence_scores"]
+
         # since `num_beam_hyps_to_keep` = 1 => only return `batch_size` x `max_length`
-        self.parent.assertListEqual(list(decoded.shape), [self.batch_size, max_length])
+        self.parent.assertListEqual(list(sequences.shape), [self.batch_size, max_length])
+        self.parent.assertListEqual(list(sequence_scores.shape), [self.batch_size])
+
+        # check sequence_scores
+        self.parent.assertFalse((sequence_scores > 0).any().item())
 
         # first batch has to finish with eos_token
-        self.parent.assertEqual(decoded[0, -1].item(), self.eos_token_id)
+        self.parent.assertEqual(sequences[0, -1].item(), self.eos_token_id)
 
         # other batches cannot finish with eos token
-        self.parent.assertNotEqual(decoded[1, -1].item(), self.eos_token_id)
-        self.parent.assertNotEqual(decoded[2, -1].item(), self.eos_token_id)
+        self.parent.assertNotEqual(sequences[1, -1].item(), self.eos_token_id)
+        self.parent.assertNotEqual(sequences[2, -1].item(), self.eos_token_id)
 
         # now test that if `num_beam_hyps_to_keep` is 3 => all beams are returned
         beam_scorer.num_beam_hyps_to_keep = self.num_beams
-        decoded = beam_scorer.finalize(
+        sequence_output = beam_scorer.finalize(
             input_ids,
             output_scores,
             output_tokens,
@@ -218,7 +226,11 @@ class BeamSearchTester:
             pad_token_id=self.pad_token_id,
             eos_token_id=self.eos_token_id,
         )
-        self.parent.assertListEqual(list(decoded.shape), [self.num_beams * self.batch_size, max_length])
+        sequences = sequence_output["sequences"]
+        sequence_scores = sequence_output["sequence_scores"]
+
+        self.parent.assertListEqual(list(sequences.shape), [self.num_beams * self.batch_size, max_length])
+        self.parent.assertListEqual(list(sequence_scores.shape), [self.num_beams * self.batch_size])
 
 
 @require_torch