Add Shuffle and sharding datapipes to datasets (#1729)

Author: parmeet

test/datasets/common.py

@@ -0,0 +1,25 @@
+from parameterized import parameterized
+from torch.utils.data.graph import traverse
+from torch.utils.data.graph_settings import get_all_graph_pipes
+from torchdata.datapipes.iter import Shuffler, ShardingFilter
+from torchtext.datasets import DATASETS
+
+from ..common.torchtext_test_case import TorchtextTestCase
+
+
+class TestShuffleShardDatasetWrapper(TorchtextTestCase):
+    # Note that for order i.e shuffle before sharding, TorchData will provide linter warning
+    # Modify this test when linter warning is available
+    @parameterized.expand(list(DATASETS.items()))
+    def test_shuffle_shard_wrapper(self, dataset_name, dataset_fn):
+        dp = dataset_fn()
+        if type(dp) == tuple:
+            dp = list(dp)
+        else:
+            dp = [dp]
+
+        for dp_split in dp:
+            dp_graph = get_all_graph_pipes(traverse(dp_split))
+            for annotation_dp_type in [Shuffler, ShardingFilter]:
+                if not any(isinstance(dp, annotation_dp_type) for dp in dp_graph):
+                    raise AssertionError(f"The dataset doesn't contain a {annotation_dp_type.__name__}() datapipe.")

torchtext/datasets/ag_news.py

@@ -71,4 +71,4 @@ def AG_NEWS(root: str, split: Union[Tuple[str], str]):
     cache_dp = cache_dp.end_caching(mode="wb", same_filepath_fn=True)
 
     data_dp = FileOpener(cache_dp, encoding="utf-8")
-    return data_dp.parse_csv().map(fn=_modify_res)
+    return data_dp.parse_csv().map(fn=_modify_res).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/amazonreviewfull.py

@@ -90,4 +90,4 @@ def AmazonReviewFull(root: str, split: Union[Tuple[str], str]):
     cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8")
-    return data_dp.parse_csv().map(fn=_modify_res)
+    return data_dp.parse_csv().map(fn=_modify_res).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/amazonreviewpolarity.py

@@ -87,4 +87,4 @@ def AmazonReviewPolarity(root: str, split: Union[Tuple[str], str]):
     cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8")
-    return data_dp.parse_csv().map(fn=_modify_res)
+    return data_dp.parse_csv().map(fn=_modify_res).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/cc100.py

@@ -176,4 +176,4 @@ def CC100(root: str, language_code: str = "en"):
     cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb")
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8").readlines(return_path=False)
-    return data_dp.map(partial(_modify_res, language_code))
+    return data_dp.map(partial(_modify_res, language_code)).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/conll2000chunking.py

@@ -80,4 +80,4 @@ def CoNLL2000Chunking(root: str, split: Union[Tuple[str], str]):
     cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8")
-    return data_dp.readlines().read_iob(sep=" ")
+    return data_dp.readlines().read_iob(sep=" ").shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/dbpedia.py

@@ -86,4 +86,4 @@ def DBpedia(root: str, split: Union[Tuple[str], str]):
     cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8")
-    return data_dp.parse_csv().map(fn=_modify_res)
+    return data_dp.parse_csv().map(fn=_modify_res).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/enwik9.py

@@ -59,4 +59,4 @@ def EnWik9(root: str):
     cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8")
-    return data_dp.readlines(return_path=False)
+    return data_dp.readlines(return_path=False).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/imdb.py

@@ -111,4 +111,4 @@ def filter_imdb_data(key, fname):
 
     data_dp = FileOpener(cache_decompressed_dp, encoding="utf-8")
     # get label from cache file, eg. "aclImdb_v1/train/neg" -> "neg"
-    return data_dp.readlines().map(_modify_res)
+    return data_dp.readlines().map(_modify_res).shuffle().set_shuffle(False).sharding_filter()

torchtext/datasets/iwslt2016.py

@@ -322,4 +322,4 @@ def IWSLT2016(
     src_lines = src_data_dp.readlines(return_path=False, strip_newline=False)
     tgt_lines = tgt_data_dp.readlines(return_path=False, strip_newline=False)
 
-    return src_lines.zip(tgt_lines)
+    return src_lines.zip(tgt_lines).shuffle().set_shuffle(False).sharding_filter()