Merge pull request #165 from chiphuyen/master

Add 'num_shards" to generate_data

Author: lukaszkaiser

tensor2tensor/bin/t2t-datagen

@@ -365,7 +365,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__":

tensor2tensor/data_generators/algorithmic.py

@@ -36,10 +36,10 @@ class AlgorithmicIdentityBinary40(problem.Problem):
   def num_symbols(self):
     return 2
 
-  def generate_data(self, data_dir, _):
+  def generate_data(self, data_dir, _, num_shards=100):
     utils.generate_dataset_and_shuffle(
         identity_generator(self.num_symbols, 40, 100000),
-        self.training_filepaths(data_dir, 100, shuffled=True),
+        self.training_filepaths(data_dir, num_shards, shuffled=True),
         identity_generator(self.num_symbols, 400, 10000),
         self.dev_filepaths(data_dir, 1, shuffled=True),
         shuffle=False)

tensor2tensor/data_generators/problem.py

@@ -113,7 +113,7 @@ class Problem(object):
   # BEGIN SUBCLASS INTERFACE
   # ============================================================================
 
-  def generate_data(self, data_dir, tmp_dir):
+  def generate_data(self, data_dir, tmp_dir, num_shards=100):
     raise NotImplementedError()
 
   def hparams(self, defaults, model_hparams):

tensor2tensor/data_generators/wmt.py

@@ -49,10 +49,10 @@ def target_vocab_size(self):
   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):
+  def generate_data(self, data_dir, tmp_dir, num_shards=100):
     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),
+        self.training_filepaths(data_dir, num_shards, shuffled=False),
         ende_wordpiece_token_generator(tmp_dir, False, self.target_vocab_size),
         self.dev_filepaths(data_dir, 1, shuffled=False))
 
@@ -92,10 +92,10 @@ def target_vocab_size(self):
   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):
+  def generate_data(self, data_dir, tmp_dir, num_shards=100):
     generator_utils.generate_dataset_and_shuffle(
         mken_wordpiece_token_generator(tmp_dir, True, self.target_vocab_size),
-        self.training_filepaths(data_dir, 100, shuffled=False),
+        self.training_filepaths(data_dir, num_shards, shuffled=False),
         mken_wordpiece_token_generator(tmp_dir, False, self.target_vocab_size),
         self.dev_filepaths(data_dir, 1, shuffled=False))