NonMatchingSplitsSizesError on Flax BART with wiki summary dataset
Author: RissyRan
Pull request: no
Labels: name: Good Second Issue

name: bug

name: jax

### System Info

Platform: TPU
Python: python3.11

### Who can help?

@sanchit-gandhi

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. set up example based on [here](https://github.com/huggingface/transformers/tree/e5eb55b88b7200e81144eaf7a2bf3001b2bd08e7/examples/flax/summarization)
2. run command below on wiki_summary

```
          "cd /tmp/transformers/examples/flax/summarization &&"
          " JAX_PLATFORM_NAME=TPU python3 run_summarization_flax.py"
          " --model_name_or_path facebook/bart-base --tokenizer_name"
          " facebook/bart-base --dataset_name wiki_summary --do_train"
          " --do_eval --do_predict --predict_with_generate --learning_rate"
          " 5e-5 --warmup_steps 0 --output_dir=/tmp/transformers/bart-base-wiki"
          " --overwrite_output_dir --num_train_epochs 3 --max_source_length"
          " 512 --max_target_length 64  --per_device_train_batch_size=64 --per_device_eval_batch_size=64"
```

### Expected behavior

We got issue below:

```
[2024-03-11, 02:09:49 UTC] {logging_mixin.py:150} WARNING - 
Generating test split:  18%|█▊        | 1000/5638 [00:00<00:00, 8781.53 examples/s]
[2024-03-11, 02:09:49 UTC] {logging_mixin.py:150} WARNING - 
Generating test split:  37%|███▋      | 2074/5638 [00:00<00:00, 9863.29 examples/s]
[2024-03-11, 02:09:49 UTC] {logging_mixin.py:150} WARNING - 
Generating test split:  55%|█████▌    | 3112/5638 [00:00<00:00, 10087.86 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating test split:  74%|███████▍  | 4188/5638 [00:00<00:00, 10345.56 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating test split:  93%|█████████▎| 5253/5638 [00:00<00:00, 10450.02 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating test split: 100%|██████████| 5638/5638 [00:00<00:00, 10140.16 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating validation split:   0%|          | 0/5074 [00:00<?, ? examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating validation split:  20%|██        | 1035/5074 [00:00<00:00, 10319.92 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating validation split:  42%|████▏     | 2110/5074 [00:00<00:00, 10566.07 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating validation split:  63%|██████▎   | 3175/5074 [00:00<00:00, 10597.24 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating validation split:  84%|████████▍ | 4253/5074 [00:00<00:00, 10663.04 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - 
Generating validation split: 100%|██████████| 5074/5074 [00:00<00:00, 10458.45 examples/s]
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - Traceback (most recent call last):
  File "/tmp/transformers/examples/flax/summarization/run_summarization_flax.py", line 1031, in <module>
    main()
  File "/tmp/transformers/examples/flax/summarization/run_summarization_flax.py", line 499, in main
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING -     dataset = load_dataset(
  File "/home/ml-auto-solutions/.local/lib/python3.10/site-packages/datasets/load.py", line 2582, in load_dataset
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING -     builder_instance.download_and_prepare(
  File "/home/ml-auto-solutions/.local/lib/python3.10/site-packages/datasets/builder.py", line 1005, in download_and_prepare
    self._download_and_prepare(
  File "/home/ml-auto-solutions/.local/lib/python3.10/site-packages/datasets/builder.py", line 1767, in _download_and_prepare
    super()._download_and_prepare(
  File "/home/ml-auto-solutions/.local/lib/python3.10/site-packages/datasets/builder.py", line 1118, in _download_and_prepare
    verify_splits(self.info.splits, split_dict)
  File "/home/ml-auto-solutions/.local/lib/python3.10/site-packages/datasets/utils/info_utils.py", line 101, in verify_splits
    raise NonMatchingSplitsSizesError(str(bad_splits))
datasets.utils.info_utils.NonMatchingSplitsSizesError: [{'expected': SplitInfo(name='train', num_bytes=207186608, num_examples=45654, shard_lengths=None, dataset_name=None), 'recorded': SplitInfo(name='train', num_bytes=0, num_examples=0, shard_lengths=None, dataset_name='wiki_summar
[2024-03-11, 02:09:50 UTC] {logging_mixin.py:150} WARNING - y')}]
[2024-03-11, 02:09:54 UTC] {taskinstance.py:1826} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/opt/python3.11/lib/python3.11/site-packages/airflow/decorators/base.py", line 220, in execute
    return_value = super().execute(context)
                   ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/python3.11/lib/python3.11/site-packages/airflow/operators/python.py", line 181, in execute
    return_value = self.execute_callable()
                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/python3.11/lib/python3.11/site-packages/airflow/operators/python.py", line 198, in execute_callable
    return self.python_callable(*self.op_args, **self.op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/gcs/dags/xlml/utils/tpu.py", line 402, in ssh_tpu
    ssh_group.run(cmds, env=env)
  File "/opt/python3.11/lib/python3.11/site-packages/fabric/group.py", line 116, in run
    return self._do("run", *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/python3.11/lib/python3.11/site-packages/fabric/group.py", line 282, in _do
    raise GroupException(results)
```

Based on this [thread](https://discuss.huggingface.co/t/nonmatchingsplitssizeserror/30033), a few flags can be added to load_dataset().

Comments:
Gentle ping @sanchit-gandhi 

Another ping @sanchit-gandhi 

Another another ping @sanchit-gandhi. Could you nominate someone to take over this whilst you're away? 

cc @ylacombe @eustlb

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.

Sorry for the ping Yoach/Eustache, I misread this for audio when it was flax

We won't have the bandwidth to fix this as of now, but if any reader has the energy and bandwidth to try and solve this, we'd welcome a PR!

@LysandreJik 
I've looked into this issue and found the issue. It appears the dataset (https://huggingface.co/datasets/m3hrdadfi/wiki_summary) is hosted on Google Drive. While this isn’t an issue in itself, the train dataset is quite large (197 MB), triggering Google Drive's "can't scan this file for viruses" warning. This forces users to manually confirm by clicking "download anyway" before the download can proceed. This is likely why only the train dataset is affected—since the validation and test datasets are smaller, they don’t trigger this warning.

I'm not sure which fix would be best. One option could be to modify the download logic itself that lives inside the Hugging Face datasets package, specifically in the DownloadManager class. We could add additional code to handle this Google Drive-specific issue. However, that might introduce new bugs down the line.

Alternatively, we could adjust the logic for providing the download link to a direct link within the DatasetModule (inside of the datasets package), bypassing the warning altogether. The trade off here is that it might require fetching the link, meaning we’d have to make an additional request before the actual download.

I’d appreciate any thoughts or suggestions.