TF generate refactor - Greedy Search
Author: patrickvonplaten
Pull request: yes

# What does this PR do?

This PR is the first step to refactor TF's `generate()` method similar in spirit to https://discuss.huggingface.co/t/big-generate-refactor/1857 . 

## Motivation

The main reasons for the refactor are:
- Make `generate` more readable and easier to understand. 
- Disentangle different components of `generate`, *e.g.* the model_input preparation, the logits processor creation and application, the sub-generation methods. Each of these components of `generate` should be as independent and extendable as possible. We've seen that such a design greatly helps to foster community contributions to the generate method for PyTorch.
- Make `generate` more generally applicable for future use cases. This mostly concerns multi-modal models, such as speech to text and image to text. 

In addition the final goal is to significantly speed up TF's generate method by making it XLA-compatible. This PR should greatly help at making the `generate()` method compatible with XLA by:
- clearly separating the model inputs preparation part and the auto-regressive part. The auto-regressive part is handled in the sub-generation method which should now be much shorter. This is the part which will require some heavy changes to work with XLA. 
- Moving out all of the logits processing logic into a `LogitsProcessorList` class. We can now simply test each of those classes separately whether they are XLA-compatible and don't have to make each of them XLA-compatible. E.g. `bad_token_ids` and `ngram_repeat` are probably not at all XLA-compatible.

## Proposed refactor design and steps to achieve it

TF's `generate` can currently handle more or less **4** different types of generation:
- a) greedy search
- b) sampling (GPT-like generation)
- c) beam search (default use case for translation, etc...)
- d) beam sample (exotic case where sampling is combined with `beam_search`

Currently a) and b) are handled by the sub-method `_generate_no_beam_search` while c) and d) are handled by `_generate_beam_search` .
This is not very readable and understandable. I propose to remove those sub-methods and instead define all four generation methods a), b), c) and d) separately. 

I propose to handle the complete generate refactor in multiple steps.
1) Define the new `generate()` design only for the simplest use case being `greedy_search` (a)). Leave all existing code for now to not break any of b), c), d)
2) At this point we can already start change greedy search to be XLA-compatible.
3) Analogues to 1), we define the sub-method b) and can then fully remove `_generate_no_beam_search`
4) Refactor c) beam search. This is a bit more complex and I'd like to handle this in a second step.
5) Finally we can refactor d) which should be relatively simple after 4) and can then remove the `_generate_beam_search` function.

This first PR lays the foundation for the complete refactor and finishes 1).

In short in does the following:

- i.) Adds more aggressive `greedy_search` tests for a decoder-only model - GPT2 - and an encoder-decoder model - T5.
- ii.) Define the TF logic for logits processors
- iii.) Define the new `generate()` method which is in a first step called `_generate()` so that it will only be used for a) `greedy_search`. It should be much more readable by adding clear `# 1. ... 2. ` comments for each model input preparation step.
- iv.) Fully refactor `greedy_search`
- v.) Clean-up some duplicated code that was heavily used by generate and related functions (*i.e.* `tf_utils.py`).


## Next steps
=> In a next step, I propose to tackle both 2) and 3) at the same time. 
@Rocketknight1 Maybe you could take care of 2) while maybe @gante maybe it would make sense if you start with 3) (should be pretty straight-forward after this PR). IMO, this would also greatly help at making more people in the team knowledgeable in how `generate()` works.
Before we go over to 4), I'm happy to do the `past`/`encoder_outputs` refactor (see all those `TODO(Patrick)` comments in `generate`). We still rely on a very ugly hack that puts `encoder_outputs` into the `past` variable. I didn't want to change this here because it require changing the `prepare_inputs_for_generation` method of multiple models and IMO the PR would have become to difficult to review.

Thoughts? @Rocketknight1 @gante @LysandreJik @sgugger


Comments:
_The documentation is not available anymore as the PR was closed or merged._

All of the following slow tests pass:

```bash
RUN_SLOW=1 pytest -sv \
tests/bart/test_modeling_tf_bart.py \
tests/t5/test_modeling_tf_t5.py \
tests/gpt2/test_modeling_tf_gpt2.py \
tests/vision_encoder_decoder/test_modeling_tf_vision_encoder_decoder.py \
tests/encoder_decoder/test_modeling_tf_encoder_decoder.py \
tests/speech_to_text/test_modeling_tf_speech_to_text.py
```

as well as

```bash
RUN_SLOW=1 pytest -sv tests/test_modeling_tf_rag.py
```
(this one is tested on brutasse as it requires a heavy index).

Just finished and it looks very cool! I had a few nits but it seems like a great refactor, which will make it easy to optimize sub-methods for performance later.

Thanks a lot for the review @gante and @Rocketknight1 . Adapted the PR according to your comments and will merge once tests are passing.

Failing tests are flaky and unrelated. Merging

@patrickvonplaten 

About the following

```
We still rely on a very ugly hack that puts encoder_outputs into the past variable.
I didn't want to change this here because it require changing the prepare_inputs_for_generation
method of multiple models and IMO the PR would have become to difficult to review.
```

I haven't looked the details, but it is something different from the PyTorch's models (and `generate`), right?
During the fix for PT/TF inconsistency, this causes some problem for a strict equivalence testing.
Maybe I can work on this once other inconsistencies are fixed (and by that time, `generate` already has some clean refactoring by you guys!).


Hey @ydshieh, 

Great initiative! Yes that's indeed an important refactor we should do soon - would you like to look into it? In short we should do the following in TF:

- Always force `return_dict=True`. The we don't need any of these `self._use_cache()` methods anymore (we can just check whether `past_key_values` is present in the output. Also this way the can clearly differentiate between `encoder_outputs` and `past_key_values` for the encoder-decoder case. Would you like to open a PR for it? Happy to help you along the way!

> Hey @ydshieh,
> 
> Great initiative! Yes that's indeed an important refactor we should do soon - would you like to look into it? In short we should do the following in TF:
> 
>     * Always force `return_dict=True`. The we don't need any of these `self._use_cache()` methods anymore (we can just check whether `past_key_values` is present in the output. Also this way the can clearly differentiate between `encoder_outputs` and `past_key_values` for the encoder-decoder case. Would you like to open a PR for it? Happy to help you along the way!

Sure - after I finish the work of adding testing `past_key_values` for TF encoder models (those can be used for causal lm and can have cross-attention`). Currently, no test for it, and making #15477 undetected for quite some time.