fix webdataset pickling

src/datasets/utils/file_utils.py

@@ -1567,18 +1567,20 @@ def xxml_dom_minidom_parse(filename_or_file, download_config: Optional[DownloadC
 class _IterableFromGenerator(TrackedIterable):
     """Utility class to create an iterable from a generator function, in order to reset the generator when needed."""
 
-    def __init__(self, generator: Callable, *args, **kwargs):
+    def __init__(self, generator: Callable, *args):
         super().__init__()
         self.generator = generator
         self.args = args
-        self.kwargs = kwargs
 
     def __iter__(self):
-        for x in self.generator(*self.args, **self.kwargs):
+        for x in self.generator(*self.args):
             self.last_item = x
             yield x
         self.last_item = None
 
+    def __reduce__(self):
+        return (self.__class__, (self.generator, *self.args))
+
 
 class ArchiveIterable(_IterableFromGenerator):
     """An iterable of (path, fileobj) from a TAR archive, used by `iter_archive`"""

src/datasets/utils/track.py

@@ -44,6 +44,6 @@ def __init__(self) -> None:
 
     def __repr__(self) -> str:
         if self.last_item is None:
-            super().__repr__()
+            return super().__repr__()
         else:
             return f"{self.__class__.__name__}(current={self.last_item})"