GitHub issue: Webdataset: KeyError: 'png' on some datasets when streaming
Author: lhoestq (MEMBER)
State: open | Created: 2024-05-07 13:09:02

reported at https://huggingface.co/datasets/tbone5563/tar_images/discussions/1

```python
>>> from datasets import load_dataset
>>> ds = load_dataset("tbone5563/tar_images")
Downloading data: 100%
 1.41G/1.41G [00:48<00:00, 17.2MB/s]
Downloading data: 100%
 619M/619M [00:11<00:00, 57.4MB/s]
Generating train split: 
 970/0 [00:02<00:00, 534.94 examples/s]
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/datasets/builder.py](https://localhost:8080/#) in _prepare_split_single(self, gen_kwargs, fpath, file_format, max_shard_size, split_info, check_duplicate_keys, job_id)
   1747                 _time = time.time()
-> 1748                 for key, record in generator:
   1749                     if max_shard_size is not None and writer._num_bytes > max_shard_size:

7 frames
[/usr/local/lib/python3.10/dist-packages/datasets/packaged_modules/webdataset/webdataset.py](https://localhost:8080/#) in _generate_examples(self, tar_paths, tar_iterators)
    108                 for field_name in image_field_names + audio_field_names:
--> 109                     example[field_name] = {"path": example["__key__"] + "." + field_name, "bytes": example[field_name]}
    110                 yield f"{tar_idx}_{example_idx}", example

KeyError: 'png'

The above exception was the direct cause of the following exception:

DatasetGenerationError                    Traceback (most recent call last)
[<ipython-input-2-8e0fbb7badc9>](https://localhost:8080/#) in <cell line: 3>()
      1 from datasets import load_dataset
      2 
----> 3 ds = load_dataset("tbone5563/tar_images")

[/usr/local/lib/python3.10/dist-packages/datasets/load.py](https://localhost:8080/#) in load_dataset(path, name, data_dir, data_files, split, cache_dir, features, download_config, download_mode, verification_mode, ignore_verifications, keep_in_memory, save_infos, revision, token, use_auth_token, task, streaming, num_proc, storage_options, trust_remote_code, **config_kwargs)
   2607 
   2608     # Download and prepare data
-> 2609     builder_instance.download_and_prepare(
   2610         download_config=download_config,
   2611         download_mode=download_mode,

[/usr/local/lib/python3.10/dist-packages/datasets/builder.py](https://localhost:8080/#) in download_and_prepare(self, output_dir, download_config, download_mode, verification_mode, ignore_verifications, try_from_hf_gcs, dl_manager, base_path, use_auth_token, file_format, max_shard_size, num_proc, storage_options, **download_and_prepare_kwargs)
   1025                         if num_proc is not None:
   1026                             prepare_split_kwargs["num_proc"] = num_proc
-> 1027                         self._download_and_prepare(
   1028                             dl_manager=dl_manager,
   1029                             verification_mode=verification_mode,

[/usr/local/lib/python3.10/dist-packages/datasets/builder.py](https://localhost:8080/#) in _download_and_prepare(self, dl_manager, verification_mode, **prepare_splits_kwargs)
   1787 
   1788     def _download_and_prepare(self, dl_manager, verification_mode, **prepare_splits_kwargs):
-> 1789         super()._download_and_prepare(
   1790             dl_manager,
   1791             verification_mode,

[/usr/local/lib/python3.10/dist-packages/datasets/builder.py](https://localhost:8080/#) in _download_and_prepare(self, dl_manager, verification_mode, **prepare_split_kwargs)
   1120             try:
   1121                 # Prepare split will record examples associated to the split
-> 1122                 self._prepare_split(split_generator, **prepare_split_kwargs)
   1123             except OSError as e:
   1124                 raise OSError(

[/usr/local/lib/python3.10/dist-packages/datasets/builder.py](https://localhost:8080/#) in _prepare_split(self, split_generator, check_duplicate_keys, file_format, num_proc, max_shard_size)
   1625             job_id = 0
   1626             with pbar:
-> 1627                 for job_id, done, content in self._prepare_split_single(
   1628                     gen_kwargs=gen_kwargs, job_id=job_id, **_prepare_split_args
   1629                 ):

[/usr/local/lib/python3.10/dist-packages/datasets/builder.py](https://localhost:8080/#) in _prepare_split_single(self, gen_kwargs, fpath, file_format, max_shard_size, split_info, check_duplicate_keys, job_id)
   1782             if isinstance(e, SchemaInferenceError) and e.__context__ is not None:
   1783                 e = e.__context__
-> 1784             raise DatasetGenerationError("An error occurred while generating the dataset") from e
   1785 
   1786         yield job_id, True, (total_num_examples, total_num_bytes, writer._features, num_shards, shard_lengths)

DatasetGenerationError: An error occurred while generating the dataset
```