Load examples from JSONLines Dataset without duplication.

src/torch_wae/cli/convert_dataset_to_web.py

@@ -61,7 +61,7 @@ def main(
     )
 
     output.mkdir(parents=True, exist_ok=True)
-    pattern = str(output / "%06d.tar")
+    pattern = str(output / "%04d.tar")
     max_size = size_shard * 1024**2
 
     with tqdm(total=n) as progress:

src/torch_wae/dataset.py

@@ -15,9 +15,17 @@ def __init__(self, path: Path) -> None:
         self.__path = path
 
     def __iter__(self) -> Iterator[Any]:
+        worker_info = data.get_worker_info()
+
         with self.__path.open() as f:
-            for line in f:
-                yield json.loads(line)
+            for i, line in enumerate(f):
+                if worker_info is None:
+                    yield json.loads(line)
+                else:
+                    worker_id = worker_info.id
+                    num_workers = worker_info.num_workers
+                    if i % num_workers == worker_id:
+                        yield json.loads(line)
 
 
 @dataclass(frozen=True)