Internal

PiperOrigin-RevId: 392122923

Author: T2T Team

tensor2tensor/data_generators/text_problems.py

@@ -835,7 +835,11 @@ def text2text_txt_tab_iterator(txt_path):
   Yields:
     {"inputs": inputs, "targets": targets}
   """
-  for line in txt_line_iterator(txt_path):
+  if txt_path.endswith(".tsv*"):
+    data_iterator = txt_line_sharded_iterator(txt_path)
+  else:
+    data_iterator = txt_line_iterator(txt_path)
+  for line in data_iterator:
     if line and "\t" in line:
       parts = line.split("\t", 1)
       inputs, targets = parts[:2]

tensor2tensor/data_generators/translate.py

@@ -344,5 +344,13 @@ def generate_text_for_vocab(self, data_dir, tmp_dir):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     data_path = self.source_data_files(dataset_split)[0]
-    assert tf.gfile.Exists(data_path)
     return text_problems.text2text_txt_tab_iterator(data_path)
+
+
+class TranslateSamanantarProblem(TranslateWmt20Problem):
+  """Base class for Samanantar Datasets."""
+
+  def generate_samples(self, data_dir, tmp_dir, dataset_split):
+    src_data_path = self.source_data_files(dataset_split)[0]
+    tgt_data_path = self.source_data_files(dataset_split)[1]
+    return text_problems.text2text_txt_iterator(src_data_path, tgt_data_path)