Fix shift_right for padded sequences in MBART
Author: BramVanroy
Pull request: yes

# What does this PR do?

While `shift_tokens_right` indicates that it accounts for padding tokens, it does not do so correctly when padded tokens are at the end of a sequence. In the example below you'll see that the current implementation does not correctly account for the padding tokens when replacing the tokens. This means that the special LID token is not removed/replaced by a padding token when there are padding tokens to start with.

Below an example to show the comparison between the current HF implementation and this PR.

```python
import torch
from transformers import MBartTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right


def shift_tokens_right_mbart(input_ids: torch.Tensor, pad_token_id: int):
    """
    Shift input ids one token to the right, and wrap the last non pad token (the <LID> token) Note that MBart does not
    have a single `decoder_start_token_id` in contrast to other Bart-like models.
    """
    prev_output_tokens = input_ids.clone()

    if pad_token_id is None:
        raise ValueError("self.model.config.pad_token_id has to be defined.")

    # replace possible -100 values in labels by `pad_token_id`
    prev_output_tokens.masked_fill_(prev_output_tokens == -100, pad_token_id)
    index_of_eos = (prev_output_tokens.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
    decoder_start_tokens = prev_output_tokens.gather(1, index_of_eos).squeeze()

    shifted = torch.full_like(input_ids, pad_token_id)
    for b_idx in range(input_ids.size(0)):
        shifted[b_idx, 1:index_of_eos[b_idx]+1] = prev_output_tokens[b_idx, :index_of_eos[b_idx]].clone()

    shifted[:, 0] = decoder_start_tokens

    return shifted


def main():
    text = ["UN Chief Says There Is No Military Solution in Syria"]

    # MBART
    tokenizer = MBartTokenizer.from_pretrained("facebook/mbart-large-cc25", src_lang="en_XX")
    input_ids = tokenizer(text)["input_ids"]
    input_ids[0] += [tokenizer.pad_token_id] * 4
    input_ids = torch.LongTensor(input_ids)

    print("original input", tokenizer.batch_decode(input_ids))
    shifted_hf = shift_tokens_right(input_ids, tokenizer.pad_token_id)
    print("HF implementation", tokenizer.batch_decode(shifted_hf))
    shifted = shift_tokens_right_mbart(input_ids, tokenizer.pad_token_id)
    print("Fixed implementation", tokenizer.batch_decode(shifted))



if __name__ == '__main__':
    main()
```

Output:

```
original input ['UN Chief Says There Is No Military Solution in Syria
en_XX<pad><pad><pad><pad>']
HF implementation ['en_XX UN Chief Says There Is No Military Solution in Syria
en_XX<pad><pad><pad>']
Fixed implementation ['en_XX UN Chief Says There Is No Military Solution in Syria
<pad><pad><pad><pad>']
```

I think this has not come up yet because one would typically train (M)BART on batches without padding (multiple sentences). But because the implementation explicitly mentions padding, I found it odd that it does not seem to work correctly when a padded sequence is used.

If this PR is accepted, the regular BART `shift_tokens_right` may also need a similar fix.

## Who can review?

- bart: @patrickvonplaten, @patil-suraj


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18985). All of your documentation changes will be reflected on that endpoint.

Hey @BramVanroy,

Thanks a lot for the great explanation :heart: - it's very easy to see the problem.

Just a quick question before digging a bit deeper into this - does it really matter?

If the original labels are:

```
'UN Chief Says There Is No Military Solution in Syria
en_XX<pad><pad><pad><pad>'
```
=> this means that we don't compute the loss on the last 4 input tokens (because the target is padding). Given that we also always use a causal mask, the last for input tokens don't influence the previous tokens. Therefore, we should compute the same loss for:
```
HF implementation ['en_XX UN Chief Says There Is No Military Solution in Syria
en_XX<pad><pad><pad>']
```
and
```
Fixed implementation ['en_XX UN Chief Says There Is No Military Solution in Syria
<pad><pad><pad><pad>']
```

no?

Yes, you are definitely right! I do not think this changes the outcome of the model because, as you say, those last padded values are ignored in CE anyway. So it does not matter. It's more a "beauty error", I guess, although it might still be good to fix though (but not urgently).

I'm slightly worried though with for-loops that we're not there before - e.g. I'm not sure ONNX is happy with it. Would it maybe be ok to just add a comment describing the problems and why "it doesn't matter" instead? 

That makes sense. A comment can be useful in case someone ever wants to use the function for other implementations. Then it is important that they are aware that the special token gets duplicated and not swapped for a padding token. I don't have time to make the changes at the moment but I'll keep this open to remind me.

If instead anyone has a vectorized solution, that's also welcome of course.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.