Re-write language_modeling datasets (PennTreebank, WikiText103, WikiText2)

README.rst

@@ -129,6 +129,16 @@ Others are planned or a work in progress:
 
 See the ``test`` directory for examples of dataset usage.
 
+Legacy Code
+===========
+
+We have currently retired several datasets and moved them under ```torchtext.legacy```:
+
+* Sentiment analysis: IMDb
+* Language modeling: abstract class + WikiText-2, WikiText103, PennTreebank
+
+These datasets are re-written with a new pattern that is introduced in `Release v0.5.0 <https://github.com/pytorch/text/releases>`_.
+
 Disclaimer on Datasets
 ======================
 

docs/source/legacy/datasets.rst

@@ -0,0 +1,69 @@
+torchtext.legacy.datasets
+====================
+
+.. currentmodule:: torchtext.legacy.datasets
+
+TorchText legacy datasets.
+
+All datasets are subclasses of :class:`torchtext.data.Dataset`, which
+inherits from :class:`torch.utils.data.Dataset` i.e, they have ``split`` and
+``iters`` methods implemented.
+
+General use cases are as follows:
+
+Approach 1, ``splits``: ::
+
+    # set up fields
+    TEXT = data.Field(lower=True, include_lengths=True, batch_first=True)
+    LABEL = data.Field(sequential=False)
+
+    # make splits for data
+    train, test = datasets.IMDB.splits(TEXT, LABEL)
+
+    # build the vocabulary
+    TEXT.build_vocab(train, vectors=GloVe(name='6B', dim=300))
+    LABEL.build_vocab(train)
+
+    # make iterator for splits
+    train_iter, test_iter = data.BucketIterator.splits(
+        (train, test), batch_size=3, device=0)
+
+Approach 2, ``iters``: ::
+
+    # use default configurations
+    train_iter, test_iter = datasets.IMDB.iters(batch_size=4)
+
+The following datasets are available:
+
+.. contents:: Datasets
+    :local:
+
+
+Language Modeling
+^^^^^^^^^^^^^^^^^
+
+Language modeling datasets are subclasses of ``LanguageModelingDataset`` class.
+
+.. autoclass:: LanguageModelingDataset
+  :members: __init__
+
+
+WikiText-2
+~~~~~~~~~~
+
+.. autoclass:: WikiText2
+  :members: splits, iters
+
+
+WikiText103
+~~~~~~~~~~~
+
+.. autoclass:: WikiText103
+  :members: splits, iters
+
+
+PennTreebank
+~~~~~~~~~~~~
+
+.. autoclass:: PennTreebank
+  :members: splits, iters

test/data/test_builtin_datasets.py

@@ -1,6 +1,6 @@
 import os
+import shutil
 import torchtext.data as data
-from torchtext.datasets import WikiText2, PennTreebank
 from torchtext.datasets import AG_NEWS
 
 from ..common.test_markers import slow
@@ -10,11 +10,14 @@
 def conditional_remove(f):
     if os.path.isfile(f):
         os.remove(f)
+    elif os.path.isdir(f):
+        shutil.rmtree(f)
 
 
 class TestDataset(TorchtextTestCase):
     @slow
-    def test_wikitext2(self):
+    def test_wikitext2_legacy(self):
+        from torchtext.legacy.datasets import WikiText2
         # smoke test to ensure wikitext2 works properly
         ds = WikiText2
         TEXT = data.Field(lower=True, batch_first=True)
@@ -27,12 +30,30 @@ def test_wikitext2(self):
                                                      bptt_len=30)
 
         # Delete the dataset after we're done to save disk space on CI
-        if os.environ.get("TRAVIS") == "true":
-            datafile = os.path.join(self.project_root, ".data", "wikitext-2")
-            conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", "wikitext-2")
+        conditional_remove(datafile)
+
+    def test_wikitext2(self):
+        from torchtext.datasets import WikiText2
+        # smoke test to ensure wikitext2 works properly
+        train_dataset, test_dataset, valid_dataset = WikiText2()
+        self.assertEqual(len(train_dataset), 2049990)
+        self.assertEqual(len(test_dataset), 241859)
+        self.assertEqual(len(valid_dataset), 214417)
+
+        vocab = train_dataset.get_vocab()
+        tokens_ids = [vocab[token] for token in 'the player characters rest'.split()]
+        self.assertEqual(tokens_ids, [2, 286, 503, 700])
+
+        # Delete the dataset after we're done to save disk space on CI
+        datafile = os.path.join(self.project_root, ".data", "wikitext-2")
+        conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", "wikitext-2-v1.zip")
+        conditional_remove(datafile)
 
     @slow
-    def test_penntreebank(self):
+    def test_penntreebank_legacy(self):
+        from torchtext.legacy.datasets import PennTreebank
         # smoke test to ensure penn treebank works properly
         TEXT = data.Field(lower=True, batch_first=True)
         ds = PennTreebank
@@ -45,9 +66,28 @@ def test_penntreebank(self):
                                                      bptt_len=30)
 
         # Delete the dataset after we're done to save disk space on CI
-        if os.environ.get("TRAVIS") == "true":
-            datafile = os.path.join(self.project_root, ".data", "penn-treebank")
-            conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", "penn-treebank")
+        conditional_remove(datafile)
+
+    def test_penntreebank(self):
+        from torchtext.datasets import PennTreebank
+        # smoke test to ensure wikitext2 works properly
+        train_dataset, test_dataset, valid_dataset = PennTreebank()
+        self.assertEqual(len(train_dataset), 924412)
+        self.assertEqual(len(test_dataset), 82114)
+        self.assertEqual(len(valid_dataset), 73339)
+
+        vocab = train_dataset.get_vocab()
+        tokens_ids = [vocab[token] for token in 'the player characters rest'.split()]
+        self.assertEqual(tokens_ids, [2, 2550, 3344, 1125])
+
+        # Delete the dataset after we're done to save disk space on CI
+        datafile = os.path.join(self.project_root, ".data", 'ptb.train.txt')
+        conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", 'ptb.test.txt')
+        conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", 'ptb.valid.txt')
+        conditional_remove(datafile)
 
     def test_text_classification(self):
         # smoke test to ensure ag_news dataset works properly
@@ -60,6 +100,7 @@ def test_text_classification(self):
         self.assertEqual(len(ag_news_test), 7600)
 
         # Delete the dataset after we're done to save disk space on CI
-        if os.environ.get("TRAVIS") == "true":
-            datafile = os.path.join(self.project_root, ".data", "AG_NEWS")
-            conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", "ag_news_csv")
+        conditional_remove(datafile)
+        datafile = os.path.join(self.project_root, ".data", "ag_news_csv.tar.gz")
+        conditional_remove(datafile)

torchtext/__init__.py

@@ -2,10 +2,12 @@
 from . import datasets
 from . import utils
 from . import vocab
+from . import legacy
 
 __version__ = '0.4.0'
 
 __all__ = ['data',
            'datasets',
            'utils',
-           'vocab']
+           'vocab',
+           'legacy']

torchtext/data/__init__.py

@@ -10,7 +10,8 @@
 from .functional import generate_sp_model, \
     load_sp_model, \
     sentencepiece_numericalizer, \
-    sentencepiece_tokenizer, custom_replace, simple_space_split
+    sentencepiece_tokenizer, custom_replace, simple_space_split, \
+    numericalize_tokens_from_iterator
 
 __all__ = ["Batch",
            "Dataset", "TabularDataset",
@@ -24,4 +25,5 @@
            "get_tokenizer", "interleave_keys",
            "generate_sp_model", "load_sp_model",
            "sentencepiece_numericalizer", "sentencepiece_tokenizer",
-           "custom_replace", "simple_space_split"]
+           "custom_replace", "simple_space_split",
+           "numericalize_tokens_from_iterator"]

torchtext/data/functional.py

@@ -1,7 +1,6 @@
 import sentencepiece as spm
 import re
 
-
 __all__ = [
     "generate_sp_model", "load_sp_model",
     "sentencepiece_numericalizer", "sentencepiece_tokenizer"
@@ -151,3 +150,31 @@ def simple_space_split(iterator):
 
     for line in iterator:
         yield line.split()
+
+
+def numericalize_tokens_from_iterator(vocab, iterator, removed_tokens=None):
+    r"""Yield a list of ids from an token iterator with a vocab.
+
+    Arguments:
+        vocab: the vocabulary convert token into id.
+        iterator: the iterator yield a list of tokens.
+        removed_tokens: removed tokens from output dataset (Default: None)
+
+    Examples:
+        >>> from torchtext.data.functional import simple_space_split
+        >>> from torchtext.data.functional import numericalize_tokens_from_iterator
+        >>> vocab = {'Sentencepiece' : 0, 'encode' : 1, 'as' : 2, 'pieces' : 3}
+        >>> ids_iter = numericalize_tokens_from_iterator(vocab,
+        >>>                               simple_space_split(["Sentencepiece as pieces",
+        >>>                                                   "as pieces"]))
+        >>> for ids in ids_iter:
+        >>>     print([num for num in ids])
+        >>> [0, 2, 3]
+        >>> [2, 3]
+    """
+    for tokens in iterator:
+        if removed_tokens is None:
+            yield iter(vocab[token] for token in tokens)
+        else:
+            yield iter(map(lambda x: vocab[x],
+                       filter(lambda x: x not in removed_tokens, tokens)))