class WebDataset(datasets.GeneratorBasedBuilder):
    DEFAULT_WRITER_BATCH_SIZE = 100
    IMAGE_EXTENSIONS: List[str]  # definition at the bottom of the script
    AUDIO_EXTENSIONS: List[str]  # definition at the bottom of the script
    VIDEO_EXTENSIONS: List[str]  # definition at the bottom of the script
    DECODERS: Dict[str, Callable[[Any], Any]]  # definition at the bottom of the script
    NUM_EXAMPLES_FOR_FEATURES_INFERENCE = 5

    @classmethod
    def _get_pipeline_from_tar(cls, tar_path, tar_iterator):
        current_example = {}
        fs: fsspec.AbstractFileSystem = fsspec.filesystem("memory")
        streaming_download_manager = datasets.StreamingDownloadManager()
        for filename, f in tar_iterator:
            example_key, field_name = base_plus_ext(filename)
            if example_key is None:
                continue
            if current_example and current_example["__key__"] != example_key:
                # reposition some keys in last position
                current_example["__key__"] = current_example.pop("__key__")
                current_example["__url__"] = current_example.pop("__url__")
                yield current_example
                current_example = {}
            current_example["__key__"] = example_key
            current_example["__url__"] = tar_path
            current_example[field_name.lower()] = f.read()
            if field_name.split(".")[-1] in SINGLE_FILE_COMPRESSION_EXTENSION_TO_PROTOCOL:
                fs.write_bytes(filename, current_example[field_name.lower()])
                extracted_file_path = streaming_download_manager.extract(f"memory://{filename}")
                with fsspec.open(extracted_file_path) as f:
                    current_example[field_name.lower()] = f.read()
                fs.delete(filename)
                data_extension = xbasename(extracted_file_path).split(".")[-1]
            else:
                data_extension = field_name.split(".")[-1]
            if data_extension in cls.DECODERS:
                current_example[field_name] = cls.DECODERS[data_extension](current_example[field_name])
        if current_example:
            yield current_example

    def _info(self) -> datasets.DatasetInfo:
        return datasets.DatasetInfo()

    def _split_generators(self, dl_manager):
        """We handle string, list and dicts in datafiles"""
        # Download the data files
        if not self.config.data_files:
            raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
        data_files = dl_manager.download(self.config.data_files)
        splits = []
        for split_name, tar_paths in data_files.items():
            if isinstance(tar_paths, str):
                tar_paths = [tar_paths]
            tar_iterators = [dl_manager.iter_archive(tar_path) for tar_path in tar_paths]
            splits.append(
                datasets.SplitGenerator(
                    name=split_name, gen_kwargs={"tar_paths": tar_paths, "tar_iterators": tar_iterators}
                )
            )
        if not self.info.features:
            # Get one example to get the feature types
            pipeline = self._get_pipeline_from_tar(tar_paths[0], tar_iterators[0])
            first_examples = list(islice(pipeline, self.NUM_EXAMPLES_FOR_FEATURES_INFERENCE))
            if any(example.keys() != first_examples[0].keys() for example in first_examples):
                raise ValueError(
                    "The TAR archives of the dataset should be in WebDataset format, "
                    "but the files in the archive don't share the same prefix or the same types."
                )
            pa_tables = [
                pa.Table.from_pylist(cast_to_python_objects([example], only_1d_for_numpy=True))
                for example in first_examples
            ]
            inferred_arrow_schema = pa.concat_tables(pa_tables, promote_options="default").schema
            features = datasets.Features.from_arrow_schema(inferred_arrow_schema)

            # Set Image types
            for field_name in first_examples[0]:
                extension = field_name.rsplit(".", 1)[-1]
                if extension in self.IMAGE_EXTENSIONS:
                    features[field_name] = datasets.Image()
            # Set Audio types
            for field_name in first_examples[0]:
                extension = field_name.rsplit(".", 1)[-1]
                if extension in self.AUDIO_EXTENSIONS:
                    features[field_name] = datasets.Audio()
            # Set Video types
            for field_name in first_examples[0]:
                extension = field_name.rsplit(".", 1)[-1]
                if extension in self.VIDEO_EXTENSIONS:
                    features[field_name] = datasets.Video()
            self.info.features = features

        return splits

    def _generate_examples(self, tar_paths, tar_iterators):
        image_field_names = [
            field_name for field_name, feature in self.info.features.items() if isinstance(feature, datasets.Image)
        ]
        audio_field_names = [
            field_name for field_name, feature in self.info.features.items() if isinstance(feature, datasets.Audio)
        ]
        all_field_names = list(self.info.features.keys())
        for tar_idx, (tar_path, tar_iterator) in enumerate(zip(tar_paths, tar_iterators)):
            for example_idx, example in enumerate(self._get_pipeline_from_tar(tar_path, tar_iterator)):
                for field_name in all_field_names:
                    if field_name not in example:
                        example[field_name] = None
                for field_name in image_field_names + audio_field_names:
                    if example[field_name] is not None:
                        example[field_name] = {
                            "path": example["__key__"] + "." + field_name,
                            "bytes": example[field_name],
                        }
                yield f"{tar_idx}_{example_idx}", example
/Users/nielsrogge/Documents/python_projecten/datasets/src/datasets/packaged_modules/webdataset/webdataset.py