Adding the state-of-the-art contrastive search decoding methods for the codebase of generation_utils.py
Author: gmftbyGMFTBY
Pull request: yes

# Adding the state-of-the-art contrastive search decoding method for the `generation_utils` codebase

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #19182

In this PR, I add the source codes of our proposed state-of-the-art decoding methods for the off-the-shelf neural text generation models. The main changes are in the following files: (1) `src/transformers/generation_utils.py`; (2) `examples/pytorch/text-generation/run_generation_contrastive_search.py`. To run the test script, please follow these commands:

```bash
cd examples/pytorch/text-generation;
CUDA_VISIBLE_DEVICES=0 python run_generation_contrastive_search.py --model_type=gpt2 --model_name_or_path=gpt2-large
```

## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests), Pull Request section?
- [x] The PR has been well discussed in [19182](https://github.com/huggingface/transformers/issues/19182)
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests? Yes, I have written the test scripts for the [contrastive search](https://github.com/gmftbyGMFTBY/transformers/blob/csearch-pr-v2/examples/pytorch/text-generation/run_generation_contrastive_search.py)


## Who can review?

According to the suggestions of @gante, @patrickvonplaten and @sgugger can review this PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->


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

@sgugger @patrickvonplaten context: this is the implementation by the authors of [this NeurIPS paper](https://arxiv.org/abs/2202.06417), as first proposed in #19182 -- a new generation strategy with very interesting results!

Hi, @sgugger, thank you so much for your suggestions. I will fix these problems quickly!

Hello @sgugger, is there any document or introduction for auto-APIs?

The [documentation](https://huggingface.co/docs/transformers/model_doc/auto) would be the place to start. You can also look at all other examples!

Hello, @sgugger, I have fixed the problems based on your valuable suggestions! Besides, I have updated the test scripts to the auto-APIs of inference.

The command line to run this test script can be found in its docstring.

Hello @sgugger, I have fixed the problems based on your suggestions.

Hey @gmftbyGMFTBY,

Super nice PR! Contrastive search decoding would be a great addition to Transformers. I feel a bit uneasy of the logic in src/transformers/generation_contrastive_search.py and would prefer to not add a new file here and instead try to adapt the code more to our existing design (try to make use of logits processor, call the model **only** in the loop and not before).

We usually haven't added any new files for generation and molding the design more into what already exists would help a lot with maintainability. 
As an example for constrained decoding, we sadly don't manage to maintain the function anymore (see https://github.com/huggingface/transformers/pull/17920) because the code was too different & too hard for us to maintain (IMO).

If possible I'd advocate quite strongly for trying to make the PR a bit more compatible with our current design (it shouldn't require much work IMO).

@gante @sgugger what do you think? 

I don't know the generate code well enough to evaluate how easy/hard it would be for this PR to fit in the existing design. I solely gave my approval based on the fact it was mimicking something that existed with constraint decoding. Sorry, I wasn't aware it wasn't maintained anymore.

Will know for the next time :-)



@patrickvonplaten replying here since I believe the multiple important points you raised are related. This will be a bit long, but bear with me - I think we can improve the quality of the PR before it gets merged 🙌 

1. `ContrastiveDecodingOneStepFast` (which should be snake cased, perhaps into `contrastive_decoding_step`) being a stand-alone function - it could be part of the loop, yes, but this is a long operation that in essence replaces picking the argmax (in `greedy_search`) or sampling (in `sample`). I agree it should not be public, but I believe we would stand to gain from a readability standpoint if we separate it somehow -- perhaps a private method in the MixIn or an inner function to `contrastive_search`? These two options would also mean that we don't need to pass `model` as an argument, which is indeed awkward.
2. `ranking_fast` being a `LogitsProcessor` - I would disagree here, despite being an operation that processes logits. In addition to breaking the fairly stable `call()` API with two new tensor inputs (`context_hidden` and `next_hidden`), it can only be used with `contrastive_search`. As such, I think we would only be adding a confusing `LogitsProcessor` to the public API.
3. Despite the above, I agree that we should pass and compute the logits processors before the decoding function, like the other `generate` methods -- this was an oversight on my end at review time! As for `top_k`, I have no strong feelings - on one hand, `contrastive search` is not meant to be used with other `LogitWarpers` (like `TemperatureLogitsWarper`), on the other hand there is probably no harm in using `LogitWarpers`.

(now the most challenging issue in design, IMO)

4. A naive implementation of `contrastive_search` needs two sets of forward passes per token -- once to get the `top_k` candidate tokens, another to get the future hidden states for each candidate (to be used as a filter). We can avoid this 2x cost if we pipe future `past_key_values` corresponding to the selected candidate into the actual `past_key_values`. It does mean, however, that the first token requires an additional forward pass. The current code does it by placing `prepare_inputs_for_generation()` (and a few other ops) in unconventional places, before the loop and at the end of the loop. Perhaps it would be better if we kept the original structure, but added an `if` at the start loop -- if we are in the first iteration, do this set of operations.

WDYT about these 4 points?

Thanks for the summary @gante !

Regarding the 4 points above:

1.) I think the `ContrastiveDecodingOneStepFast` is the core of the contrastive generation algorithm and I don't see any advantage having it it's own "sub-function" (it'll never be called from another function than `contrastive_search`).  Also, at the moment, the method returns 6 tensors and accepts 8 or so arguments so it's much more than just replacing the greedy argmax operation IMO. To begin with, I think it'd be very helpful to just copy-paste all the code into the `contrastive_search` and then later down the road I think we could see if parts of it could be moved out (don't think it's necessary though).

2.) I think we would have the following gains to having it implemented as a logit processor:
- We understand the method better & easier to maintain. If we mold `ranking` fast into a logit processor, it's much easier to understand it
- Very easy to test this method
- It is to me a method that processes logits -> so logically it should be a logits processor to me
- Now regarding the API, IMO `logits_processor` inputs are not just restricted to `input_ids` and `scores`. IMO, the API is rather (`input_ids`, `scores`, <other-args-that-are-needed>) => `scores` . Also, we've already "broken" the "input_ids" & "scores" - only API here: https://github.com/huggingface/transformers/blob/d2e5b19b821f0cf43c7cf4f01be5faa1cb20aa64/src/transformers/generation_utils.py#L2989 
=> I understand your point here and ok for me to not make it a logits processor, but IMO it would be better/cleaner
3.) Yes I think it's very important to align the structure of the generate method as much as possible to the other ones
4.) Totally fine for me to have an extra forward pass in the beginning, it's just important that we make the structure the same as greedy or sample (I don't have a problem with if statements here at all). For maintenance and also to understand the method, it's very important IMO that one could compare greedy search and constrastive search line-by-line and then see quickly how the two methods differ

Thank you for your valuable responses! I will revise this PR according to your reviews.

@gmftbyGMFTBY let's go with the core of @patrickvonplaten's suggestions. Patrick and I also talked on Slack to align a few details regarding `ranking_fast` :D

Here's a summary of the main changes we are requesting, to ensure the code you're adding remains easy to maintain:
1. Move the contents of `ContrastiveDecodingOneStepFast` to where the function is called, as opposed to being a function call;
2. Keep `ranking_fast` as it is now (i.e. NOT a logits processor);
3. Let's apply the logits processors at the start of each iteration, and move `top_k` to the logits warpers (like `sample` [does it](https://github.com/huggingface/transformers/blob/d2e5b19b821f0cf43c7cf4f01be5faa1cb20aa64/src/transformers/generation_utils.py#L2059)). This ensures that there are minimal differences between generation strategies;
4. Let's rearrange the order of operations such that all model forward passes happen inside the generation loop (with an `if` for the operations that are only supposed to happen on the first iterations).

Here is a diagram with the expected code structure, to ensure we are all on the same page: https://miro.com/app/board/uXjVPMiVFFg=/

Finally, thank you for your cooperation with us 🙌 This back and forth may be a bit frustrating, but it will ensure your contribution will be long-lived!


@gante Nice! but I just saw your message after I finished writing the `logits_processor` corresponding to `ranking_fast` function.
The followings are the implementation of the `ranking_fast`'s `logits_processor`:
![image](https://user-images.githubusercontent.com/27548710/196190366-5cb9c697-051d-49de-903d-b4279b3d8571.png)

It is initialized in `_get_logits_processor` function:
![image](https://user-images.githubusercontent.com/27548710/196190683-70561c32-c38a-4991-bc01-e058b3378a34.png)

and will be called:
![image](https://user-images.githubusercontent.com/27548710/196190797-fb7568d1-99d2-400d-a45b-205522b8309f.png)

Can I keep this implementation or just still employ the `ranking_fast` function?



@gante Oh, I got your reason about not using the `logits_processor`. I will follow the instruction in the `sample` function!

The revisions have been updated!

Hello, @gante, I have updated the PR based on your suggestions.

Oh, I am still working on the integration test.

@gmftbyGMFTBY you probably have to add the `@slow` decorator to the test, and run it locally with `RUN_SLOW=1 py.test (...)` to confirm that it is working. 

Our CI doesn't run tests with `@slow` on push (and fails if the test doesn't have the decorator and is actually slow), but we run them every 24h and track them internally :)

Ok, I got it!


Okay, I am working on it! Thanks a lot for your reviews!

BTW @gmftbyGMFTBY,

Just read a through your extremely nice issue! It seems like you experimented with OPT as well, so maybe let's add a test for OPT as well then ? :-) OPT's `past_key_values` are slightly different compared to GPT2's `past_key_values` so maybe instead of adding a test for GPT-J and GPT-2, it would make more sense to add a test for OPT in addition to GPT2?

Also, if the paper is only concerned with open-ended generation (so less with encoder-decoder architectures), I'm also totally fine with **not** testing for T5 and BART (it's a nice to have, but if it takes too much time and it's not too important - happy to skip it!).

Regarding the fast dummy test, could you maybe make use of those dummy models: 
- https://huggingface.co/hf-internal-testing/tiny-random-gpt2
- https://huggingface.co/hf-internal-testing/tiny-random-gptj
- https://huggingface.co/hf-internal-testing/tiny-random-t5
- https://huggingface.co/hf-internal-testing/tiny-random-bart

The tests colud look very similar to:
https://github.com/huggingface/transformers/blob/71ca79448cd334970fa2893f4faaa094ca13ca6f/tests/generation/test_generation_utils.py#L2053

just much shorter, *i.e.* they only need to test for shape equality. 

Yeah, we have already tested the OPT models, and it works fine. I will supply more tests to the pre-trained models that you mentioned.

@patrickvonplaten more tests about these models are added: 
* gpt2-large
* gpt-j (EleutherAI/gpt-j-6B)
* opt (facebook/opt-6.7b)
* BART (facebook/bart-large-cnn)
* T5 (flax-community/t5-base-cnn-dm)

These tests are passed successfully. Can you do the final check about this PR?

Thank you for being part of this process @gmftbyGMFTBY 🙌 All queries have been addressed and the PR looks in a good state, merging! 

@gante @patrickvonplaten @sgugger  Wow, Thank you very much for your help and support. Love huggingface team! 

@gante @patrickvonplaten @sgugger -- Many thanks for your kind help throughout the process! It means a great deal to me and @gmftbyGMFTBY. Huggingface is the best!

Great work @gmftbyGMFTBY and @yxuansu, thanks for bearing with us through the PR :-)