Diverse beam search
Author: ayushtiku5
Pull request: yes

# What does this PR do?

<!--
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 -->

Implementation of diverse beam search decoding as described in the paper: https://arxiv.org/pdf/1610.02424.pdf 

diversity function reference taken from: https://github.com/ashwinkalyan/dbs

## Implementation details
Consider a T5 summarization task.
`article="Justin Timberlake and Jessica Biel, welcome to parenthood. 
The celebrity couple announced the arrival of their son, Silas Randall Timberlake, in statements to People. 
"Silas was the middle name of Timberlake's maternal grandfather Bill Bomar, who died in 2012, while Randall is the musician's own middle name, as well as his father's first," People reports. 
The couple announced the pregnancy in January, with an Instagram post. It is the first baby for both."`

Generation using normal beam search can be done as:
`model.generate(
   input_ids=input_ids,
   num_beams=2,
   num_return_sequences=2
)`

This generates:
`['the couple announced the pregnancy in January. it is the first baby for both.', 'the couple announced the pregnancy in January. it is the first baby for both of them ']`

Generation using diverse beam search can be done as:
`model.generate(
   input_ids=input_ids,
   num_beams=2,
   num_return_sequences=2,
   beam_groups=2,
   diversity_penalty=1.5
)`

This generates:
`['the couple announced the pregnancy in January. it is the first baby for both.', 'Justin Timberlake and Jessica Biel have welcomed their son, Silas Randall ']`

This means that 2 beams will be divided into 2 groups of 1 beam each, ensuring diversity between each group. NOTE: If `beam_groups=1`, then it will be same as the normal beam search as all the beams belong to the same group. Higher `diversity_penalty` will ensure more diversity between the groups of beams. When doing generation using diverse beam search, we need to ensure that `num_beams>=beam_groups` and also `num_beams` is divisible by `beam_groups`.
 
## Who can review?

@patrickvonplaten, @TevenLeScao

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 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.

 albert, bert, XLM: @LysandreJik
 GPT2: @LysandreJik, @patrickvonplaten
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @patrickvonplaten, @TevenLeScao
 Blenderbot, Bart, Marian, Pegasus: @patrickvonplaten
 T5: @patrickvonplaten
 Rag: @patrickvonplaten, @lhoestq
 EncoderDecoder: @patrickvonplaten
 Longformer, Reformer: @patrickvonplaten
 TransfoXL, XLNet: @TevenLeScao, @patrickvonplaten
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 FSTM: @stas00
 -->


Comments:
@patrickvonplaten I am implementing diverse beam search. Please do suggest code design for this. 😃 

> @patrickvonplaten I am implementing diverse beam search. Please do suggest code design for this.

Awesome that you work on this!

I think this looks like the right approach! However, I'd also recommend creating a new beam_scorer to be sure to not break backwards compatilibily. We can see at a later stage if we can try to merge some code together with the current beam search code :-) 

Also, can you add a link to the paper in this PR ? this would be great :-) 

@patrickvonplaten please review. I have made the required changes :)

@patrickvonplaten just a gentle reminder to review the PR. Thanks!

> @patrickvonplaten just a gentle reminder to review the PR. Thanks!

Sorry, I'll review the PR this week! Also wondering how this PR relates to this one: #8840

@patrickvonplaten I think #8840 ensures that first token of every predicted sequence is different. This PR ensures diversity between group of beams at every time step of sequence generation. I think this will be more generic. Also we can change extent of diversity using `diversity_penalty` parameter.

@patrickvonplaten Also I was thinking that currently I am subtracting the diversity penalty directly from the `beam_scores`. So, finally when we are doing `beam_scorer.finalize()`, the `final_beam_scores` will also include the effect of `diversity_penalty`. 

I was thinking maybe we should penalise the `beam_scores` with diversity penalty only when we are selecting top `2*group_size` beam candidates:
`next_token_scores, next_tokens = torch.topk(
                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
                )`

But for choosing the final beams in the end the scores shouldn't include the penalty due to diversity. What do you think?

> @patrickvonplaten Also I was thinking that currently I am subtracting the diversity penalty directly from the `beam_scores`. So, finally when we are doing `beam_scorer.finalize()`, the `final_beam_scores` will also include the effect of `diversity_penalty`.
> 
> I was thinking maybe we should penalise the `beam_scores` with diversity penalty only when we are selecting top `2*group_size` beam candidates:
> `next_token_scores, next_tokens = torch.topk( next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True )`
> 
> But for choosing the final beams in the end the scores shouldn't include the penalty due to diversity. What do you think?

Hey @ayushtiku5,

That's a good point! I do think though that we should leave the `beam_scores` as there are in the end as well. My main arguments are:

1) It helps to have more diversity in the output. If we only use the diversity penalty for choosing the next beam_token, but not add it to the `_beam_scores`, the beam_scores will be very high for beams of similar tokens, which I think is what we want to prevent here. I think `beam_scores` should be penalized for every  token in the corresponding `beam_idx` that is also present in another `beam_idx` of the same `beam_group`. It's also more consistent and logical IMO: We should update the `beam_score` with the `probability` that the current beam_id was selected. 

2) It would be very ugly to implement and I'd like to avoid it...

Is that fine for you?

> > @patrickvonplaten Also I was thinking that currently I am subtracting the diversity penalty directly from the `beam_scores`. So, finally when we are doing `beam_scorer.finalize()`, the `final_beam_scores` will also include the effect of `diversity_penalty`.
> > I was thinking maybe we should penalise the `beam_scores` with diversity penalty only when we are selecting top `2*group_size` beam candidates:
> > `next_token_scores, next_tokens = torch.topk( next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True )`
> > But for choosing the final beams in the end the scores shouldn't include the penalty due to diversity. What do you think?
> 
> Hey @ayushtiku5,
> 
> That's a good point! I do think though that we should leave the `beam_scores` as there are in the end as well. My main arguments are:
> 
> 1. It helps to have more diversity in the output. If we only use the diversity penalty for choosing the next beam_token, but not add it to the `_beam_scores`, the beam_scores will be very high for beams of similar tokens, which I think is what we want to prevent here. I think `beam_scores` should be penalized for every  token in the corresponding `beam_idx` that is also present in another `beam_idx` of the same `beam_group`. It's also more consistent and logical IMO: We should update the `beam_score` with the `probability` that the current beam_id was selected.
> 2. It would be very ugly to implement and I'd like to avoid it...
> 
> Is that fine for you?

@patrickvonplaten yeah sure, I am fine with this.

@ayushtiku5 - hope it's ok that I fiddled quite a bit with your PR. The functionality is kept 1:1 the same (I added an integration test in the very beginning to be sure of that), but the design is slightly different with the main goal to keep the method as general as possible. 

IMO, the PR is now good to merge :-) Could you take a final look at whether the new names and design is ok for you? 

Afterward, we can think about a nice code snippet / use case to advertise the big new feature of `transformers` :-) 
Awesome job!

@ayushtiku5 do you think the following code snippet could be a nice use case of diverse beam search?

```python
from transformers import pipeline
summarizer = pipeline("summarization", model="sshleifer/distilbart-xsum-12-6")

ARTICLE = """Part of the Broad Road was closed to traffic on Sunday at about 18:00 GMT.
The three adults and three children have been taken to Altnagelvin Hospital with non
life-threatening injuries. The Fire Service, Northern Ireland Ambulance Service
and police attended the crash. The Broad Road has since been reopened."""

# normal beam search
summarizer(ARTICLE, num_return_sequences=2)
# => [' Five people, including three children, have been taken to hospital following a two-vehicle crash in Londonderry.',
# ' Five people, including three children, have been taken to hospital after a two-vehicle crash in Londonderry.']

# diverse beam search
summarizer(ARTICLE, num_return_sequences=2, num_beam_groups=6, diversity_penalty=10.0)
# => ['Three men are in hospital after a car and a lorry crashed in Londonderry.',
# 'Six pedestrians were injured when a car and two vehicles crashed in County Antrim.']
```

> @ayushtiku5 - hope it's ok that I fiddled quite a bit with your PR. The functionality is kept 1:1 the same (I added an integration test in the very beginning to be sure of that), but the design is slightly different with the main goal to keep the method as general as possible.
> 
> IMO, the PR is now good to merge :-) Could you take a final look at whether the new names and design is ok for you?
> 
> Afterward, we can think about a nice code snippet / use case to advertise the big new feature of `transformers` :-)
> Awesome job!

Hey @patrickvonplaten ,

Just one thing. In the `BeamScorer`'s `finalize()` method, we are directly selecting top `num_beams` beams from the `final_beam_scores`. This assumes that the beam scores in `final_beam_scores` will be sorted in decreasing order for a particular `batch_idx`. However, this will not be the case for our diverse beam search. `final_beam_scores` will be sorted for the beams inside a particular group, but not necessarily for all the beams for a particular `batch_idx`. So, I think we will have to sort the `final_beam_scores` for every `batch_idx`. I did this previously [here](https://github.com/huggingface/transformers/pull/8627/commits/14d5b6ca6e527eac2cdb9e9400d4c00f6d7add01#diff-098eb3834a12a0788445325f6795950fc5d59ec8fc8d34fef115ae5e379e18f2R292)

The rest looks good to me. Thanks for refactoring!

[UPDATE]: added this in [this](https://github.com/huggingface/transformers/pull/8627/commits/c99eb5a8dc57a7b0d33a8ac06d8c6a32a7812ad4) commit

Hey @ayushtiku5,

sorry I forgot to mention on why I deleted those lines. IMO we don't need to add this functionality because it doesn't matter whether the scores are sorted or not. In this line: https://github.com/huggingface/transformers/blob/72d6c9c68ba19b2e991b0d7a32989410399b33f5/src/transformers/generation_beam_search.py#L330 you can see that the `add(...)` method automatically keeps the best scores and throws out the worse scores. Since the loop goes through all scores anyway it does not matter IMO whether they are sorted or not. 

What do you think? IMO, we can revert the last commit.

> Hey @ayushtiku5,
> 
> sorry I forgot to mention on why I deleted those lines. IMO we don't need to add this functionality because it doesn't matter whether the scores are sorted or not. In this line:
> 
> https://github.com/huggingface/transformers/blob/72d6c9c68ba19b2e991b0d7a32989410399b33f5/src/transformers/generation_beam_search.py#L330
> 
> you can see that the `add(...)` method automatically keeps the best scores and throws out the worse scores. Since the loop goes through all scores anyway it does not matter IMO whether they are sorted or not.
> What do you think? IMO, we can revert the last commit.

Yeah sorry! I completely missed it. Reverted the commit.

> > Hey @ayushtiku5,
> > sorry I forgot to mention on why I deleted those lines. IMO we don't need to add this functionality because it doesn't matter whether the scores are sorted or not. In this line:
> > https://github.com/huggingface/transformers/blob/72d6c9c68ba19b2e991b0d7a32989410399b33f5/src/transformers/generation_beam_search.py#L330
> > 
> > you can see that the `add(...)` method automatically keeps the best scores and throws out the worse scores. Since the loop goes through all scores anyway it does not matter IMO whether they are sorted or not.
> > What do you think? IMO, we can revert the last commit.
> 
> Yeah sorry! I completely missed it. Reverted the commit.

No worries :-) The comment wasn't the best either - I updated it. Think it's a bit clearer now.

@ayushtiku5 - super sorry, we messed up the previous branch yesterday. I opened a new PR with the same authorship -> so it should be good to merge :-)