Repository: huggingface/transformers

Issue:
Bug when combining grouped beam search and constrained prefix decoding
## Environment info

- `transformers` version: 4.3.3
- Platform: Linux-5.8.0-38-generic-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.5.1+cu101 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help

@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): T5

The problem arises when using: my own modified scripts

## To reproduce

Steps to reproduce the behavior: run this simple script

```python
from transformers import T5TokenizerFast, T5ForConditionalGeneration

tokenizer = T5TokenizerFast.from_pretrained('t5-small')

inp = 'The <extra_id_0> walks in <extra_id_1> park'
enc_inp = tokenizer(inp, return_tensors='pt')
model = T5ForConditionalGeneration.from_pretrained('t5-small')


def prefix_allowed_tokens_fn(batch_id, input_ids):
    return [2]  # dummy value


out = model.generate(
    **enc_inp,
    num_beams=2,
    num_beam_groups=2,
    diversity_penalty=0.2,
    prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
)
```

This produces the following error:
```
Traceback (most recent call last):
  File "debugging/grouped_beam_search.py", line 14, in <module>
    out = model.generate(
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/mounts/Users/student/martin/.local/lib/python3.8/site-packages/transformers/generation_utils.py", line 1041, in generate
    return self.group_beam_search(
  File "/mounts/Users/student/martin/.local/lib/python3.8/site-packages/transformers/generation_utils.py", line 2161, in group_beam_search
    next_token_scores = logits_processor(
  File "/mounts/Users/student/martin/.local/lib/python3.8/site-packages/transformers/generation_logits_process.py", line 89, in __call__
    scores = processor(input_ids, scores)
  File "/mounts/Users/student/martin/.local/lib/python3.8/site-packages/transformers/generation_logits_process.py", line 458, in __call__
    for batch_id, beam_sent in enumerate(input_ids.view(-1, self._num_beams, input_ids.shape[-1])):
RuntimeError: shape '[-1, 2, 1]' is invalid for input of size 1
```

## Expected behavior

No error.

As far as I can tell, the `PrefixConstrainedLogitsProcessor` still receives the original number of beams even when grouped beam search is used. But it should be the number of subbeams. So replacing `num_beams` with `num_beams // num_beam_groups` in the constructor of `PrefixConstrainedLogitsProcessor` in method `_get_logits_processor` in file `generation_utils.py` should fix it.
What do you think?


Discussion:
Hey @mnschmit,

thanks for your bug report! Yes, you're right -> I think we should indeed replace `num_beams` by `num_beams // num_beam_groups`. Do you want to open a PR to fix it? :-) Otherwise, I can do it as well

Resolving patch:
diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py
--- a/src/transformers/generation_utils.py
+++ b/src/transformers/generation_utils.py
@@ -605,7 +605,7 @@ def _get_logits_processor(
         if min_length is not None and eos_token_id is not None and min_length > -1:
             processors.append(MinLengthLogitsProcessor(min_length, eos_token_id))
         if prefix_allowed_tokens_fn is not None:
-            processors.append(PrefixConstrainedLogitsProcessor(prefix_allowed_tokens_fn, num_beams))
+            processors.append(PrefixConstrainedLogitsProcessor(prefix_allowed_tokens_fn, num_beams // num_beam_groups))
         if forced_bos_token_id is not None:
             processors.append(ForcedBOSTokenLogitsProcessor(forced_bos_token_id))
         if forced_eos_token_id is not None:


Test patch: