Add ByT5 option to example run_t5_mlm_flax.py
Author: mapmeld
Pull request: yes

Small change adding ByT5 option to the Flax T5 training example

When model_type is `byt5`, use ByT5Tokenizer in place of T5TokenizerFast

Example: https://colab.research.google.com/drive/1WcDRPYyvuMZDbWuhsS3hTaVyXxjqryPz?usp=sharing

Comments:
I was trying this on TPU, but it ended with the following error message:

```bash
Traceback (most recent call last):
  File "run_t5_mlm_flax.py", line 728, in <module>
    model_inputs = data_collator(samples)
  File "run_t5_mlm_flax.py", line 276, in __call__
    batch["decoder_input_ids"] = shift_tokens_right(
  File "/home/stefan/transformers/src/transformers/models/t5/modeling_flax_t5.py", line 55, in shift_tokens_right
    shifted_input_ids = jax.ops.index_update(shifted_input_ids, (..., 0), decoder_start_token_id)
  File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/ops/scatter.py", line 352, in index_update
    return _scatter_update(
  File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/ops/scatter.py", line 64, in _scatter_update
    y = jnp.asarray(y)
  File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3082, in asarray
    return array(a, dtype=dtype, copy=False, order=order)
  File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3042, in array
    lax._check_user_dtype_supported(_inferred_dtype, "array")
  File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 6963, in _check_user_dtype_supported
    raise TypeError(msg)
TypeError: JAX only supports number and bool dtypes, got dtype object in array
```

I was using this config:

```
https://huggingface.co/google/byt5-base/raw/main/config.json
```

with the following parameters:

```bash
python run_t5_mlm_flax.py     --output_dir="${MODEL_DIR}"     --model_type="t5"     --config_name="${MODEL_DIR}"     --tokenizer_name="google/byt5-base"     --max_seq_length="512"     --per_device_train_batch_size="16"     --per_device_eval_batch_size="16"     --learning_rate="1e-3"     --weight_decay="0.001"     --warmup_steps="5000"     --overwrite_output_dir     --num_train_epochs="10"     --logging_steps="500"     --save_steps="2500"     --eval_steps="2500"     --train_file /mnt/datasets/train.txt     --validation_file /mnt/datasets/validation.txt
```


@patrickvonplaten do you have any hint how to fix this :thinking: 

> I was trying this on TPU, but it ended with the following error message:
> 
> ```shell
> Traceback (most recent call last):
>   File "run_t5_mlm_flax.py", line 728, in <module>
>     model_inputs = data_collator(samples)
>   File "run_t5_mlm_flax.py", line 276, in __call__
>     batch["decoder_input_ids"] = shift_tokens_right(
>   File "/home/stefan/transformers/src/transformers/models/t5/modeling_flax_t5.py", line 55, in shift_tokens_right
>     shifted_input_ids = jax.ops.index_update(shifted_input_ids, (..., 0), decoder_start_token_id)
>   File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/ops/scatter.py", line 352, in index_update
>     return _scatter_update(
>   File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/ops/scatter.py", line 64, in _scatter_update
>     y = jnp.asarray(y)
>   File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3082, in asarray
>     return array(a, dtype=dtype, copy=False, order=order)
>   File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 3042, in array
>     lax._check_user_dtype_supported(_inferred_dtype, "array")
>   File "/home/stefan/dev/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 6963, in _check_user_dtype_supported
>     raise TypeError(msg)
> TypeError: JAX only supports number and bool dtypes, got dtype object in array
> ```
> 
> I was using this config:
> 
> ```
> https://huggingface.co/google/byt5-base/raw/main/config.json
> ```
> 
> with the following parameters:
> 
> ```shell
> python run_t5_mlm_flax.py     --output_dir="${MODEL_DIR}"     --model_type="t5"     --config_name="${MODEL_DIR}"     --tokenizer_name="google/byt5-base"     --max_seq_length="512"     --per_device_train_batch_size="16"     --per_device_eval_batch_size="16"     --learning_rate="1e-3"     --weight_decay="0.001"     --warmup_steps="5000"     --overwrite_output_dir     --num_train_epochs="10"     --logging_steps="500"     --save_steps="2500"     --eval_steps="2500"     --train_file /mnt/datasets/train.txt     --validation_file /mnt/datasets/validation.txt
> ```
> 
> @patrickvonplaten do you have any hint how to fix this

Let me try!

@stefan-it I cannot reproduce the error. Can you try running the following:

```bash
./run_t5_mlm_flax.py \
     --output_dir="${MODEL_DIR}" \
     --model_type="t5" \
     --config_name="${MODEL_DIR}" \
     --tokenizer_name="google/byt5-base" \
     --max_seq_length="128" \
     --per_device_train_batch_size="1" \
     --per_device_eval_batch_size="1" \
     --learning_rate="1e-3" \
     --weight_decay="0.001" \
     --warmup_steps="5000" \
     --overwrite_output_dir \
     --num_train_epochs="10" \
     --logging_steps="500" \
     --save_steps="2500" \
     --eval_steps="2500" \
     --dataset_name="oscar" \
     --dataset_config_name="unshuffled_deduplicated_als"
```

using 
`https://huggingface.co/google/byt5-base/raw/main/config.json` as your config?

This uses a very small oscar dataset just to check that the script is correct. As you can see the script should run just fine.

Hi @patrickvonplaten , your command is working - even with a sequence length of 512 and a batch size of 16. I'll check my dataset now 😅 Maybe some lines are too short...

That's really interesting, I just filtered lines that contain less than five tokens and training is working. Thanks for your help :hugs: