How to do mask prediction with ByT5?
Author: patrickvonplaten
Pull request: no

# 🚀 Feature request

It is currently not very clear on how to do mask prediction using ByT5.
For T5, we, *e.g.* know that the first 100 sentinel tokens are used as masks, e.g. the following code snippet somewhat shows how `t5-base` does mask denoising: 

```python
from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("t5-base")
model = T5ForConditionalGeneration.from_pretrained("t5-base")

input_ids = tokenizer("The dog <extra_id_0> a ball <extra_id_1> park", return_tensors="pt").input_ids
print("Input_ids", input_ids)  # <- gives tensor([[   37,  1782, 32099,     3,     9,  1996, 32098,  2447,     1]])


output_ids = model.generate(input_ids)
print("Output ids", output_ids)  # <- gives tensor([[    0, 32099,    19, 32098,    16,     8, 32097,    19, 32096,    19,
         32095,    16,     8, 32094,    19, 32093,    16,     8, 32092,    16]])


print("Output string", tokenizer.batch_decode(output_ids))  # <- gives ['<pad> <extra_id_0> is<extra_id_1> in the<extra_id_2> is<extra_id_3> is<extra_id_4> in the<extra_id_5> is<extra_id_6> in the<extra_id_7> in']
```

Here we can see that the model learns to do the correct masking pattern and also gives somewhat sensible predictions for `<extra_id_1>` and `<extra_id_2>`, being "is" and "in the".

Since the ByT5 paper states: 

```bash
Second, we modify the pre-training task. mT5 uses the “span corruption” pre-training objective first proposed by Raffel et al. (2020) where spans of tokens in unlabeled text data are replaced with a single “sentinel” ID and the model must fill in the missing spans. Rather than adding 100 new tokens for the sentinels, we find it sufficient to reuse the final 100 byte IDs. While mT5 uses an average span length of 3 subword tokens, we find that masking longer byte-spans is valuable. Specifically, we set our mean mask span length to 20 bytes, and show ablations of this value in section 6.
```

IMO one should use the masking ids starting from 258 (256 + 3 special tokens - 1). E.g. the following code snippet should work more or less for ByT5:

```python
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

tokenizer = AutoTokenizer.from_pretrained("google/byt5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("google/byt5-base")

input_ids_prompt = "The dog chases a ball in the park."

input_ids = tokenizer(input_ids_prompt).input_ids

# Now Mask
# Note that we can add "{extra_id_...}" to the string directly 
# as the Byte tokenizer would incorrectly merge the tokens
# We need to work on the character level directly here
# => mask to "The dog [258]a ball [257]park."
input_ids = torch.tensor([input_ids[:8] + [258] + input_ids[14:21] + [257] + input_ids[28:]])
print("In IDs", input_ids)  # <- gives tensor([[ 87, 107, 104,  35, 103, 114, 106,  35, 258,  35, 100,  35, 101, 100,
         111, 111, 257,  35, 115, 100, 117, 110,  49,   1]])

# ByT5 produces only one char at a time so we need to produce many more output characters here -> set `max_length=100`.
output_ids = model.generate(input_ids, max_length=100)[0].tolist()

print("Out IDs", output_ids)  
"""^ gives [0, 258, 108, 118,  35, 119, 107, 104,  35, 114, 113, 104,  35, 122,
         107, 114,  35, 103, 114, 104, 118, 257,  35, 108, 113,  35, 119, 107,
         104,  35, 103, 108, 118, 102, 114, 256, 108, 113,  35, 119, 107, 104,
          35, 115, 100, 117, 110,  49,  35,  87, 107, 104,  35, 103, 114, 106,
          35, 108, 118,  35, 119, 107, 104,  35, 114, 113, 104,  35, 122, 107,
         114,  35, 103, 114, 104, 118,  35, 100,  35, 101, 100, 111, 111,  35,
         108, 113, 255,  35, 108, 113,  35, 119, 107, 104,  35, 115, 100, 117,
         110,  49] <- Note how 258 descends to 257, 256, 255
"""

# Now we need to split on the sentinel tokens, let's write a short loop for this
output_ids_list = []
start_token = 0
sentinel_token = 258
while sentinel_token in output_ids:
    split_idx = output_ids.index(sentinel_token)
    output_ids_list.append(output_ids[start_token: split_idx])
    start_token = split_idx
    sentinel_token -= 1

output_ids_list.append(output_ids[start_token:])

output_string = tokenizer.batch_decode(output_ids_list)
print("Out string", output_string)  # <- gives ['<pad>', 'is the one who does', ' in the disco', 'in the park. The dog is the one who does a ball in', ' in the park.']
```

We can see here that this also makes somewhat sense, e.g. "The dog [is the one who does] a ball [in the disco park]"

I think it would be a good idea to document this somewhere in our Transformers docs. Happy to do so.

@craffel lintingxue - could you take a quick look at whether my description above is correct here? 

E.g. for ByT5 the masking starts with the token `258` and then descends down to at most `257` no? 

Comments:
Also we should probably change the ByT5 tokenizer configs: https://huggingface.co/google/byt5-small/blob/main/tokenizer_config.json as its at best very confusing to write that the token ids from 259 to 383 are sentinel tokens. This is just wrong no? There are no sentinel tokens and the last 124 tokens are just there to round the number up to 3 * 2**7 = 384 for TPU no?

Yes, this is correct to the best of my recollection - I think we just forgot to allocate additional sentinel token IDs. @adarob can probably confirm.

Putting it in the official ByT5 docs here: https://github.com/huggingface/transformers/pull/16646

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.