Added new SST2 dataset class

requirements.txt

@@ -4,6 +4,9 @@ tqdm
 # Downloading data and other files
 requests
 
+# Torchdata
+git+https://github.com/pytorch/data.git
+
 # Optional NLP tools
 nltk
 spacy

setup.py

@@ -88,7 +88,10 @@ def run(self):
     license='BSD',
 
     install_requires=[
-        'tqdm', 'requests', pytorch_package_dep, 'numpy'
+        'tqdm', 'requests', pytorch_package_dep, 'numpy', 'torchdata==0.1.0a0+7772406'
+    ],
+    dependency_links=[
+        "https://github.com/pytorch/data.git#egg=torchdata",
     ],
     python_requires='>=3.5',
     classifiers=[

test/experimental/test_datasets.py

@@ -0,0 +1,32 @@
+import hashlib
+import json
+
+from torchtext.experimental.datasets import sst2
+
+from ..common.torchtext_test_case import TorchtextTestCase
+
+
+class TestDataset(TorchtextTestCase):
+    def test_sst2_dataset(self):
+        split = ("train", "dev", "test")
+        train_dp, dev_dp, test_dp = sst2.SST2(split=split)
+
+        # verify hashes of first line in dataset
+        self.assertEqual(
+            hashlib.md5(
+                json.dumps(next(iter(train_dp)), sort_keys=True).encode("utf-8")
+            ).hexdigest(),
+            sst2._FIRST_LINE_MD5["train"],
+        )
+        self.assertEqual(
+            hashlib.md5(
+                json.dumps(next(iter(dev_dp)), sort_keys=True).encode("utf-8")
+            ).hexdigest(),
+            sst2._FIRST_LINE_MD5["dev"],
+        )
+        self.assertEqual(
+            hashlib.md5(
+                json.dumps(next(iter(test_dp)), sort_keys=True).encode("utf-8")
+            ).hexdigest(),
+            sst2._FIRST_LINE_MD5["test"],
+        )

torchtext/experimental/datasets/__init__.py

@@ -1,3 +1,4 @@
 from . import raw
+from . import sst2
 
-__all__ = ['raw']
+__all__ = ["raw", "sst2"]

torchtext/experimental/datasets/sst2.py

@@ -0,0 +1,93 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import os
+
+from torchdata.datapipes.iter import (
+    HttpReader,
+    IterableWrapper,
+)
+from torchtext.data.datasets_utils import (
+    _add_docstring_header,
+    _create_dataset_directory,
+    _wrap_split_argument,
+)
+
+
+NUM_LINES = {
+    "train": 67349,
+    "dev": 872,
+    "test": 1821,
+}
+
+MD5 = "9f81648d4199384278b86e315dac217c"
+URL = "https://dl.fbaipublicfiles.com/glue/data/SST-2.zip"
+
+_EXTRACTED_FILES = {
+    "train": f"{os.sep}".join(["SST-2", "train.tsv"]),
+    "dev": f"{os.sep}".join(["SST-2", "dev.tsv"]),
+    "test": f"{os.sep}".join(["SST-2", "test.tsv"]),
+}
+
+_EXTRACTED_FILES_MD5 = {
+    "train": "da409a0a939379ed32a470bc0f7fe99a",
+    "dev": "268856b487b2a31a28c0a93daaff7288",
+    "test": "3230e4efec76488b87877a56ae49675a",
+}
+
+_FIRST_LINE_MD5 = {
+    "train": "2552b8cecd57b2e022ef23411c688fa8",
+    "dev": "1b0ffd6aa5f2bf0fd9840a5f6f1a9f07",
+    "test": "f838c81fe40bfcd7e42e9ffc4dd004f7",
+}
+
+DATASET_NAME = "SST2"
+
+
+@_add_docstring_header(num_lines=NUM_LINES, num_classes=2)
+@_create_dataset_directory(dataset_name=DATASET_NAME)
+@_wrap_split_argument(("train", "dev", "test"))
+def SST2(root, split):
+    return SST2Dataset(root, split).get_datapipe()
+
+
+class SST2Dataset:
+    """The SST2 dataset uses torchdata datapipes end-2-end.
+    To avoid download at every epoch, we cache the data on-disk
+    We do sanity check on dowloaded and extracted data
+    """
+
+    def __init__(self, root, split):
+        self.root = root
+        self.split = split
+
+    def get_datapipe(self):
+        # cache data on-disk
+        cache_dp = IterableWrapper([URL]).on_disk_cache(
+            HttpReader,
+            op_map=lambda x: (x[0], x[1].read()),
+            filepath_fn=lambda x: os.path.join(self.root, os.path.basename(x)),
+        )
+
+        # do sanity check
+        check_cache_dp = cache_dp.check_hash(
+            {os.path.join(self.root, "SST-2.zip"): MD5}, "md5"
+        )
+
+        # extract data from zip
+        extracted_files = check_cache_dp.read_from_zip()
+
+        # Filter extracted files and do sanity check
+        check_extracted_files = extracted_files.filter(
+            lambda x: self.split in x[0]
+        ).check_hash(
+            {
+                os.path.join(
+                    self.root, _EXTRACTED_FILES[self.split]
+                ): _EXTRACTED_FILES_MD5[self.split]
+            },
+            "md5",
+        )
+
+        # Parse CSV file and yield data samples
+        return check_extracted_files.parse_csv(skip_lines=1, delimiter="\t").map(
+            lambda x: (x[0], x[1])
+        )