Refactoring the generate() function
Author: patrickvonplaten
Pull request: yes

## Generation refactor

This is a possible design that IMO would make generate much more readable and flexible for future code changes. 
The code shown in this PR is more or less pseudo code, but I'm sure that this design is fully backwards compatible (for the `generate()` function, I don't plan on keep `beam_search_generation` and `_no_beam_search_generation` -> I don't think anybody directly accessed these functions).

It's probably better to look at the code directly, here: https://github.com/huggingface/transformers/blob/3c8a12a3bf44aa7845675852e28c47d0a9cb808e/src/transformers/generation_utils.py then at the diff.

The following major changes are made:

### Split the generate method into four generate functions:

a) `greedy_search` (corresponds to num_beams = 1, do_sample=False)
b) `sample` (corresponds to num_beams = 1, do_sample=True)
c) `beam_search` (corresponds to num_beams > 1, do_sample = False)
d) `beam_sample` (corresponds to num_beams > 1, do_sample = True)

It is split in a way that the functions can be used on their own and don't necessarily have to be accessed by the `generate()` method. This allows for much more flexibility for the user - he can decide what kind of distribution warper he wants to use and what kind of "Beam scorer" he wants to use (more on "distribution warper" and "beam scorer" later.
     Also mostly a model only uses one of these methods, `EncoderDecoder` models usually use `beam_search()`, `...ForCausalLM` models usually use `sample()`. Because `generate` is now split to the corresponding functions relevant for a model, it makes the code much more readable for the users because they will mostly only look into  one of the four functions.
    Splitting `generate()` into four functions removes **a lot** of `if-else` statements

### Create `LogitsProcessor` and `LogitsWarper` objects. 

Instead of adding each of these "logit warpers / processors", such as `bad_token_words` with `if-else` statements, a list of these objects is created in the beginning and then called in the code. This is largely copied from this very nice PR: https://github.com/huggingface/transformers/pull/5420 . The arguments are 1) easier to test, 2) adds more flexbility because users can easily add their own "logit warpers / processors". 
Note that we need both a `pre_processor` and a `dist_warper` list to make `beam_sample` work correctly. This comment explains why in more detail: https://github.com/huggingface/transformers/pull/5420#discussion_r449779867

### Create `beam_scorer` class.

The `generate_beam_search` function has become extremely hard to read because most of the beam search logic is directly written into the function. I would propose to move this code into a `beam_scorer` class that would also replace the `BeamHypotheses` class. The class would essentially expose the following functions: 

```python
beam_scorer.update(next_scores, next_tokens)

next_beam_scores = beam_scorer.get_next_scores()
next_beam_tokens = beam_scorer.get_next_tokens()
next_beam_idx = beam_scorer.get_next_beam_idx()

beam_scorer.is_done()
```

and IMO all the beam search relevant code can be handled in the `beam_scorer.update(...)` function.
Besides better readability, such a class could also be replaced by another "Beam seach scorer" logic, which makes the beam search code much easier to extend IMO.

The con of this PR is obviously that some code would be copy pasted along the four functions (but it should not be too much).

Let me know if something is not clear in the design. Would be super happy to hear your feedback @LysandreJik @thomwolf @sshleifer @yjernite @JetRunner @sgugger 

## 1st Review:

The current state shows how the refactored code would look like. All important tests are passing now (the slow GPT2, Bart & T5 tests), but the PR is not at all finished yet. It would be nice if @sshleifer @yjernite @thomwolf (and @LysandreJik @sgugger) could take a first look at the complete "new" architecture and give some feedback. If you guys are ok with the new design, I will add a bunch of tests and clean up the code more to make sure we have 100% backward compatibily. 

## TODO:
- [x] Make all slow tests pass
- [x] Rename according to discussion
- [x] Add tests for all processors
- [x] Add more and better generation tests
- [x] Do speed comparison
- [x] Add docstring
- [x] Final thoughts about design

## Final review:

I'm happy with the changes I've done now. Better tests (more aggressive and should also be faster since only 3 tokens are generated instead of 10 now) have been added and docstrings have hopefully become more understandable. 
A complete explanation of new "generation" philosophy is described in the forum: https://discuss.huggingface.co/t/big-generate-refactor/1857

Comments:
I think this is a massive improvement and really hope it gets merged. Everything is much better encapsulated and it's super easy to find what you want.

I have some naming ideas, but the code seems to not be finished so I will wait.

I am very unfamiliar with the generate code because the code has scared me way, this one is way more inviting to read and I can actually understand it so I think that's a very nice achievement :-)

I don't know if what's hidden in `# add necessary encoder decoder code` and `# add all post processing functions` is different for the four versions of decode, but might be nice to refactor it in some helper methods if possible.

Seconding Sam and Sylvain. Really excited for the improved legibility but we should also make sure that code isn't copy-pasted four times or we'll start having inconsistencies soon :) 

We should probably allow `dist_warper` in the beam search too: the current ones wouldn't make a difference but some use cases will (e.g. noisy channel which uses a backward distribution to reweigh the next token scores)

Related: https://github.com/huggingface/transformers/issues/7626

Can you please add this functionality too https://github.com/huggingface/transformers/issues/5164?

This is really cool! I love the additional clarity and think you're building a really strong base for the generation code. Here's what I think we should still chat about before finalizing the PR:
1. We still have inconsistencies in how we handle the end of generation. I really think that `max_length` should be handled the same way as `min_length` with a sampler in `pre_processor ` which forces a Dirac on `eos_token_id` when `cur_len==max_length`
2. Similarly, we should get rid of `adjust_logits_during_generation` and have e.g. a Bart-specific sampler in `pre_processor `
3. We're missing out on supporting some really interesting research by not giving a good option to return and backprop through the generation scores. Also, I think we should extend self-documenting outputs to the `generate` function (cc @sgugger ). My proposal would be to add a `return_dict` argument, with an `output_generation_scores` (and possibly `with_grad`) option. (And we just return the generated ids if `return_dict=False` to stay backwards compatible)
4. Generation with `decoder_prefix_ids` for encoder-decoder models :D ! That can be a future PR though, and can also be handled with `pre_processor` to force the output (a bit wasteful but will do in a pinch)

I also agree with @sgugger that we should come up with a better name than `Sampler`, will think about it

@yjernite has great feature requests, but this PR is already huge and I don't see why they need to be handled here.

> @yjernite has great feature requests, but this PR is already huge and I don't see why they need to be handled here.

Good point, we can definitely look at 3. and 4. later, and I know from experience that 2. is probably the wrong kind of rabbit hole to get into right now.

I am a little concerned about 1. though. The samplers will definitely need to consider more than just `input_ids` and `scores` in the future, and we should make sure that we don't need to rebuild them from the ground up when that happens. 

@patrickvonplaten what are your thoughts on changing e.g. [generation_utils.py#L451](https://github.com/huggingface/transformers/blob/5df79e2c41bf4b47ab4a36f903be163252714fe3/src/transformers/generation_utils.py#L451) when we need to also pass `cur_len` to enforce the max length or if we want the samples to look back at the `encoder_input_ids` ?

Thanks for the feedback! @yjernite - regarding your points: 

1) I think I see your point that `max_length` should be treated the same way as `min_length`. If we follow this appoarch, we would replace the `while cur_len < max_length:` with `while True:` and then break if `max_length` is hit. I'm a bit worried the people that will use one of the four functions directly, such as: 

```python 

pre_processor = # create your list of pre_processors here <= THIS MUST INCLUDE max_length
dist_warper = # create your dist warper here

outputs = model.sample(input_ids, pre_processor, dist_warper, max_length, pad_token_id, eos_token_id, **model_kwargs)
```

will forget to put a `max_length` "processor" in `pre_processor` and then the `while True:` loop would run forever. 

For me, the difference between `max_length` and and *e.g.* `min_length` is that `max_length` is a mandatory parameter for generation, which is why I left it as an input to the "sub" generation functions.
For you, what would be the big advantage of moving `max_length` to a preprocessor item (besides consistency?).

I'm 100% fine with extending the pre-processors or warpers (trying to not use the word samplers anymore :D) to accept more input arguments, but I think we could also do this in a future PR as it would not break backwards compatibility.

Happy to discuss what opinions the others have on this!

2) Agree - will have to see how to make that non-breaking, but should be possible!

3) Agree to add a `ModelOutputs` class to `generate()` that would include `attentions` and `hidden_states`. Regarding being able to backprop through `generate()` - I don't really think that this is super important. It would probably also require a lot of tweaking the way generate() is done currently so I'd prefer if ppl would just use the own fork/branch for these kind of things

4) This should already be possible I think. `decoder_input_ids` can be passed to generate.

## Speed comparison:

sample search / greedy search yields equivalent results (TESTED on GPT2)

beam search yields ~5 % speed up thanks to the use of tensors instead of lists in `beam_scorer` (TESTED on BART and T5)

All slow tests that pass on master, also pass in this PR now. Ran all those tests: 

```
run_generation_integration_tests () {
        RUN_SLOW=1 pytest tests/test_modeling_pegasus.py
        RUN_SLOW=1 pytest tests/test_modeling_bart.py
        RUN_SLOW=1 pytest tests/test_modeling_t5.py
        RUN_SLOW=1 pytest tests/test_modeling_reformer.py
        RUN_SLOW=1 pytest tests/test_modeling_marian.py
        RUN_SLOW=1 pytest tests/test_modeling_mbart.py
        RUN_SLOW=1 pytest tests/test_modeling_prophetnet.py
        RUN_SLOW=1 pytest tests/test_modeling_xlm_prophetnet.py
        RUN_SLOW=1 pytest tests/test_modeling_encoder_decoder.py
        RUN_PIPELINE_TESTS=1 RUN_SLOW=1 pytest tests/test_pipelines_conversational.py
        RUN_PIPELINE_TESTS=1 RUN_SLOW=1 pytest tests/test_pipelines_text2text_generation.py
        RUN_PIPELINE_TESTS=1 RUN_SLOW=1 pytest tests/test_pipelines_text_generation.py
        RUN_PIPELINE_TESTS=1 RUN_SLOW=1 pytest tests/test_pipelines_summarization.py
        RUN_PIPELINE_TESTS=1 RUN_SLOW=1 pytest tests/test_pipelines_translation.py
        RUN_PIPELINE_TESTS=1 RUN_SLOW=1 pytest tests/test_pipelines_dialog.py
        RUN_SLOW=1 pytest tests/test_modeling_gpt2.py
        RUN_SLOW=1 pytest tests/test_modeling_xlnet.py
        RUN_SLOW=1 pytest tests/test_modeling_transfo_xl.py
        RUN_SLOW=1 pytest tests/test_modeling_rag.py
        RUN_SLOW=1 pytest tests/test_modeling_fsmt.py
        RUN_SLOW=1 pytest tests/test_modeling_blenderbot.py
}
```

Think I ran all of the important ones (cc @sshleifer )

> I don't know if this PR was ready for review or not but I still went ahead ;-). I still like this a lot, I added suggestions for the docs and I think it would be great if all your awesome work was documented. `GenerationMixin` is already documented in `main_classes/model` so all its public methods should have nice doc.
> 
> Then I think all you added warrants an `internal/generation` where we could document all the tools you added.
> 
> Last thing: I'm not super fan of `DistProcessor` as a name, mainly because I don't know that dist means distribution (I think distributed but I'm obsessed with Trainer ;-) ). `DistributionProcessor` might be a bit long. The class seems to be mainly doing some preprocessing on the logits though, so why not `LogitsProcessor`?

Changed to `LogitsProcessor` - like that name! Yeah the docs weren't ready yet, but thanks for the feedback :-) 

Hmmm OK so after thinking about it for a bit I realize that having a `max_length` logits processor would change the current beam search behavior in some cases.

Currently, the beam gets re-ordered at `max_length` without penalizing sequences that aren't "well-formed" (haven't reached EOS within the allotted time). If we switched to a `LogitsProcessor` however, the beam search would up-weight sequences that were more likely to lead to an EOS at time step (`max_length`-1)

I think this new behavior is better and definitely would like the option to implement it (@srush would love your input on that especially), but I'm not sure how breaking it is (conversely, it might also account for some of the different scores we've had from other libs). We can also just add a `BeamSearchMaxLengthLogitsProcessor` now or later.

What do you think @thomwolf @patrickvonplaten @sshleifer ?

> We can also just add a BeamSearchMaxLengthLogitsProcessor now or later.

This sounds useful! Bart forces all the generations of length `max_length -1` to end in EOS, which helps performance. Might help other models too.  Would lean towards adding it later. 

> This all looks great to me. I have final nits on the docs, some general rules not to forget:
> 
> * no abbreviation in the documentation as we have all the space we want to explain things to the user
> * no lines > 119 pretty please, the script takes care of everything except the examples and some of the examples have veeeeeery long lines.



> This all looks great to me. I have final nits on the docs, some general rules not to forget:
> 
> * no abbreviation in the documentation as we have all the space we want to explain things to the user
> * no lines > 119 pretty please, the script takes care of everything except the examples and some of the examples have veeeeeery long lines.

Noo, I broke the 119 rule again - I though the script would now save me from everything :D Will correct this! Should I also add the `>>>` for the example code or is this just for the model's forward function? 

> Should I also add the `>>>` for the example code or is this just for the model's forward function?

It's only used as a marker for doctests, so you should do this for examples that are not slow. We still need to resuscitate the doctests with @LysandreJik for it to have any use, though.


The doctest run as slow tests so you can still add them even if they're slow!

> The doctest run as slow tests so you can still add them even if they're slow!

Okey added `>>>` to all examples and made them pretty

Hi, thanks for creating this PR and the code looks great! Is it possible to decode and meanwhile return the token probabilities? It would be very helpful in the following scenarios -

I. Calculating perplexities of generated texts
II. Reinforcement learning for text generation.

Is it possible to backpropagate the conditional text generation model parameters on the loss defined via this .generate() function?

Hi @patrickvonplaten
Can you share how to calculate the output probabilities of each token given by generate() ?