| """WebDataset decoding and local shard resolution.""" | |
| import json | |
| import os | |
| from itertools import islice | |
| from typing import Any, Optional | |
| import datasets | |
| import pyarrow as pa | |
| from datasets.data_files import ( | |
| _get_origin_metadata, | |
| DataFilesList as DataFilesListBase, | |
| DownloadConfig, | |
| SingleOriginMetadata, | |
| ) | |
| from datasets.features.features import cast_to_python_objects | |
| from datasets.packaged_modules.webdataset.webdataset import WebDataset as WebDatasetBase | |
| class WebDataset(WebDatasetBase): | |
| def _split_generators(self, dl_manager: Any) -> list[datasets.SplitGenerator]: | |
| """Handle string/list/dict data files and infer features from one sample stream.""" | |
| 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: list[datasets.SplitGenerator] = [] | |
| first_tar_path: str | None = None | |
| first_tar_iterator: Any | None = None | |
| 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] | |
| if first_tar_path is None: | |
| first_tar_path = tar_paths[0] | |
| first_tar_iterator = tar_iterators[0] | |
| splits.append( | |
| datasets.SplitGenerator( | |
| name=split_name, gen_kwargs={"tar_paths": tar_paths, "tar_iterators": tar_iterators} | |
| ) | |
| ) | |
| if not self.info.features: | |
| if first_tar_path is None or first_tar_iterator is None: | |
| raise ValueError("No TAR files were resolved from data_files.") | |
| pipeline = self._get_pipeline_from_tar(first_tar_path, first_tar_iterator) | |
| first_examples = list(islice(pipeline, self.NUM_EXAMPLES_FOR_FEATURES_INFERENCE)) | |
| if not first_examples: | |
| raise ValueError("Unable to infer features: no examples found in TAR pipeline.") | |
| 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) | |
| for field_name in first_examples[0]: | |
| extension = field_name.rsplit(".", 1)[-1] | |
| if extension in self.IMAGE_EXTENSIONS: | |
| features[field_name] = datasets.Image() | |
| for field_name in first_examples[0]: | |
| extension = field_name.rsplit(".", 1)[-1] | |
| if extension in self.VIDEO_EXTENSIONS: | |
| features[field_name] = datasets.Video(decode=False) | |
| for field_name in first_examples[0]: | |
| extension = field_name.rsplit(".", 1)[-1] | |
| if extension in ["json"]: | |
| features[field_name] = datasets.Value("string") | |
| self.info.features = features | |
| return splits | |
| def json_loads(data: str | bytes | bytearray) -> str: | |
| """Parse JSON payload and normalize it back to a string value.""" | |
| data_obj = json.loads(data) | |
| json_str = json.dumps(data_obj) | |
| return json_str | |
| DECODERS = { | |
| "json": json_loads, | |
| } | |
| WebDataset.DECODERS = DECODERS | |
| def _get_single_origin_metadata( | |
| data_file: str, | |
| download_config: Optional[DownloadConfig] = None, | |
| ) -> SingleOriginMetadata: | |
| _ = data_file, download_config | |
| return () | |
| class DataFilesList(DataFilesListBase): | |
| def from_patterns( | |
| cls, | |
| patterns: list[str], | |
| base_path: Optional[str] = None, | |
| allowed_extensions: Optional[list[str]] = None, | |
| download_config: Optional[DownloadConfig] = None, | |
| ) -> "DataFilesList": | |
| _ = base_path, allowed_extensions | |
| data_files = patterns | |
| if not data_files: | |
| raise ValueError("`patterns` must contain at least one path.") | |
| if any("://" in str(path) for path in data_files): | |
| raise ValueError("Dataset manifests must list local TAR files, not remote URLs.") | |
| missing_files = [path for path in data_files if not os.path.exists(path)] | |
| if missing_files: | |
| preview = ", ".join(missing_files[:3]) | |
| if len(missing_files) > 3: | |
| preview += f" ... (+{len(missing_files) - 3} more)" | |
| raise FileNotFoundError(f"Missing data file(s): {preview}") | |
| origin_metadata = _get_origin_metadata(data_files, download_config=download_config) | |
| return cls(data_files, origin_metadata) | |
| datasets.data_files._get_single_origin_metadata = _get_single_origin_metadata | |
| datasets.data_files.DataFilesList = DataFilesList | |
| datasets.packaged_modules.webdataset.WebDataset = WebDataset | |