[BC-breaking] Standardize raw dataset doc strings and argument order. (#1151)

Author: cpuhrsch

torchtext/experimental/datasets/raw/common.py

@@ -1,4 +1,6 @@
 import torch
+import inspect
+import functools
 
 
 def check_default_set(split, target_select, dataset_name):
@@ -25,6 +27,109 @@ def wrap_datasets(datasets, split):
     return datasets
 
 
+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)
+    """
+    argspec = inspect.getfullargspec(fn)
+    if not (argspec.args[0] == "root" and
+            argspec.args[1] == "split" and
+            argspec.args[2] == "offset"):
+        raise ValueError("Internal Error: Given function {} did not adhere to standard signature.".format(fn))
+    default_split = argspec.defaults[1]
+
+    if isinstance(default_split, tuple):
+        example_subset = default_split[:2]
+        if len(default_split) < 3:
+            example_subset = (default_split[1],)
+        return """{} dataset
+
+        Separately returns the {} split
+
+        Args:
+            root: Directory where the datasets are saved.
+                Default: ".data"
+            split: split or splits to be returned. Can be a string or tuple of strings.
+                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)) + fn.__doc__
+
+    if isinstance(default_split, str):
+        return """{} dataset
+
+        Only returns the {default_split} split
+
+        Args:
+            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) + fn.__doc__
+
+    raise ValueError("default_split type expected to be of string or tuple but got {}".format(type(default_split)))
+
+
+def add_docstring_header(fn):
+    fn.__doc__ = dataset_docstring_header(fn)
+    return fn
+
+
+def wrap_split_argument(fn):
+    """
+    Wraps given function of specific signature to extend behavior of split
+    to support individual strings. The given function is expected to have a split
+    kwarg that accepts tuples of strings, e.g. ('train', 'valid') and the returned
+    function will have a split argument that also accepts strings, e.g. 'train', which
+    are then turned single entry tuples. Furthermore, the return value of the wrapped
+    function is unpacked if split is only a single string to enable behavior such as
+
+    train = AG_NEWS(split='train')
+    train, valid = AG_NEWS(split=('train', 'valid'))
+    """
+
+    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
+            argspec.kwonlydefaults is None and
+            len(argspec.annotations) == 0
+            ):
+        raise ValueError("Internal Error: Given function {} did not adhere to standard signature.".format(fn))
+
+    # functools.wraps only forwards __module__, __name__, etc
+    # (see https://docs.python.org/3/library/functools.html#functools.update_wrapper)
+    # but not default values of arguments. The wrapped function fn is assumed to have
+    # 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, argspec.defaults):
+        fn_kwargs_dict[arg] = default
+
+    @functools.wraps(fn)
+    def new_fn(**kwargs):
+        for arg in fn_kwargs_dict:
+            if arg not in kwargs:
+                kwargs[arg] = fn_kwargs_dict[arg]
+        orig_split = kwargs["split"]
+        kwargs["split"] = check_default_set(orig_split, argspec.defaults[1], fn.__name__)
+        result = fn(**kwargs)
+        return wrap_datasets(tuple(result), orig_split)
+
+    return new_fn
+
+
 class RawTextIterableDataset(torch.utils.data.IterableDataset):
     """Defines an abstraction for raw text iterable datasets.
     """

torchtext/experimental/datasets/raw/language_modeling.py

@@ -2,8 +2,8 @@
 import io
 from torchtext.utils import download_from_url, extract_archive
 from torchtext.experimental.datasets.raw.common import RawTextIterableDataset
-from torchtext.experimental.datasets.raw.common import check_default_set
-from torchtext.experimental.datasets.raw.common import wrap_datasets
+from torchtext.experimental.datasets.raw.common import wrap_split_argument
+from torchtext.experimental.datasets.raw.common import add_docstring_header
 
 URLS = {
     'WikiText2':
@@ -19,12 +19,7 @@
 }
 
 
-def _setup_datasets(dataset_name, root, split_, year, language, offset):
-    if dataset_name == 'WMTNewsCrawl':
-        split = check_default_set(split_, ('train',), dataset_name)
-    else:
-        split = check_default_set(split_, ('train', 'test', 'valid'), dataset_name)
-
+def _setup_datasets(dataset_name, root, split, year, language, offset):
     if dataset_name == 'PennTreebank':
         extracted_files = [download_from_url(URLS['PennTreebank'][key],
                                              root=root, hash_value=MD5['PennTreebank'][key],
@@ -49,23 +44,13 @@ def _setup_datasets(dataset_name, root, split_, year, language, offset):
         datasets.append(RawTextIterableDataset(dataset_name,
                                                NUM_LINES[dataset_name][item], iter(io.open(path[item], encoding="utf8")), offset=offset))
 
-    return wrap_datasets(tuple(datasets), split_)
+    return datasets
 
 
+@wrap_split_argument
+@add_docstring_header
 def WikiText2(root='.data', split=('train', 'valid', 'test'), offset=0):
-    """ Defines WikiText2 datasets.
-
-    Create language modeling dataset: WikiText2
-    Separately returns the train/test/valid set
-
-    Args:
-        root: Directory where the datasets are saved. Default: ".data"
-        split: a string or tuple for the returned datasets. Default: ('train', 'valid, 'test')
-            By default, all the three datasets (train, test, valid) are generated. Users
-            could also choose any one or two of them, for example ('train', 'test') or
-            just a string 'train'.
-        offset: the number of the starting line. Default: 0
-
+    """
     Examples:
         >>> from torchtext.experimental.raw.datasets import WikiText2
         >>> train_dataset, valid_dataset, test_dataset = WikiText2()
@@ -76,19 +61,10 @@ def WikiText2(root='.data', split=('train', 'valid', 'test'), offset=0):
     return _setup_datasets("WikiText2", root, split, None, None, offset)
 
 
+@wrap_split_argument
+@add_docstring_header
 def WikiText103(root='.data', split=('train', 'valid', 'test'), offset=0):
-    """ Defines WikiText103 datasets.
-
-    Create language modeling dataset: WikiText103
-    Separately returns the train/test/valid set
-
-    Args:
-        root: Directory where the datasets are saved. Default: ".data"
-        split: the returned datasets. Default: ('train', 'valid','test')
-            By default, all the three datasets (train, test, valid) are generated. Users
-            could also choose any one or two of them, for example ('train', 'test').
-        offset: the number of the starting line. Default: 0
-
+    """
     Examples:
         >>> from torchtext.experimental.datasets.raw import WikiText103
         >>> train_dataset, valid_dataset, test_dataset = WikiText103()
@@ -98,21 +74,10 @@ def WikiText103(root='.data', split=('train', 'valid', 'test'), offset=0):
     return _setup_datasets("WikiText103", root, split, None, None, offset)
 
 
+@wrap_split_argument
+@add_docstring_header
 def PennTreebank(root='.data', split=('train', 'valid', 'test'), offset=0):
-    """ Defines PennTreebank datasets.
-
-    Create language modeling dataset: PennTreebank
-    Separately returns the train/test/valid set
-
-    Args:
-        root: Directory where the datasets are saved. Default: ".data"
-        split: a string or tuple for the returned datasets
-            (Default: ('train', 'test','valid'))
-            By default, all the three datasets ('train', 'valid', 'test') are generated. Users
-            could also choose any one or two of them, for example ('train', 'test') or
-            just a string 'train'.
-        offset: the number of the starting line. Default: 0
-
+    """
     Examples:
         >>> from torchtext.experimental.datasets.raw import PennTreebank
         >>> train_dataset, valid_dataset, test_dataset = PennTreebank()
@@ -123,18 +88,11 @@ def PennTreebank(root='.data', split=('train', 'valid', 'test'), offset=0):
     return _setup_datasets("PennTreebank", root, split, None, None, offset)
 
 
-def WMTNewsCrawl(root='.data', split=('train'), year=2010, language='en', offset=0):
-    """ Defines WMT News Crawl.
-
-    Create language modeling dataset: WMTNewsCrawl
-
-    Args:
-        root: Directory where the datasets are saved. Default: ".data"
-        split: a string or tuple for the returned datasets.
-            (Default: 'train')
-        year: the year of the dataset (Default: 2010)
+@wrap_split_argument
+@add_docstring_header
+def WMTNewsCrawl(root='.data', split='train', offset=0, year=2010, language='en'):
+    """    year: the year of the dataset (Default: 2010)
         language: the language of the dataset (Default: 'en')
-        offset: the number of the starting line. Default: 0
 
     Note: WMTNewsCrawl provides datasets based on the year and language instead of train/valid/test.
     """
@@ -148,12 +106,14 @@ def WMTNewsCrawl(root='.data', split=('train'), year=2010, language='en', offset
     'PennTreebank': PennTreebank,
     'WMTNewsCrawl': WMTNewsCrawl
 }
+
 NUM_LINES = {
     'WikiText2': {'train': 36718, 'valid': 3760, 'test': 4358},
     'WikiText103': {'train': 1801350, 'valid': 3760, 'test': 4358},
     'PennTreebank': {'train': 42068, 'valid': 3370, 'test': 3761},
     'WMTNewsCrawl': {'train': 17676013}
 }
+
 MD5 = {
     'WikiText2': '542ccefacc6c27f945fb54453812b3cd',
     'WikiText103': '9ddaacaf6af0710eda8c456decff7832',

torchtext/experimental/datasets/raw/question_answer.py

@@ -1,8 +1,8 @@
 from torchtext.utils import download_from_url
 import json
 from torchtext.experimental.datasets.raw.common import RawTextIterableDataset
-from torchtext.experimental.datasets.raw.common import check_default_set
-from torchtext.experimental.datasets.raw.common import wrap_datasets
+from torchtext.experimental.datasets.raw.common import wrap_split_argument
+from torchtext.experimental.datasets.raw.common import add_docstring_header
 
 URLS = {
     'SQuAD1':
@@ -30,59 +30,48 @@ def _create_data_from_json(data_path):
                     yield (_context, _question, _answers, _answer_start)
 
 
-def _setup_datasets(dataset_name, root, split_, offset):
-    split = check_default_set(split_, ('train', 'dev'), dataset_name)
+def _setup_datasets(dataset_name, root, split, offset):
     extracted_files = {key: download_from_url(URLS[dataset_name][key], root=root,
                                               hash_value=MD5[dataset_name][key], hash_type='md5') for key in split}
-    return wrap_datasets(tuple(RawTextIterableDataset(dataset_name, NUM_LINES[dataset_name][item],
-                                                      _create_data_from_json(extracted_files[item]), offset=offset) for item in split), split_)
+    return [RawTextIterableDataset(dataset_name, NUM_LINES[dataset_name][item],
+                                   _create_data_from_json(extracted_files[item]), offset=offset) for item in split]
 
 
+@wrap_split_argument
+@add_docstring_header
 def SQuAD1(root='.data', split=('train', 'dev'), offset=0):
-    """ A dataset iterator yields the data of Stanford Question Answering dataset - SQuAD1.0.
+    """
+    Examples:
+        >>> train_dataset, dev_dataset = torchtext.experimental.datasets.raw.SQuAD1()
+        >>> for idx, (context, question, answer, ans_pos) in enumerate(train_dataset):
+        >>>     print(idx, (context, question, answer, ans_pos))
+
     The iterator yields a tuple of (raw context, raw question, a list of raw answer,
     a list of answer positions in the raw context).
     For example, ('Architecturally, the school has a Catholic character. Atop the ...',
                   'To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?',
                   ['Saint Bernadette Soubirous'],
                   [515])
-
-    Args:
-        root: Directory where the datasets are saved. Default: ".data"
-        split: a string or tuple for the returned datasets (Default: ('train', 'dev'))
-            By default, both datasets (train, dev) are generated. Users could also choose any one or two of them,
-            for example ('train', 'dev') or just a string 'train'.
-        offset: the number of the starting line. Default: 0
-
-    Examples:
-        >>> train_dataset, dev_dataset = torchtext.experimental.datasets.raw.SQuAD1()
-        >>> for idx, (context, question, answer, ans_pos) in enumerate(train_dataset):
-        >>>     print(idx, (context, question, answer, ans_pos))
     """
 
     return _setup_datasets("SQuAD1", root, split, offset)
 
 
+@wrap_split_argument
+@add_docstring_header
 def SQuAD2(root='.data', split=('train', 'dev'), offset=0):
-    """ A dataset iterator yields the data of Stanford Question Answering dataset - SQuAD2.0.
+    """
+    Examples:
+        >>> train_dataset, dev_dataset = torchtext.experimental.datasets.raw.SQuAD2()
+        >>> for idx, (context, question, answer, ans_pos) in enumerate(train_dataset):
+        >>>     print(idx, (context, question, answer, ans_pos))
+
     The iterator yields a tuple of (raw context, raw question, a list of raw answer,
     a list of answer positions in the raw context).
     For example, ('Beyoncé Giselle Knowles-Carter (/biːˈjɒnseɪ/ bee-YON-say) (born September 4, 1981) is an ...',
                   'When did Beyonce start becoming popular?',
                   ['in the late 1990s'],
                   [269])
-
-    Args:
-        root: Directory where the datasets are saved. Default: ".data"
-        split: a string or tuple for the returned datasets (Default: ('train', 'dev'))
-            By default, both datasets (train, dev) are generated. Users could also choose any one or two of them,
-            for example ('train', 'dev') or just a string 'train'.
-        offset: the number of the starting line. Default: 0
-
-    Examples:
-        >>> train_dataset, dev_dataset = torchtext.experimental.datasets.raw.SQuAD2()
-        >>> for idx, (context, question, answer, ans_pos) in enumerate(train_dataset):
-        >>>     print(idx, (context, question, answer, ans_pos))
     """
 
     return _setup_datasets("SQuAD2", root, split, offset)
@@ -92,10 +81,12 @@ def SQuAD2(root='.data', split=('train', 'dev'), offset=0):
     'SQuAD1': SQuAD1,
     'SQuAD2': SQuAD2
 }
+
 NUM_LINES = {
     'SQuAD1': {'train': 87599, 'dev': 10570},
     'SQuAD2': {'train': 130319, 'dev': 11873}
 }
+
 MD5 = {
     'SQuAD1': {'train': '981b29407e0affa3b1b156f72073b945', 'dev': '3e85deb501d4e538b6bc56f786231552'},
     'SQuAD2': {'train': '62108c273c268d70893182d5cf8df740', 'dev': '246adae8b7002f8679c027697b0b7cf8'}