take in account FLAGS.num_shards

Author: Huyen Nguyen

tensor2tensor/bin/t2t-datagen

@@ -356,7 +356,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/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))