more num_shards

Author: Huyen Nguyen

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

@@ -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))