Repository: huggingface/transformers (Python)
Date: 2022-09-01 11:11:16+00:00
Task type: Bug Fix

Issue:
BUG for beam_indices from model.generate()
### System Info

- `transformers` version: 4.22.0.dev0
- Platform: Linux-5.8.0-51-generic-x86_64-with-glibc2.10
- Python version: 3.8.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@patil-suraj, @patrickvonplaten, @LysandreJik

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python
from transformers import BartTokenizer,BartForConditionalGeneration
model_path = "/data/pretrained_model/bart_base"
toker = BartTokenizer.from_pretrained(model_path)
model = BartForConditionalGeneration.from_pretrained(model_path)

input_tokens = ["what do you think it ? huggingface is a great library. And I enjoy it very much",
                "transformers is so good"]
batch_size = 2
num_beams = 10
max_length = 10
num_return_sequences = 5
input_ids = toker(input_tokens,return_tensors='pt',padding=True).input_ids
output=model.generate(input_ids,max_length=max_length,\
						num_beams=num_beams,num_return_sequences=num_return_sequences,\
						return_dict_in_generate=True,output_scores=True)
print(output.beam_indices)
```
![image](https://user-images.githubusercontent.com/38466901/187733097-195fda80-3b1f-4b59-898f-e2eacf10729d.png)
![image](https://user-images.githubusercontent.com/38466901/187734309-9fde1b06-3172-4730-97d6-42e953cbffc9.png)



### Expected behavior

This is super weird that `beam_indices` of second batch has indices in the first 10 beams. If calculate the average logits across the sentence according to this `beam_indices`, we won't get the `output.sequences_scores` So I think the number in the red box of the first picture should be added 10 (num_beams), if we add 10, we can get the correct token to be generated in `output.sequences[5]` as shown in the second picture


Issue discussion:
Also, could you please check this ? https://discuss.huggingface.co/t/larger-sum-logits-larger-sum-probability/22358
Also cc @gante for `generate` :)

Solution patch:
diff --git a/src/transformers/generation_beam_search.py b/src/transformers/generation_beam_search.py
--- a/src/transformers/generation_beam_search.py
+++ b/src/transformers/generation_beam_search.py
@@ -259,7 +259,7 @@ def process(
                         continue
                     if beam_indices is not None:
                         beam_index = beam_indices[batch_beam_idx]
-                        beam_index = beam_index + (next_index,)
+                        beam_index = beam_index + (batch_beam_idx,)
                     else:
                         beam_index = None
 


Test patch:
diff --git a/tests/generation/test_generation_beam_search.py b/tests/generation/test_generation_beam_search.py
--- a/tests/generation/test_generation_beam_search.py
+++ b/tests/generation/test_generation_beam_search.py
@@ -172,7 +172,7 @@ def cut_expected_tensor(tensor):
                 input_ids[correct_idx].tolist(), beam_scorer._beam_hyps[batch_idx].beams[0][1].tolist()
             )
             self.parent.assertListEqual(
-                expected_beam_indices + [next_indices[batch_idx, 1].item()],
+                expected_beam_indices + [correct_idx],
                 torch.tensor(beam_scorer._beam_hyps[batch_idx].beams[0][2]).tolist(),
             )