BART example does not produce expected masks
Author: BramVanroy
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.21.2
- Platform: Windows-10-10.0.19041-SP0
- Python version: 3.8.8
- Huggingface_hub version: 0.9.1
- PyTorch version (GPU?): 1.12.1+cu113 (True)


### Who can help?

@sgugger @patil-suraj
The example was added by @duongna21 and @sanchit-gandhi 
### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Here is a minimal example that I extracted from [the example](https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_bart_dlm_flax.py) (turned into functions, added example data).

```python
import math
from itertools import chain
from typing import Dict, List

import numpy as np
from transformers import AutoTokenizer, BatchEncoding, PreTrainedTokenizerBase


def shift_tokens_right(input_ids: np.array, pad_token_id: int, decoder_start_token_id: int) -> np.ndarray:
    """
    Shift input ids one token to the right.
    """
    shifted_input_ids = np.zeros_like(input_ids)
    shifted_input_ids[:, 1:] = input_ids[:, :-1]
    shifted_input_ids[:, 0] = decoder_start_token_id

    shifted_input_ids = np.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
    return shifted_input_ids


def collate(examples: List[Dict[str, List[int]]], tokenizer: PreTrainedTokenizerBase, decoder_start_token_id=2, permute_sentence_ratio=1.0,
            mask_ratio=0.3, poisson_lambda=3.5) -> BatchEncoding:
    # convert list to dict and tensorize input
    batch = BatchEncoding(
        {k: np.array([examples[i][k] for i in range(len(examples))], dtype=int) for k, v in examples[0].items()}
    )

    batch["labels"] = batch["input_ids"].copy()
    batch["decoder_input_ids"] = shift_tokens_right(
        batch["labels"], tokenizer.pad_token_id, decoder_start_token_id
    )
    # permuting sentences
    do_permute = False
    if permute_sentence_ratio > 0.0:
        batch["input_ids"] = permute_sentences(batch["input_ids"], tokenizer, permute_sentence_ratio=permute_sentence_ratio)
        do_permute = True

    # masking span of tokens (text infilling in the paper)
    if mask_ratio:
        batch["input_ids"], batch["labels"] = span_mask_tokens(
            batch["input_ids"], batch["labels"], tokenizer, do_permute=do_permute, poisson_lambda=poisson_lambda, mask_ratio=mask_ratio
        )

    # ignore pad tokens
    batch["attention_mask"] = (batch["input_ids"] != tokenizer.pad_token_id).astype(int)
    batch["decoder_attention_mask"] = (batch["decoder_input_ids"] != tokenizer.pad_token_id).astype(int)
    return batch

def permute_sentences(input_ids, tokenizer, permute_sentence_ratio=1.0):
    """
    Shuffle sentences in each document.
    """
    results = input_ids.copy()

    # find end locations of sentences
    end_sentence_mask = input_ids == tokenizer.pad_token_id
    sentence_ends = np.argwhere(end_sentence_mask)
    sentence_ends[:, 1] += 1
    example_has_multiple_sentences, num_sentences = np.unique(sentence_ends[:, 0], return_counts=True)
    num_sentences_map = {sent_idx: count for sent_idx, count in zip(example_has_multiple_sentences, num_sentences)}

    num_to_permute = np.ceil(num_sentences * permute_sentence_ratio).astype(int)
    num_to_permute_map = {
        sent_idx: count for sent_idx, count in zip(example_has_multiple_sentences, num_to_permute)
    }

    sentence_ends = np.split(sentence_ends[:, 1], np.unique(sentence_ends[:, 0], return_index=True)[1][1:])
    sentence_ends_map = {sent_idx: count for sent_idx, count in zip(example_has_multiple_sentences, sentence_ends)}

    for i in range(input_ids.shape[0]):
        if i not in example_has_multiple_sentences:
            continue
        substitutions = np.random.permutation(num_sentences_map[i])[: num_to_permute_map[i]]
        ordering = np.arange(0, num_sentences_map[i])
        ordering[substitutions] = substitutions[np.random.permutation(num_to_permute_map[i])]

        # write shuffled sentences into results
        index = 0
        for j in ordering:
            sentence = input_ids[i, (sentence_ends_map[i][j - 1] if j > 0 else 0) : sentence_ends_map[i][j]]
            results[i, index : index + sentence.shape[0]] = sentence
            index += sentence.shape[0]
    return results


def span_mask_tokens(input_ids, labels, tokenizer, do_permute=True, poisson_lambda=3.5, mask_ratio=0.3):
    """
    Sampling text spans with span lengths drawn from a Poisson distribution and masking them.
    """
    special_tokens_mask_labels = [
        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
    ]
    special_tokens_mask_inputs = [
        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in input_ids.tolist()
    ]
    special_tokens_mask_labels = np.array(special_tokens_mask_labels, dtype=bool)
    special_tokens_mask_inputs = np.array(special_tokens_mask_inputs, dtype=bool)

    # determine how many tokens we need to mask in total
    is_token_mask = ~(input_ids == tokenizer.pad_token_id) & ~special_tokens_mask_inputs
    num_tokens_to_mask = int(math.ceil(is_token_mask.astype(float).sum() * mask_ratio))
    if num_tokens_to_mask == 0:
        return input_ids, labels

    # generate a sufficient number of span lengths
    span_lengths = np.random.poisson(lam=poisson_lambda, size=(num_tokens_to_mask,))
    while np.cumsum(span_lengths, 0)[-1] < num_tokens_to_mask:
        span_lengths = np.concatenate(
            [span_lengths, np.random.poisson(lam=poisson_lambda, size=(num_tokens_to_mask,))]
        )

    # remove all spans of length 0
    # note that BART inserts additional mask tokens where length == 0,
    # which we do not implement for now as it adds additional complexity
    span_lengths = span_lengths[span_lengths > 0]

    # trim to about num_tokens_to_mask tokens
    cutoff_idx = np.argmin(np.abs(np.cumsum(span_lengths, 0) - num_tokens_to_mask)) + 1
    span_lengths = span_lengths[:cutoff_idx]

    # randomly choose starting positions for masking
    token_indices = np.argwhere(is_token_mask == 1)
    span_starts = np.random.permutation(token_indices.shape[0])[: span_lengths.shape[0]]
    # prepare mask
    masked_indices = np.array(token_indices[span_starts])
    mask = np.full_like(input_ids, fill_value=False)

    # mask starting positions
    for mi in masked_indices:
        mask[tuple(mi)] = True
    span_lengths -= 1

    # fill up spans
    max_index = input_ids.shape[1] - 1
    remaining = (span_lengths > 0) & (masked_indices[:, 1] < max_index)
    while np.any(remaining):
        masked_indices[remaining, 1] += 1
        for mi in masked_indices:
            mask[tuple(mi)] = True
        span_lengths -= 1
        remaining = (span_lengths > 0) & (masked_indices[:, 1] < max_index)

    # place the mask tokens
    mask[np.where(special_tokens_mask_inputs)] = False
    input_ids[np.where(mask)] = tokenizer.mask_token_id
    if not do_permute:
        labels[np.where(mask == 0)] = -100
    else:
        labels[np.where(special_tokens_mask_labels)] = -100

    # remove mask tokens that are not starts of spans
    to_remove = (mask == 1) & np.roll((mask == 1), 1, 1)
    new_input_ids = np.full_like(input_ids, fill_value=tokenizer.pad_token_id)
    for i, example in enumerate(input_ids):
        new_example = example[~to_remove[i]]
        new_input_ids[i, : new_example.shape[0]] = new_example

    return new_input_ids, labels


def group_texts(examples, max_seq_length):
    # Concatenate all texts.
    concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
    # customize this part to your needs.
    if total_length >= max_seq_length:
        total_length = (total_length // max_seq_length) * max_seq_length
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
        for k, t in concatenated_examples.items()
    }
    return result


def main():
    pass

if __name__ == '__main__':
    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
    text = [f"I have never seen a man eating such a large cookie in one sitting.{tokenizer.pad_token}Wow!",
            "In the evening he often breaks into the bakery to much on cookies and milk"]

    # For the example we first need to work with a batch (because group_texts is batched)
    # and then convert it back to a list of samples (to test the the collator)
    encoded = tokenizer(text, padding=True, return_attention_mask=False)
    max_len = max(len(e) for e in encoded["input_ids"])
    encoded = group_texts(encoded, max_len//2)
    print(f"input_ids after group (len: {max_len//2})", tokenizer.batch_decode(encoded["input_ids"]))

    # Convert batch back to list of samples and collate
    pre_collate = [{k: seq} for k, v in encoded.items() for seq in v]
    expected_mask_ratio = 0.3
    batch = collate(pre_collate, tokenizer, mask_ratio=expected_mask_ratio)

    print("input_ids after collate", tokenizer.batch_decode(batch["input_ids"]))

    # Calculate mask coverage
    unique, counts = np.unique(batch["input_ids"], return_counts=True)
    maskcounts = dict(zip(unique, counts))[tokenizer.mask_token_id]
    # Count the number of tokens but exclude special tokens
    ntokens = sum([len([t for t in e if t not in tokenizer.all_special_ids]) for e in encoded["input_ids"]])
    print(f"produced masks coverage (expected {expected_mask_ratio})",
          maskcounts / ntokens, f"({maskcounts}/{ntokens})")

    ntokens_after_collate = sum([len([t for t in e if t not in tokenizer.all_special_ids]) for e in batch["input_ids"]])
    print(f"no. tokens before {ntokens}, no. tokens after collate {ntokens_after_collate}")
```

The output will be something like this:

```
input_ids after group (len: 10) ['
I have never seen a man eating such a', ' large cookie in one sitting.<pad>Wow!
', '
In the evening he often breaks into the bakery', ' to much on cookies and milk
<pad><pad><pad>']
input_ids after collate ['
I have never seen a<mask> eating<mask><pad>', ' large cookie<mask> one sitting.<pad>Wow!
', '
In the evening he<mask> the bakery<pad><pad>', '<pad><pad> to much on cookies and milk
<pad>']
produced masks coverage (expected 0.3) 0.125 (4/32)
no. tokens before 32, no. tokens after collate 25
```

### Expected behavior

Am I wrong in expecting a masking coverage that is closer to 0.3 (here it is only 0.125)?

In addition, I find it very odd that the last sequence suddenly has _prepended_ padding tokens?! I found that this does not always happen (there is some random sampling happening so results vary). But I don't think this should every happen?

Comments:
Seems that it was a combination of: too short sequence length and not the best way of calculating the actual coverage (not accounting for span masks). Revised example that yields the expected results:

```python
import math
from itertools import chain
from typing import Dict, List

import numpy as np
from transformers import AutoTokenizer, BatchEncoding, PreTrainedTokenizerBase


def shift_tokens_right(input_ids: np.array, pad_token_id: int, decoder_start_token_id: int) -> np.ndarray:
    """
    Shift input ids one token to the right.
    """
    shifted_input_ids = np.zeros_like(input_ids)
    shifted_input_ids[:, 1:] = input_ids[:, :-1]
    shifted_input_ids[:, 0] = decoder_start_token_id

    shifted_input_ids = np.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
    return shifted_input_ids


def collate(examples: List[Dict[str, List[int]]], tokenizer: PreTrainedTokenizerBase, decoder_start_token_id=2, permute_sentence_ratio=1.0,
            mask_ratio=0.3, poisson_lambda=3.5) -> BatchEncoding:
    # convert list to dict and tensorize input
    batch = BatchEncoding(
        {k: np.array([examples[i][k] for i in range(len(examples))], dtype=int) for k, v in examples[0].items()}
    )

    batch["labels"] = batch["input_ids"].copy()
    batch["decoder_input_ids"] = shift_tokens_right(
        batch["labels"], tokenizer.pad_token_id, decoder_start_token_id
    )
    # permuting sentences
    do_permute = False
    if permute_sentence_ratio > 0.0:
        batch["input_ids"] = permute_sentences(batch["input_ids"], tokenizer, permute_sentence_ratio=permute_sentence_ratio)
        do_permute = True

    # masking span of tokens (text infilling in the paper)
    if mask_ratio:
        batch["input_ids"], batch["labels"] = span_mask_tokens(
            batch["input_ids"], batch["labels"], tokenizer, do_permute=do_permute, poisson_lambda=poisson_lambda, mask_ratio=mask_ratio
        )

    # ignore pad tokens
    batch["attention_mask"] = (batch["input_ids"] != tokenizer.pad_token_id).astype(int)
    batch["decoder_attention_mask"] = (batch["decoder_input_ids"] != tokenizer.pad_token_id).astype(int)
    return batch

def permute_sentences(input_ids, tokenizer, permute_sentence_ratio=1.0):
    """
    Shuffle sentences in each document.
    """
    results = input_ids.copy()

    # find end locations of sentences
    end_sentence_mask = input_ids == tokenizer.pad_token_id
    sentence_ends = np.argwhere(end_sentence_mask)
    sentence_ends[:, 1] += 1
    example_has_multiple_sentences, num_sentences = np.unique(sentence_ends[:, 0], return_counts=True)
    num_sentences_map = {sent_idx: count for sent_idx, count in zip(example_has_multiple_sentences, num_sentences)}

    num_to_permute = np.ceil(num_sentences * permute_sentence_ratio).astype(int)
    num_to_permute_map = {
        sent_idx: count for sent_idx, count in zip(example_has_multiple_sentences, num_to_permute)
    }

    sentence_ends = np.split(sentence_ends[:, 1], np.unique(sentence_ends[:, 0], return_index=True)[1][1:])
    sentence_ends_map = {sent_idx: count for sent_idx, count in zip(example_has_multiple_sentences, sentence_ends)}

    for i in range(input_ids.shape[0]):
        if i not in example_has_multiple_sentences:
            continue
        substitutions = np.random.permutation(num_sentences_map[i])[: num_to_permute_map[i]]
        ordering = np.arange(0, num_sentences_map[i])
        ordering[substitutions] = substitutions[np.random.permutation(num_to_permute_map[i])]

        # write shuffled sentences into results
        index = 0
        for j in ordering:
            sentence = input_ids[i, (sentence_ends_map[i][j - 1] if j > 0 else 0) : sentence_ends_map[i][j]]
            results[i, index : index + sentence.shape[0]] = sentence
            index += sentence.shape[0]
    return results


def span_mask_tokens(input_ids, labels, tokenizer, do_permute=True, poisson_lambda=3.5, mask_ratio=0.3):
    """
    Sampling text spans with span lengths drawn from a Poisson distribution and masking them.
    """
    special_tokens_mask_labels = [
        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
    ]
    special_tokens_mask_inputs = [
        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in input_ids.tolist()
    ]
    special_tokens_mask_labels = np.array(special_tokens_mask_labels, dtype=bool)
    special_tokens_mask_inputs = np.array(special_tokens_mask_inputs, dtype=bool)

    # determine how many tokens we need to mask in total
    is_token_mask = ~(input_ids == tokenizer.pad_token_id) & ~special_tokens_mask_inputs
    num_tokens_to_mask = int(math.ceil(is_token_mask.astype(float).sum() * mask_ratio))
    if num_tokens_to_mask == 0:
        return input_ids, labels

    # generate a sufficient number of span lengths
    span_lengths = np.random.poisson(lam=poisson_lambda, size=(num_tokens_to_mask,))
    while np.cumsum(span_lengths, 0)[-1] < num_tokens_to_mask:
        span_lengths = np.concatenate(
            [span_lengths, np.random.poisson(lam=poisson_lambda, size=(num_tokens_to_mask,))]
        )

    # remove all spans of length 0
    # note that BART inserts additional mask tokens where length == 0,
    # which we do not implement for now as it adds additional complexity
    span_lengths = span_lengths[span_lengths > 0]

    # trim to about num_tokens_to_mask tokens
    cutoff_idx = np.argmin(np.abs(np.cumsum(span_lengths, 0) - num_tokens_to_mask)) + 1
    span_lengths = span_lengths[:cutoff_idx]

    # randomly choose starting positions for masking
    token_indices = np.argwhere(is_token_mask == 1)
    span_starts = np.random.permutation(token_indices.shape[0])[: span_lengths.shape[0]]
    # prepare mask
    masked_indices = np.array(token_indices[span_starts])
    mask = np.full_like(input_ids, fill_value=False)

    # mask starting positions
    for mi in masked_indices:
        mask[tuple(mi)] = True
    span_lengths -= 1

    # fill up spans
    max_index = input_ids.shape[1] - 1
    remaining = (span_lengths > 0) & (masked_indices[:, 1] < max_index)
    while np.any(remaining):
        masked_indices[remaining, 1] += 1
        for mi in masked_indices:
            mask[tuple(mi)] = True
        span_lengths -= 1
        remaining = (span_lengths > 0) & (masked_indices[:, 1] < max_index)

    # place the mask tokens
    mask[np.where(special_tokens_mask_inputs)] = False
    input_ids[np.where(mask)] = tokenizer.mask_token_id
    if not do_permute:
        labels[np.where(mask == 0)] = -100
    else:
        labels[np.where(special_tokens_mask_labels)] = -100

    # remove mask tokens that are not starts of spans
    to_remove = (mask == 1) & np.roll((mask == 1), 1, 1)
    new_input_ids = np.full_like(input_ids, fill_value=tokenizer.pad_token_id)
    for i, example in enumerate(input_ids):
        new_example = example[~to_remove[i]]
        new_input_ids[i, : new_example.shape[0]] = new_example

    return new_input_ids, labels


def group_texts(examples, max_seq_length):
    # Concatenate all texts.
    concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
    # customize this part to your needs.
    if total_length >= max_seq_length:
        total_length = (total_length // max_seq_length) * max_seq_length
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
        for k, t in concatenated_examples.items()
    }
    return result


def main():
    pass

def get_n_mask_tokens(tokens, mask_token_id):
    unique, counts = np.unique(tokens, return_counts=True)
    counter = dict(zip(unique, counts))
    return counter[mask_token_id]


def get_n_nonspecial_tokens(tokens, all_special_ids):
    return len([t for t in tokens if t not in all_special_ids])


if __name__ == '__main__':
    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
    text = "On september 2nd the Group of Seven (G7) countries launched a new attempt to regain the advantage in the" \
           " West’s energy confrontation with Russia: imposing a price cap on purchases of Russian oil and oil" \
           " products, probably to take effect on December 5th. "
    encoded = tokenizer(text)
    input_ids = encoded["input_ids"]
    n_input_toks = get_n_nonspecial_tokens(input_ids, tokenizer.all_special_ids)
    print("DECODED INPUT", tokenizer.decode(input_ids))
    processed = collate([encoded], tokenizer)
    input_ids_out = processed["input_ids"].squeeze().tolist()
    n_output_toks = get_n_nonspecial_tokens(input_ids_out, tokenizer.all_special_ids)
    print("DECODED OUTPUT", tokenizer.decode(input_ids_out))

    n_masks_out = get_n_mask_tokens(input_ids_out, tokenizer.mask_token_id) + (n_input_toks-n_output_toks)
    print(f"MASK RATIO ({n_masks_out}/{n_input_toks})", n_masks_out/n_input_toks)
```

Output:

```
DECODED INPUT 
On september 2nd the Group of Seven (G7) countries launched a new attempt to regain the advantage in the West’s energy confrontation with Russia: imposing a price cap on purchases of Russian oil and oil products, probably to take effect on December 5th. 

DECODED OUTPUT 
On september 2nd the<mask>) countries launched a new attempt to regain the advantage in the West’s<mask> price cap on purchases of Russian oil and oil products, probably to take effect on December<mask>th. 
<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
MASK RATIO (17/57) 0.2982456140350877
```

Hey @BramVanroy! Thanks for posting this issue with such clear and concise code-snippets! In your opinion, would it be worth implementing these changes in the example script? Or are they specific to your use case? Perhaps you could summarise quickly the changes you made to computing the actual coverage! Thanks!

Hey @sanchit-gandhi. The issue was with how I calculated the coverage for reporting, so there was a mistake on my end not on yours!