[WebDataset] Support compressed files (#6931)

* support compressed files for in webdataset

* ignore windows behavior

* again

Author: lhoestq

src/datasets/features/features.py

@@ -312,21 +312,17 @@ def _cast_to_python_objects(obj: Any, only_1d_for_numpy: bool, optimize_list_cas
                 True,
             )
     elif config.TORCH_AVAILABLE and "torch" in sys.modules and isinstance(obj, torch.Tensor):
+        if obj.dtype == torch.bfloat16:
+            return _cast_to_python_objects(
+                obj.detach().to(torch.float).cpu().numpy(),
+                only_1d_for_numpy=only_1d_for_numpy,
+                optimize_list_casting=optimize_list_casting,
+            )[0], True
         if obj.ndim == 0:
             return obj.detach().cpu().numpy()[()], True
         elif not only_1d_for_numpy or obj.ndim == 1:
             return obj.detach().cpu().numpy(), True
         else:
-            if obj.dtype == torch.bfloat16:
-                return (
-                    [
-                        _cast_to_python_objects(
-                            x, only_1d_for_numpy=only_1d_for_numpy, optimize_list_casting=optimize_list_casting
-                        )[0]
-                        for x in obj.detach().to(torch.float).cpu().numpy()
-                    ],
-                    True,
-                )
             return (
                 [
                     _cast_to_python_objects(

src/datasets/packaged_modules/webdataset/webdataset.py

@@ -3,11 +3,13 @@
 from itertools import islice
 from typing import Any, Callable, Dict, List
 
+import fsspec
 import numpy as np
 import pyarrow as pa
 
 import datasets
 from datasets.features.features import cast_to_python_objects
+from datasets.utils.file_utils import SINGLE_FILE_COMPRESSION_EXTENSION_TO_PROTOCOL, xbasename
 
 
 logger = datasets.utils.logging.get_logger(__name__)
@@ -23,6 +25,8 @@ class WebDataset(datasets.GeneratorBasedBuilder):
     @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:
             if "." in filename:
                 example_key, field_name = filename.split(".", 1)
@@ -32,8 +36,17 @@ def _get_pipeline_from_tar(cls, tar_path, tar_iterator):
                 current_example["__key__"] = example_key
                 current_example["__url__"] = tar_path
                 current_example[field_name.lower()] = f.read()
-                if field_name.split(".")[-1] in cls.DECODERS:
-                    current_example[field_name] = cls.DECODERS[field_name.split(".")[-1]](current_example[field_name])
+                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
 

src/datasets/utils/file_utils.py

@@ -762,6 +762,9 @@ def readline(f: io.RawIOBase):
     # archive compression
     "zip": "zip",
 }
+SINGLE_FILE_COMPRESSION_EXTENSION_TO_PROTOCOL = {
+    fs_class.extension.lstrip("."): fs_class.protocol for fs_class in COMPRESSION_FILESYSTEMS
+}
 SINGLE_FILE_COMPRESSION_PROTOCOLS = {fs_class.protocol for fs_class in COMPRESSION_FILESYSTEMS}
 SINGLE_SLASH_AFTER_PROTOCOL_PATTERN = re.compile(r"(?<!:):/")
 

tests/packaged_modules/test_webdataset.py

@@ -10,6 +10,16 @@
 from ..utils import require_pil, require_sndfile, require_torch
 
 
+@pytest.fixture
+def gzipped_text_wds_file(tmp_path, text_gz_path):
+    filename = tmp_path / "file.tar"
+    num_examples = 3
+    with tarfile.open(str(filename), "w") as f:
+        for example_idx in range(num_examples):
+            f.add(text_gz_path, f"{example_idx:05d}.txt.gz")
+    return str(filename)
+
+
 @pytest.fixture
 def image_wds_file(tmp_path, image_file):
     json_file = tmp_path / "data.json"
@@ -64,6 +74,29 @@ def tensor_wds_file(tmp_path, tensor_file):
     return str(filename)
 
 
+@require_pil
+def test_gzipped_text_webdataset(gzipped_text_wds_file, text_path):
+    data_files = {"train": [gzipped_text_wds_file]}
+    webdataset = WebDataset(data_files=data_files)
+    split_generators = webdataset._split_generators(DownloadManager())
+    assert webdataset.info.features == Features(
+        {
+            "__key__": Value("string"),
+            "__url__": Value("string"),
+            "txt.gz": Value("string"),
+        }
+    )
+    assert len(split_generators) == 1
+    split_generator = split_generators[0]
+    assert split_generator.name == "train"
+    generator = webdataset._generate_examples(**split_generator.gen_kwargs)
+    _, examples = zip(*generator)
+    assert len(examples) == 3
+    assert isinstance(examples[0]["txt.gz"], str)
+    with open(text_path, "r") as f:
+        assert examples[0]["txt.gz"].replace("\r\n", "\n") == f.read().replace("\r\n", "\n")
+
+
 @require_pil
 def test_image_webdataset(image_wds_file):
     import PIL.Image