Remove offset option from torchtext.datasets, move torchtext.datasets.common to torchtext.data.dataset_utils (#1188)

Author: cpuhrsch

test/data/test_builtin_datasets.py

@@ -186,11 +186,6 @@ def test_raw_text_classification(self, info):
             self.assertEqual(torchtext.datasets.MD5[dataset_name], info['MD5'])
         del data_iter
 
-    def test_num_lines_of_dataset(self):
-        train_iter, test_iter = torchtext.datasets.AG_NEWS(offset=10)
-        _data = [item for item in train_iter]
-        self.assertEqual(len(_data), 119990)
-
     @parameterized.expand(list(sorted(torchtext.datasets.DATASETS.keys())))
     def test_raw_datasets_split_argument(self, dataset_name):
         if dataset_name in GOOGLE_DRIVE_BASED_DATASETS:
@@ -223,13 +218,6 @@ def test_datasets_split_argument(self, dataset_name):
         # Exercise default constructor
         _ = dataset()
 
-    def test_offset_dataset(self):
-        train_iter, test_iter = torchtext.datasets.AG_NEWS(split=('train', 'test'), offset=10)
-        container = [text[:20] for idx, (label, text) in enumerate(train_iter) if idx < 5]
-        self.assertEqual(container, ['Oil and Economy Clou', 'No Need for OPEC to ',
-                                     'Non-OPEC Nations Sho', 'Google IPO Auction O',
-                                     'Dollar Falls Broadly'])
-
     def test_next_method_dataset(self):
         train_iter, test_iter = torchtext.datasets.AG_NEWS()
         for_count = 0

torchtext/datasets/common.py renamed to torchtext/data/datasets_utils.py

@@ -44,12 +44,11 @@ def dataset_docstring_header(fn):
     """
     Returns docstring for a dataset based on function arguments.
 
-    Assumes function signature of form (root='.data', split=<some tuple of strings>, offset=0, **kwargs)
+    Assumes function signature of form (root='.data', split=<some tuple of strings>, **kwargs)
     """
     argspec = inspect.getfullargspec(fn)
     if not (argspec.args[0] == "root" and
-            argspec.args[1] == "split" and
-            argspec.args[2] == "offset"):
+            argspec.args[1] == "split"):
         raise ValueError("Internal Error: Given function {} did not adhere to standard signature.".format(fn))
     default_split = argspec.defaults[1]
 
@@ -68,8 +67,6 @@ def dataset_docstring_header(fn):
                 By default, all three datasets are generated. Users
                 could also choose any subset of them, for example {} or just 'train'.
                 Default: {}
-            offset: the number of the starting line.
-                Default: 0
         """.format(fn.__name__, "/".join(default_split), str(example_subset), str(default_split))
 
     if isinstance(default_split, str):
@@ -81,9 +78,7 @@ def dataset_docstring_header(fn):
             root: Directory where the datasets are saved.
                 Default: ".data"
             split: Only {default_split} is available.
-                Default: {default_split}
-            offset: the number of the starting line.
-                Default: 0""".format(fn.__name__, default_split=default_split)
+                Default: {default_split}""".format(fn.__name__, default_split=default_split)
 
     raise ValueError("default_split type expected to be of string or tuple but got {}".format(type(default_split)))
 
@@ -116,9 +111,7 @@ def wrap_split_argument(fn):
     argspec = inspect.getfullargspec(fn)
     if not (argspec.args[0] == "root" and
             argspec.args[1] == "split" and
-            argspec.args[2] == "offset" and
             argspec.defaults[0] == ".data" and
-            argspec.defaults[2] == 0 and
             argspec.varargs is None and
             argspec.varkw is None and
             len(argspec.kwonlyargs) == 0 and
@@ -133,16 +126,15 @@ def wrap_split_argument(fn):
     # keyword arguments with default values only, so only  a dictionary of default
     # values is needed to support that behavior for new_fn as well.
     fn_kwargs_dict = {}
-    for arg, default in zip(argspec.args[3:], argspec.defaults[3:]):
+    for arg, default in zip(argspec.args[2:], argspec.defaults[2:]):
         fn_kwargs_dict[arg] = default
 
     @functools.wraps(fn)
-    def new_fn(root='.data', split=argspec.defaults[1], offset=0, **kwargs):
+    def new_fn(root='.data', split=argspec.defaults[1], **kwargs):
         for arg in fn_kwargs_dict:
             if arg not in kwargs:
                 kwargs[arg] = fn_kwargs_dict[arg]
         kwargs["root"] = root
-        kwargs["offset"] = offset
         kwargs["split"] = check_default_set(split, argspec.defaults[1], fn.__name__)
         result = fn(**kwargs)
         return wrap_datasets(tuple(result), split)
@@ -154,32 +146,28 @@ class RawTextIterableDataset(torch.utils.data.IterableDataset):
     """Defines an abstraction for raw text iterable datasets.
     """
 
-    def __init__(self, name, full_num_lines, iterator, offset=0):
+    def __init__(self, name, full_num_lines, iterator):
         """Initiate text-classification dataset.
         """
         super(RawTextIterableDataset, self).__init__()
         self.name = name
         self.full_num_lines = full_num_lines
         self._iterator = iterator
-        self.start = offset
-        if offset < 0:
-            raise ValueError("Given offset must be non-negative, got {} instead.".format(offset))
-        self.num_lines = full_num_lines - offset
+        self.num_lines = full_num_lines
+        self.current_pos = None
 
     def __iter__(self):
-        for i, item in enumerate(self._iterator):
-            if i < self.start:
-                continue
-            if self.num_lines and i >= (self.start + self.num_lines):
-                break
-            yield item
+        return self
 
     def __next__(self):
+        if self.current_pos == self.num_lines - 1:
+            raise StopIteration
         item = next(self._iterator)
+        if self.current_pos is None:
+            self.current_pos = 0
+        else:
+            self.current_pos += 1
         return item
 
     def __len__(self):
         return self.num_lines
-
-    def get_iterator(self):
-        return self._iterator

torchtext/datasets/ag_news.py

@@ -1,7 +1,7 @@
 from torchtext.utils import download_from_url, unicode_csv_reader
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
 import os
 import io
 
@@ -23,7 +23,7 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def AG_NEWS(root='.data', split=('train', 'test'), offset=0):
+def AG_NEWS(root='.data', split=('train', 'test')):
     def _create_data_from_csv(data_path):
         with io.open(data_path, encoding="utf8") as f:
             reader = unicode_csv_reader(f)
@@ -37,5 +37,5 @@ def _create_data_from_csv(data_path):
                                  hash_value=MD5[item],
                                  hash_type='md5')
         datasets.append(RawTextIterableDataset("AG_NEWS", NUM_LINES[item],
-                                               _create_data_from_csv(path), offset=offset))
+                                               _create_data_from_csv(path)))
     return datasets

torchtext/datasets/amazonreviewfull.py

@@ -1,8 +1,8 @@
 from torchtext.utils import download_from_url, extract_archive, unicode_csv_reader
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
-from torchtext.datasets.common import find_match
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
+from torchtext.data.datasets_utils import find_match
 import os
 import io
 import logging
@@ -21,7 +21,7 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def AmazonReviewFull(root='.data', split=('train', 'test'), offset=0):
+def AmazonReviewFull(root='.data', split=('train', 'test')):
     def _create_data_from_csv(data_path):
         with io.open(data_path, encoding="utf8") as f:
             reader = unicode_csv_reader(f)
@@ -37,5 +37,5 @@ def _create_data_from_csv(data_path):
         path = find_match(item + '.csv', extracted_files)
         logging.info('Creating {} data'.format(item))
         datasets.append(RawTextIterableDataset("AmazonReviewFull", NUM_LINES[item],
-                                               _create_data_from_csv(path), offset=offset))
+                                               _create_data_from_csv(path)))
     return datasets

torchtext/datasets/amazonreviewpolarity.py

@@ -1,8 +1,8 @@
 from torchtext.utils import download_from_url, extract_archive, unicode_csv_reader
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
-from torchtext.datasets.common import find_match
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
+from torchtext.data.datasets_utils import find_match
 import os
 import io
 
@@ -20,7 +20,7 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def AmazonReviewPolarity(root='.data', split=('train', 'test'), offset=0):
+def AmazonReviewPolarity(root='.data', split=('train', 'test')):
     def _create_data_from_csv(data_path):
         with io.open(data_path, encoding="utf8") as f:
             reader = unicode_csv_reader(f)
@@ -35,5 +35,5 @@ def _create_data_from_csv(data_path):
     for item in split:
         path = find_match(item + '.csv', extracted_files)
         datasets.append(RawTextIterableDataset("AmazonReviewPolarity", NUM_LINES[item],
-                                               _create_data_from_csv(path), offset=offset))
+                                               _create_data_from_csv(path)))
     return datasets

torchtext/datasets/conll2000chunking.py

@@ -1,8 +1,8 @@
 from torchtext.utils import download_from_url, extract_archive
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
-from torchtext.datasets.common import find_match
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
+from torchtext.data.datasets_utils import find_match
 
 URL = {
     'train': "https://www.clips.uantwerpen.be/conll2000/chunking/train.txt.gz",
@@ -40,12 +40,12 @@ def _create_data_from_iob(data_path, separator):
 
 @wrap_split_argument
 @add_docstring_header()
-def CoNLL2000Chunking(root='.data', split=('train', 'test'), offset=0):
+def CoNLL2000Chunking(root='.data', split=('train', 'test')):
     datasets = []
     for item in split:
         dataset_tar = download_from_url(URL[item], root=root, hash_value=MD5[item], hash_type='md5')
         extracted_files = extract_archive(dataset_tar)
         data_filename = find_match(item + ".txt", extracted_files)
         datasets.append(RawTextIterableDataset("CoNLL2000Chunking", NUM_LINES[item],
-                                               _create_data_from_iob(data_filename, " "), offset=offset))
+                                               _create_data_from_iob(data_filename, " ")))
     return datasets

torchtext/datasets/dbpedia.py

@@ -1,8 +1,8 @@
 from torchtext.utils import download_from_url, extract_archive, unicode_csv_reader
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
-from torchtext.datasets.common import find_match
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
+from torchtext.data.datasets_utils import find_match
 import os
 import io
 
@@ -20,7 +20,7 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def DBpedia(root='.data', split=('train', 'test'), offset=0):
+def DBpedia(root='.data', split=('train', 'test')):
     def _create_data_from_csv(data_path):
         with io.open(data_path, encoding="utf8") as f:
             reader = unicode_csv_reader(f)
@@ -35,5 +35,5 @@ def _create_data_from_csv(data_path):
     for item in split:
         path = find_match(item + '.csv', extracted_files)
         datasets.append(RawTextIterableDataset("DBpedia", NUM_LINES[item],
-                                               _create_data_from_csv(path), offset=offset))
+                                               _create_data_from_csv(path)))
     return datasets

torchtext/datasets/enwik9.py

@@ -1,8 +1,8 @@
 import logging
 from torchtext.utils import download_from_url, extract_archive
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
 import io
 
 URL = 'http://mattmahoney.net/dc/enwik9.zip'
@@ -16,10 +16,10 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def EnWik9(root='.data', split='train', offset=0):
+def EnWik9(root='.data', split='train'):
     dataset_tar = download_from_url(URL, root=root, hash_value=MD5, hash_type='md5')
     extracted_files = extract_archive(dataset_tar)
     path = extracted_files[0]
     logging.info('Creating {} data'.format(split[0]))
     return [RawTextIterableDataset('EnWik9',
-                                   NUM_LINES[split[0]], iter(io.open(path, encoding="utf8")), offset=offset)]
+                                   NUM_LINES[split[0]], iter(io.open(path, encoding="utf8")))]

torchtext/datasets/imdb.py

@@ -1,7 +1,7 @@
 from torchtext.utils import download_from_url, extract_archive
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
 import io
 
 URL = 'http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz'
@@ -18,7 +18,7 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def IMDB(root='.data', split=('train', 'test'), offset=0):
+def IMDB(root='.data', split=('train', 'test')):
     def generate_imdb_data(key, extracted_files):
         for fname in extracted_files:
             if 'urls' in fname:
@@ -33,5 +33,5 @@ def generate_imdb_data(key, extracted_files):
     datasets = []
     for item in split:
         iterator = generate_imdb_data(item, extracted_files)
-        datasets.append(RawTextIterableDataset("IMDB", NUM_LINES[item], iterator, offset=offset))
+        datasets.append(RawTextIterableDataset("IMDB", NUM_LINES[item], iterator))
     return datasets

torchtext/datasets/iwslt.py

@@ -3,9 +3,9 @@
 import codecs
 import xml.etree.ElementTree as ET
 from torchtext.utils import (download_from_url, extract_archive)
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
 
 URL = 'https://drive.google.com/uc?id=1l5y6Giag9aRPwGtuZHswh3w5v3qEz8D8'
 
@@ -278,6 +278,6 @@ def _iter(src_data_iter, tgt_data_iter):
                 yield item
 
         datasets.append(
-            RawTextIterableDataset("IWSLT", NUM_LINES[key], _iter(src_data_iter, tgt_data_iter), offset=offset))
+            RawTextIterableDataset("IWSLT", NUM_LINES[key], _iter(src_data_iter, tgt_data_iter)))
 
     return datasets

torchtext/datasets/multi30k.py

@@ -1,9 +1,9 @@
 import io
 import os
 from torchtext.utils import (download_from_url, extract_archive)
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
 
 _URL_BASE_ = 'https://raw.githubusercontent.com/multi30k/dataset/master/data/task'
 
@@ -171,6 +171,6 @@ def _iter(src_data_iter, tgt_data_iter):
                 yield item
 
         datasets.append(
-            RawTextIterableDataset("Multi30k", NUM_LINES[key], _iter(src_data_iter, tgt_data_iter), offset=offset))
+            RawTextIterableDataset("Multi30k", NUM_LINES[key], _iter(src_data_iter, tgt_data_iter)))
 
     return datasets

torchtext/datasets/penntreebank.py

@@ -1,8 +1,8 @@
 import logging
 from torchtext.utils import download_from_url
-from torchtext.datasets.common import RawTextIterableDataset
-from torchtext.datasets.common import wrap_split_argument
-from torchtext.datasets.common import add_docstring_header
+from torchtext.data.datasets_utils import RawTextIterableDataset
+from torchtext.data.datasets_utils import wrap_split_argument
+from torchtext.data.datasets_utils import add_docstring_header
 import io
 
 URL = {
@@ -26,7 +26,7 @@
 
 @wrap_split_argument
 @add_docstring_header()
-def PennTreebank(root='.data', split=('train', 'valid', 'test'), offset=0):
+def PennTreebank(root='.data', split=('train', 'valid', 'test')):
     datasets = []
     for item in split:
         path = download_from_url(URL[item],
@@ -35,6 +35,5 @@ def PennTreebank(root='.data', split=('train', 'valid', 'test'), offset=0):
         logging.info('Creating {} data'.format(item))
         datasets.append(RawTextIterableDataset('PennTreebank',
                                                NUM_LINES[item],
-                                               iter(io.open(path, encoding="utf8")),
-                                               offset=offset))
+                                               iter(io.open(path, encoding="utf8"))))
     return datasets