Move algorithmic and WMT problems to Problem class, correct summaries.

PiperOrigin-RevId: 162424062

Author: Lukasz Kaiser

tensor2tensor/bin/t2t-datagen

@@ -35,7 +35,6 @@ import tempfile
 
 import numpy as np
 
-from tensor2tensor.data_generators import algorithmic
 from tensor2tensor.data_generators import algorithmic_math
 from tensor2tensor.data_generators import all_problems  # pylint: disable=unused-import
 from tensor2tensor.data_generators import audio
@@ -60,6 +59,8 @@ flags.DEFINE_string("tmp_dir", "/tmp/t2t_datagen",
                     "Temporary storage directory.")
 flags.DEFINE_string("problem", "",
                     "The name of the problem to generate data for.")
+flags.DEFINE_string("exclude_problems", "",
+                    "Comma-separates list of problems to exclude.")
 flags.DEFINE_integer("num_shards", 10, "How many shards to use.")
 flags.DEFINE_integer("max_cases", 0,
                      "Maximum number of cases to generate (unbounded if 0).")
@@ -74,37 +75,6 @@ flags.DEFINE_string("t2t_usr_dir", "",
 # Mapping from problems that we can generate data for to their generators.
 # pylint: disable=g-long-lambda
 _SUPPORTED_PROBLEM_GENERATORS = {
-    "algorithmic_shift_decimal40": (
-        lambda: algorithmic.shift_generator(20, 10, 40, 100000),
-        lambda: algorithmic.shift_generator(20, 10, 80, 10000)),
-    "algorithmic_reverse_binary40": (
-        lambda: algorithmic.reverse_generator(2, 40, 100000),
-        lambda: algorithmic.reverse_generator(2, 400, 10000)),
-    "algorithmic_reverse_decimal40": (
-        lambda: algorithmic.reverse_generator(10, 40, 100000),
-        lambda: algorithmic.reverse_generator(10, 400, 10000)),
-    "algorithmic_addition_binary40": (
-        lambda: algorithmic.addition_generator(2, 40, 100000),
-        lambda: algorithmic.addition_generator(2, 400, 10000)),
-    "algorithmic_addition_decimal40": (
-        lambda: algorithmic.addition_generator(10, 40, 100000),
-        lambda: algorithmic.addition_generator(10, 400, 10000)),
-    "algorithmic_multiplication_binary40": (
-        lambda: algorithmic.multiplication_generator(2, 40, 100000),
-        lambda: algorithmic.multiplication_generator(2, 400, 10000)),
-    "algorithmic_multiplication_decimal40": (
-        lambda: algorithmic.multiplication_generator(10, 40, 100000),
-        lambda: algorithmic.multiplication_generator(10, 400, 10000)),
-    "algorithmic_reverse_nlplike_decimal8K": (
-        lambda: algorithmic.reverse_generator_nlplike(8000, 70, 100000,
-                                                      10, 1.300),
-        lambda: algorithmic.reverse_generator_nlplike(8000, 70, 10000,
-                                                      10, 1.300)),
-    "algorithmic_reverse_nlplike_decimal32K": (
-        lambda: algorithmic.reverse_generator_nlplike(32000, 70, 100000,
-                                                      10, 1.050),
-        lambda: algorithmic.reverse_generator_nlplike(32000, 70, 10000,
-                                                      10, 1.050)),
     "algorithmic_algebra_inverse": (
         lambda: algorithmic_math.algebra_inverse(26, 0, 2, 100000),
         lambda: algorithmic_math.algebra_inverse(26, 3, 3, 10000)),
@@ -124,29 +94,9 @@ _SUPPORTED_PROBLEM_GENERATORS = {
                                                     2**14, 2**9),
         lambda: wsj_parsing.parsing_token_generator(FLAGS.tmp_dir, False,
                                                     2**14, 2**9)),
-    "wmt_enfr_characters": (
-        lambda: wmt.enfr_character_generator(FLAGS.tmp_dir, True),
-        lambda: wmt.enfr_character_generator(FLAGS.tmp_dir, False)),
-    "wmt_enfr_tokens_8k": (
-        lambda: wmt.enfr_wordpiece_token_generator(FLAGS.tmp_dir, True, 2**13),
-        lambda: wmt.enfr_wordpiece_token_generator(FLAGS.tmp_dir, False, 2**13)
-    ),
-    "wmt_enfr_tokens_32k": (
-        lambda: wmt.enfr_wordpiece_token_generator(FLAGS.tmp_dir, True, 2**15),
-        lambda: wmt.enfr_wordpiece_token_generator(FLAGS.tmp_dir, False, 2**15)
-    ),
-    "wmt_ende_characters": (
-        lambda: wmt.ende_character_generator(FLAGS.tmp_dir, True),
-        lambda: wmt.ende_character_generator(FLAGS.tmp_dir, False)),
     "wmt_ende_bpe32k": (
         lambda: wmt.ende_bpe_token_generator(FLAGS.tmp_dir, True),
         lambda: wmt.ende_bpe_token_generator(FLAGS.tmp_dir, False)),
-    "wmt_zhen_tokens_32k": (
-        lambda: wmt.zhen_wordpiece_token_generator(FLAGS.tmp_dir, True,
-                                                   2**15, 2**15),
-        lambda: wmt.zhen_wordpiece_token_generator(FLAGS.tmp_dir, False,
-                                                   2**15, 2**15)
-    ),
     "lm1b_32k": (
         lambda: lm1b.generator(FLAGS.tmp_dir, True),
         lambda: lm1b.generator(FLAGS.tmp_dir, False)
@@ -285,6 +235,9 @@ def main(_):
   # Calculate the list of problems to generate.
   problems = sorted(
       list(_SUPPORTED_PROBLEM_GENERATORS) + registry.list_problems())
+  for exclude in FLAGS.exclude_problems.split(","):
+    if exclude:
+      problems = [p for p in problems if exclude not in p]
   if FLAGS.problem and FLAGS.problem[-1] == "*":
     problems = [p for p in problems if p.startswith(FLAGS.problem[:-1])]
   elif FLAGS.problem:
@@ -364,7 +317,7 @@ def generate_data_for_problem(problem):
 
 def generate_data_for_registered_problem(problem_name):
   problem = registry.problem(problem_name)
-  problem.generate_data(FLAGS.data_dir, FLAGS.tmp_dir)
+  problem.generate_data(FLAGS.data_dir, FLAGS.tmp_dir, FLAGS.num_shards)
 
 
 if __name__ == "__main__":