

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
Note that we cannot add "{extra_id_}" to the string directly
as the Byte tokenizer would incorrectly merge the tokens
For ByT5, we need to work directly on the character level
Contrary to T5, ByT5 does not use sentinel tokens for masking, but instead
uses final utf character ids.
UTF-8 is represented by 8 bits and ByT5 has 3 special tokens.
=> There are 2**8+2 = 259 input ids and mask tokens count down from index 258.
=> 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:]])
input_ids
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()
output_ids
[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)
output_string
['', 'is the one who does', ' in the disco', 'in the park. The dog is the one who does a ball in', ' in the park.']