Port WMT en-de tokens 8k/32k to new Problem registry

PiperOrigin-RevId: 162025600

Author: Ryan Sepassi

tensor2tensor/bin/t2t-datagen

@@ -134,14 +134,6 @@ _SUPPORTED_PROBLEM_GENERATORS = {
     "wmt_ende_bpe32k": (
         lambda: wmt.ende_bpe_token_generator(FLAGS.tmp_dir, True),
         lambda: wmt.ende_bpe_token_generator(FLAGS.tmp_dir, False)),
-    "wmt_ende_tokens_8k": (
-        lambda: wmt.ende_wordpiece_token_generator(FLAGS.tmp_dir, True, 2**13),
-        lambda: wmt.ende_wordpiece_token_generator(FLAGS.tmp_dir, False, 2**13)
-    ),
-    "wmt_ende_tokens_32k": (
-        lambda: wmt.ende_wordpiece_token_generator(FLAGS.tmp_dir, True, 2**15),
-        lambda: wmt.ende_wordpiece_token_generator(FLAGS.tmp_dir, False, 2**15)
-    ),
     "wmt_zhen_tokens_32k": (
         lambda: wmt.zhen_wordpiece_token_generator(FLAGS.tmp_dir, True,
                                                    2**15, 2**15),

tensor2tensor/data_generators/algorithmic.py

@@ -37,12 +37,12 @@ def num_symbols(self):
     return 2
 
   def generate_data(self, data_dir, _):
-    utils.generate_files(
+    utils.generate_dataset_and_shuffle(
         identity_generator(self.num_symbols, 40, 100000),
-        self.training_filepaths(data_dir, 100))
-    utils.generate_files(
+        self.training_filepaths(data_dir, 100, shuffled=True),
         identity_generator(self.num_symbols, 400, 10000),
-        self.dev_filepaths(data_dir, 1))
+        self.dev_filepaths(data_dir, 1, shuffled=True),
+        shuffle=False)
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults

tensor2tensor/data_generators/generator_utils.py

@@ -359,6 +359,17 @@ def write_records(records, out_filename):
   writer.close()
 
 
+def generate_dataset_and_shuffle(train_gen,
+                                 train_paths,
+                                 dev_gen,
+                                 dev_paths,
+                                 shuffle=True):
+  generate_files(train_gen, train_paths)
+  generate_files(dev_gen, dev_paths)
+  if shuffle:
+    shuffle_dataset(train_paths + dev_paths)
+
+
 def shuffle_dataset(filenames):
   tf.logging.info("Shuffling data...")
   for fname in filenames:

tensor2tensor/data_generators/problem.py

@@ -78,12 +78,18 @@ class Problem(object):
   New problems are specified by the following methods:
 
   Data generation:
-    * generate_data(data_dir)
+    * generate_data(data_dir, tmp_dir)
         - Generate training and dev datasets into data_dir.
         - Additonal files, e.g. vocabulary files, should also be written to
           data_dir.
+        - Downloads and other files can be written to tmp_dir
+        - If you have a training and dev generator, you can generate the
+          training and dev datasets with
+          generator_utils.generate_dataset_and_shuffle.
         - Use the self.training_filepaths and self.dev_filepaths functions to
-          get sharded filenames.
+          get sharded filenames. If shuffled=False, the filenames will contain
+          an "unshuffled" suffix; you should then shuffle the data
+          shard-by-shard with generator_utils.shuffle_dataset.
         - Subclasses must override
     * dataset_filename()
         - Base filename for problem.
@@ -125,13 +131,17 @@ def feature_encoders(self, data_dir):
   # END SUBCLASS INTERFACE
   # ============================================================================
 
-  def training_filepaths(self, data_dir, num_shards):
-    return utils.train_data_filenames(self.dataset_filename(), data_dir,
-                                      num_shards)
-
-  def dev_filepaths(self, data_dir, num_shards):
-    return utils.dev_data_filenames(self.dataset_filename(), data_dir,
-                                    num_shards)
+  def training_filepaths(self, data_dir, num_shards, shuffled):
+    file_basename = self.dataset_filename()
+    if not shuffled:
+      file_basename += utils.UNSHUFFLED_SUFFIX
+    return utils.train_data_filenames(file_basename, data_dir, num_shards)
+
+  def dev_filepaths(self, data_dir, num_shards, shuffled):
+    file_basename = self.dataset_filename()
+    if not shuffled:
+      file_basename += utils.UNSHUFFLED_SUFFIX
+    return utils.dev_data_filenames(file_basename, data_dir, num_shards)
 
   def __init__(self, was_reversed=False, was_copy=False):
     """Create a Problem.

tensor2tensor/data_generators/problem_hparams.py

@@ -456,26 +456,6 @@ def wmt_ende_characters(unused_model_hparams):
   return p
 
 
-def wmt_ende_tokens(model_hparams, wrong_vocab_size):
-  """English to German translation benchmark."""
-  p = default_problem_hparams()
-  # This vocab file must be present within the data directory.
-  vocab_filename = os.path.join(model_hparams.data_dir,
-                                "tokens.vocab.%d" % wrong_vocab_size)
-  subtokenizer = text_encoder.SubwordTextEncoder(vocab_filename)
-  p.input_modality = {
-      "inputs": (registry.Modalities.SYMBOL, subtokenizer.vocab_size)
-  }
-  p.target_modality = (registry.Modalities.SYMBOL, subtokenizer.vocab_size)
-  p.vocabulary = {
-      "inputs": subtokenizer,
-      "targets": subtokenizer,
-  }
-  p.input_space_id = 3
-  p.target_space_id = 8
-  return p
-
-
 def wmt_zhen_tokens(model_hparams, wrong_vocab_size):
   """Chinese to English translation benchmark."""
   p = default_problem_hparams()
@@ -751,9 +731,6 @@ def img2img_imagenet(unused_model_hparams):
     "wmt_enfr_tokens_32k_combined": lambda p: wmt_enfr_tokens(p, 2**15),
     "wmt_enfr_tokens_128k": lambda p: wmt_enfr_tokens(p, 2**17),
     "wmt_ende_characters": wmt_ende_characters,
-    "wmt_ende_tokens_8k": lambda p: wmt_ende_tokens(p, 2**13),
-    "wmt_ende_tokens_32k": lambda p: wmt_ende_tokens(p, 2**15),
-    "wmt_ende_tokens_128k": lambda p: wmt_ende_tokens(p, 2**17),
     "wmt_ende_bpe32k": wmt_ende_bpe32k,
     "wmt_zhen_tokens_32k": lambda p: wmt_zhen_tokens(p, 2**15),
     "image_cifar10_tune": image_cifar10,

tensor2tensor/data_generators/wmt.py

@@ -24,20 +24,64 @@
 # Dependency imports
 
 from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import problem
 from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import wsj_parsing
+from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
-
 tf.flags.DEFINE_string("ende_bpe_path", "", "Path to BPE files in tmp_dir."
                        "Download from https://drive.google.com/open?"
                        "id=0B_bZck-ksdkpM25jRUN2X2UxMm8")
 
-
 FLAGS = tf.flags.FLAGS
 
 
+@registry.register_problem("wmt_ende_tokens_8k")
+class WMTEnDeTokens8k(problem.Problem):
+  """Problem spec for WMT En-De translation."""
+
+  @property
+  def target_vocab_size(self):
+    return 2**13  # 8192
+
+  def feature_encoders(self, data_dir):
+    return _default_wmt_feature_encoders(data_dir, self.target_vocab_size)
+
+  def generate_data(self, data_dir, tmp_dir):
+    generator_utils.generate_dataset_and_shuffle(
+        ende_wordpiece_token_generator(tmp_dir, True, self.target_vocab_size),
+        self.training_filepaths(data_dir, 100, shuffled=False),
+        ende_wordpiece_token_generator(tmp_dir, False, self.target_vocab_size),
+        self.dev_filepaths(data_dir, 1, shuffled=False))
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    vocab_size = self._encoders["inputs"].vocab_size
+    p.input_modality = {"inputs": (registry.Modalities.SYMBOL, vocab_size)}
+    p.target_modality = (registry.Modalities.SYMBOL, vocab_size)
+    p.input_space_id = problem.SpaceID.EN_TOK
+    p.target_space_id = problem.SpaceID.DE_TOK
+
+
+@registry.register_problem("wmt_ende_tokens_32k")
+class WMTEnDeTokens32k(WMTEnDeTokens8k):
+
+  @property
+  def target_vocab_size(self):
+    return 2**15  # 32768
+
+
+def _default_wmt_feature_encoders(data_dir, target_vocab_size):
+  vocab_filename = os.path.join(data_dir, "tokens.vocab.%d" % target_vocab_size)
+  subtokenizer = text_encoder.SubwordTextEncoder(vocab_filename)
+  return {
+      "inputs": subtokenizer,
+      "targets": subtokenizer,
+  }
+
+
 # End-of-sentence marker.
 EOS = text_encoder.EOS_TOKEN
 
@@ -130,7 +174,8 @@ def token_generator(source_path, target_path, token_vocab, eos=None):
         source, target = source_file.readline(), target_file.readline()
 
 
-def bi_vocabs_token_generator(source_path, target_path,
+def bi_vocabs_token_generator(source_path,
+                              target_path,
                               source_token_vocab,
                               target_token_vocab,
                               eos=None):
@@ -184,8 +229,8 @@ def ende_bpe_token_generator(tmp_dir, train):
   train_path = _get_wmt_ende_dataset(tmp_dir, dataset_path)
   token_path = os.path.join(tmp_dir, "vocab.bpe.32000")
   token_vocab = text_encoder.TokenTextEncoder(vocab_filename=token_path)
-  return token_generator(train_path + ".en", train_path + ".de",
-                         token_vocab, EOS)
+  return token_generator(train_path + ".en", train_path + ".de", token_vocab,
+                         EOS)
 
 
 _ENDE_TRAIN_DATASETS = [
@@ -240,22 +285,15 @@ def ende_bpe_token_generator(tmp_dir, train):
     ],
 ]
 
-_ZHEN_TRAIN_DATASETS = [
-    [
-        ("http://data.statmt.org/wmt17/translation-task/"
-         "training-parallel-nc-v12.tgz"),
-        ("training/news-commentary-v12.zh-en.zh",
-         "training/news-commentary-v12.zh-en.en")
-    ]
-]
+_ZHEN_TRAIN_DATASETS = [[("http://data.statmt.org/wmt17/translation-task/"
+                          "training-parallel-nc-v12.tgz"),
+                         ("training/news-commentary-v12.zh-en.zh",
+                          "training/news-commentary-v12.zh-en.en")]]
 
-_ZHEN_TEST_DATASETS = [
-    [
-        "http://data.statmt.org/wmt17/translation-task/dev.tgz",
-        ("dev/newsdev2017-zhen-src.zh",
-         "dev/newsdev2017-zhen-ref.en")
-    ]
-]
+_ZHEN_TEST_DATASETS = [[
+    "http://data.statmt.org/wmt17/translation-task/dev.tgz",
+    ("dev/newsdev2017-zhen-src.zh", "dev/newsdev2017-zhen-ref.en")
+]]
 
 
 def _compile_data(tmp_dir, datasets, filename):
@@ -317,23 +355,21 @@ def ende_character_generator(tmp_dir, train):
                              character_vocab, EOS)
 
 
-def zhen_wordpiece_token_generator(tmp_dir, train,
-                                   source_vocab_size,
+def zhen_wordpiece_token_generator(tmp_dir, train, source_vocab_size,
                                    target_vocab_size):
   """Wordpiece generator for the WMT'17 zh-en dataset."""
   datasets = _ZHEN_TRAIN_DATASETS if train else _ZHEN_TEST_DATASETS
   source_datasets = [[item[0], [item[1][0]]] for item in datasets]
   target_datasets = [[item[0], [item[1][1]]] for item in datasets]
   source_vocab = generator_utils.get_or_generate_vocab(
-      tmp_dir, "tokens.vocab.zh.%d" % source_vocab_size,
-      source_vocab_size, source_datasets)
+      tmp_dir, "tokens.vocab.zh.%d" % source_vocab_size, source_vocab_size,
+      source_datasets)
   target_vocab = generator_utils.get_or_generate_vocab(
-      tmp_dir, "tokens.vocab.en.%d" % target_vocab_size,
-      target_vocab_size, target_datasets)
+      tmp_dir, "tokens.vocab.en.%d" % target_vocab_size, target_vocab_size,
+      target_datasets)
   tag = "train" if train else "dev"
   data_path = _compile_data(tmp_dir, datasets, "wmt_zhen_tok_%s" % tag)
-  return bi_vocabs_token_generator(data_path + ".lang1",
-                                   data_path + ".lang2",
+  return bi_vocabs_token_generator(data_path + ".lang1", data_path + ".lang2",
                                    source_vocab, target_vocab, EOS)
 
 
@@ -366,17 +402,15 @@ def parsing_character_generator(tmp_dir, train):
   return character_generator(text_filepath, tags_filepath, character_vocab, EOS)
 
 
-def tabbed_parsing_token_generator(tmp_dir, train, prefix,
-                                   source_vocab_size, target_vocab_size):
+def tabbed_parsing_token_generator(tmp_dir, train, prefix, source_vocab_size,
+                                   target_vocab_size):
   """Generate source and target data from a single file."""
   source_vocab = generator_utils.get_or_generate_tabbed_vocab(
       tmp_dir, "parsing_train.pairs", 0,
-      prefix + "_source.tokens.vocab.%d" % source_vocab_size,
-      source_vocab_size)
+      prefix + "_source.tokens.vocab.%d" % source_vocab_size, source_vocab_size)
   target_vocab = generator_utils.get_or_generate_tabbed_vocab(
       tmp_dir, "parsing_train.pairs", 1,
-      prefix + "_target.tokens.vocab.%d" % target_vocab_size,
-      target_vocab_size)
+      prefix + "_target.tokens.vocab.%d" % target_vocab_size, target_vocab_size)
   filename = "parsing_%s" % ("train" if train else "dev")
   pair_filepath = os.path.join(tmp_dir, filename + ".pairs")
   return tabbed_generator(pair_filepath, source_vocab, target_vocab, EOS)
@@ -395,5 +429,5 @@ def parsing_token_generator(tmp_dir, train, vocab_size):
       tmp_dir, "tokens.vocab.%d" % vocab_size, vocab_size)
   filename = "%s_%s.trees" % (FLAGS.parsing_path, "train" if train else "dev")
   tree_filepath = os.path.join(tmp_dir, filename)
-  return wsj_parsing.token_generator(tree_filepath,
-                                     symbolizer_vocab, symbolizer_vocab, EOS)
+  return wsj_parsing.token_generator(tree_filepath, symbolizer_vocab,
+                                     symbolizer_vocab, EOS)

tensor2tensor/utils/trainer_utils_test.py

@@ -34,13 +34,13 @@
 @registry.register_problem
 class TinyAlgo(algorithmic.AlgorithmicIdentityBinary40):
 
-  def generate_data(self, data_dir):
+  def generate_data(self, data_dir, _):
     generator_utils.generate_files(
         algorithmic.identity_generator(self.num_symbols, 40, 100000),
-        self.training_filepaths(data_dir, 1), 100)
+        self.training_filepaths(data_dir, 1, shuffled=True), 100)
     generator_utils.generate_files(
         algorithmic.identity_generator(self.num_symbols, 400, 10000),
-        self.dev_filepaths(data_dir, 1), 100)
+        self.dev_filepaths(data_dir, 1, shuffled=True), 100)
 
 
 @registry.register_hparams
@@ -61,7 +61,8 @@ def setUpClass(cls):
     # Generate a small test dataset
     FLAGS.problems = "tiny_algo"
     TrainerUtilsTest.data_dir = tf.test.get_temp_dir()
-    registry.problem(FLAGS.problems).generate_data(TrainerUtilsTest.data_dir)
+    registry.problem(FLAGS.problems).generate_data(TrainerUtilsTest.data_dir,
+                                                   None)
 
   def testModelsImported(self):
     models = registry.list_models()