Repository: google-research/text-to-text-transfer-transformer

Issue:
Enable caching of additional task types
e.g., TextLineTask (#282)


Discussion:
This is awesome, thanks a lot @adarob .
I was just going to reopen issue #282  because I couldn't integrate TFDS with the T5 Tasks.
I have tested this commit, and the dataset generation function must include a new parameter "shuffle" in order to work.
Old:
`def xxx_dataset_fn(split, shuffle_files=True):`
New:
`def xxx_dataset_fn(split, shuffle_files=True, shuffle=False):`

I am not sure if it intended.
Hmm, seems to be it doesn't support multithreading, even after changing the parameters:
```
python -m t5.data.cache_tasks_main \
 --tasks=unsupervised_bio_denoising_task \
 --output_cache_dir=xxxxxxx \
 --module_import=bio_task \
--pipeline_options=["--runner=DirectRunner","--direct_num_workers 40"]
```

Also:
`--pipeline_options=["--runner=DirectRunner","--direct_num_workers 40","--direct_running_mode multi_threading"]`
Doesn't work.
I just fixed the shuffle issue you pointed out.

Currently, this only parallelizes based on how you've sharded the inputs. So if you have 10 text files, it will run on 10 workers in parallel. I could make it split it out the line by line, but there is more overhead and it's a bigger refactor. Let me see if there is an elegant way to do it.
Great, thanks for the shuffle fix.

hmmm, in my case, I have 6k files, and it runs a single thread:

```
trainFilesPath = tf.io.gfile.glob('gs://xxxx/train_*')
validFilesPath = tf.io.gfile.glob('gs://xxxxx/valid_*')

dataset_path = {
    "train": trainFilesPath,
    "validation": validFilesPath
}

def protein_dataset_fn(split, shuffle_files=True, shuffle=False):

  if shuffle_files == True and dataset_path[split] is list:
    random.shuffle(dataset_path[split])

  ds = tf.data.TextLineDataset(dataset_path[split]).skip(0)
  ds = ds.map(lambda *ex: dict(zip(["targets"], ex)))

  return ds
```

Did I miss something in my code ?
Ah yes, a generic task cannot be distributed, as mentioned in the comment. I could add a warning to this effect as well. If you use a TextLineTask, it will be distributed. 

Resolving patch:
diff --git a/t5/data/cache_tasks_main.py b/t5/data/cache_tasks_main.py
--- a/t5/data/cache_tasks_main.py
+++ b/t5/data/cache_tasks_main.py
@@ -86,17 +86,15 @@ def _import_modules(modules):
       importlib.import_module(module)
 
 
-class PreprocessAndTokenize(beam.PTransform):
-  """Preprocesses and tokenizes a Task.
-
-  Expects input to be a PCollection of sharded file paths.
+class BasePreprocessAndTokenizeTask(beam.PTransform):
+  """Abstract base class to preprocess and tokenize a Task.
 
   Returns a PCollection of tokenized example dicts containing Tensors.
   """
 
   def __init__(
       self, task, split, max_input_examples=None, modules_to_import=()):
-    """PreprocessAndTokenize constructor.
+    """BasePreprocessAndTokenizeTask constructor.
 
     Args:
       task: Task, the task to process.
@@ -109,26 +107,32 @@ def __init__(
     self._max_input_examples = max_input_examples
     self._split = split
     self._modules_to_import = modules_to_import
-    self.files = self._task.tfds_dataset.files(split)
+    self.shards = self._get_shards()
     logging.info(
-        "%s %s files: %s", task.name, split, ", ".join(
-            ["%s" % inst for inst in self.files]))
+        "%s %s shards: %s", task.name, split, ", ".join(
+            ["%s" % f for f in self.shards]))
+
+  def _get_shards(self):
+    raise NotImplementedError
+
+  def _load_shard(self, shard):
+    raise NotImplementedError
 
   def _increment_counter(self, name):
     metrics.Metrics.counter(
         str("%s_%s" % (self._task.name, self._split)), name).inc()
 
-  def _emit_tokenized_examples(self, shard_instruction):
+  def _emit_tokenized_examples(self, shard):
     """Emits examples keyed by shard path and index for a single shard."""
     _import_modules(self._modules_to_import)
-    logging.info("Processing shard: %s", shard_instruction)
+    logging.info("Processing shard: %s", shard)
     self._increment_counter("input-shards")
 
-    ds = self._task.tfds_dataset.load_shard(shard_instruction)
+    ds = self._load_shard(shard)
 
     if self._max_input_examples:
       num_shard_examples = int(
-          self._max_input_examples / len(self.files))
+          self._max_input_examples / len(self.shards))
       ds = ds.repeat().take(num_shard_examples)
 
     ds = self._task.preprocess_text(ds)
@@ -138,16 +142,58 @@ def _emit_tokenized_examples(self, shard_instruction):
 
     for ex in tfds.as_numpy(ds):
       self._increment_counter("examples")
+      logging.info(ex)
       yield ex
 
   def expand(self, pipeline):
     return (
         pipeline
-        | beam.Create(self.files)
+        | beam.Create(self.shards)
         | beam.FlatMap(self._emit_tokenized_examples)
         | beam.Reshuffle())  # Allows for additional parallelization.
 
 
+class PreprocessAndTokenizeTfdsTask(BasePreprocessAndTokenizeTask):
+  """Preprocesses and tokenizes a TfdsTask.
+
+  Returns a PCollection of tokenized example dicts containing Tensors.
+  """
+
+  def _get_shards(self):
+    return self._task.tfds_dataset.files(self._split)
+
+  def _load_shard(self, shard):
+    return self._task.tfds_dataset.load_shard(shard)
+
+
+class PreprocessAndTokenizeTextLineTask(BasePreprocessAndTokenizeTask):
+  """Preprocesses and tokenizes a TextLineTask.
+
+  Returns a PCollection of tokenized example dicts containing Tensors.
+  """
+
+  def _get_shards(self):
+    return tf.io.gfile.glob(self._task._split_to_filepattern[self._split])  # pylint:disable=protected-access
+
+  def _load_shard(self, shard):
+    return tf.data.TextLineDataset(shard).skip(self._task._skip_header_lines)  # pylint:disable=protected-access
+
+
+class PreprocessAndTokenizeGenericTask(BasePreprocessAndTokenizeTask):
+  """Preprocesses and tokenizes a generic Task.
+
+  Cannot be distributed.
+
+  Returns a PCollection of tokenized example dicts containing Tensors.
+  """
+
+  def _get_shards(self):
+    return ["Unkown"]
+
+  def _load_shard(self, unused_shard):
+    return self._task._dataset_fn(self._split, shuffle_files=False)  # pylint:disable=protected-access
+
+
 class WriteExampleTfRecord(beam.PTransform):
   """Writes examples (dicts) to a TFRecord of tf.Example protos."""
 
@@ -295,10 +341,6 @@ def run_pipeline(
       if included_regex.match(t) and not excluded_regex.match(t)]
   for task_name in task_names:
     task = t5.data.TaskRegistry.get(task_name)
-    if not isinstance(task, t5.data.TfdsTask):
-      # TODO(adarob): Add support for non-TfdsTasks.
-      logging.info("Skipping non-`TfdsTask`: '%s'", task.name)
-      continue
     if not task.supports_caching:
       logging.info(
           "Skipping task that does not support caching: '%s'", task.name)
@@ -336,11 +378,21 @@ def run_pipeline(
 
     output_dirs.append(output_dir)
 
+    if isinstance(task, t5.data.TfdsTask):
+      pat_cls = PreprocessAndTokenizeTfdsTask
+    elif isinstance(task, t5.data.TextLineTask):
+      pat_cls = PreprocessAndTokenizeTextLineTask
+    else:
+      logging.warning(
+          "Generic Task '%s' cannot be distributed. To speed up preprocessing, "
+          "use a TfdsTask or TextLineTask instead.", task.name)
+      pat_cls = PreprocessAndTokenizeGenericTask
+
     for split in task.splits:
       label = "%s_%s" % (task.name, split)
-      pat = PreprocessAndTokenize(
-          task, split, max_input_examples, modules_to_import)
-      num_shards = len(pat.files)
+
+      pat = pat_cls(task, split, max_input_examples, modules_to_import)
+      num_shards = len(pat.shards)
       examples = (
           pipeline
           | "%s_pat" % label >> pat)
diff --git a/t5/data/utils.py b/t5/data/utils.py
--- a/t5/data/utils.py
+++ b/t5/data/utils.py
@@ -367,7 +367,7 @@ def __init__(self,
                token_preprocessor=None,
                output_features=None,
                num_input_examples=None,
-               supports_caching=False,
+               supports_caching=True,
                sentencepiece_model_path=None,
                shuffle_buffer_size=_SHUFFLE_BUFFER_SIZE):
     """Task constructor.
@@ -772,7 +772,6 @@ def __init__(
       metric_fns,
       tfds_data_dir=None,
       splits=None,
-      supports_caching=True,
       **task_kwargs):
     """TfdsTask constructor.
 
@@ -793,7 +792,6 @@ def __init__(
         allowable canonical splits (e.g., 'validation') to TFDS splits or slices
         (e.g., 'train[':1%']), or None. The default, None, uses all available
         splits from the TFDS dataset info.
-      supports_caching: bool, whether or not this task supports offline caching.
       **task_kwargs: dict, additional keyword arguments for the parent `Task`
         class.
     """
@@ -815,7 +813,6 @@ def dataset_fn(split, shuffle_files):
         splits=list(splits) if splits else None,
         text_preprocessor=text_preprocessor,
         metric_fns=metric_fns,
-        supports_caching=supports_caching,
         **task_kwargs)
 
   @property
@@ -865,6 +862,9 @@ def __init__(
       **task_kwargs: dict, additional keyword arguments for the parent `Task`
         class.
     """
+    self._split_to_filepattern = split_to_filepattern
+    self._skip_header_lines = skip_header_lines
+
     def dataset_fn(split, shuffle_files):
       filepattern = split_to_filepattern[split]
 
@@ -880,7 +880,7 @@ def _read_file(fname):
     super().__init__(
         name,
         dataset_fn=dataset_fn,
-        splits=split_to_filepattern.keys(),
+        splits=list(split_to_filepattern.keys()),
         text_preprocessor=text_preprocessor,
         metric_fns=metric_fns,
         **task_kwargs)


Test patch:
diff --git a/t5/data/cache_tasks_test.py b/t5/data/cache_tasks_test.py
--- a/t5/data/cache_tasks_test.py
+++ b/t5/data/cache_tasks_test.py
@@ -49,24 +49,28 @@ def test_get_info(self):
             }
         }))
 
-  def test_pipeline(self):
+  def validate_pipeline(self, task_name):
     self.assertTrue(TaskRegistry.get("cached_task").cache_dir)
-    self.assertFalse(TaskRegistry.get("uncached_task").cache_dir)
+    task = TaskRegistry.get(task_name)
+    self.assertFalse(task.cache_dir)
 
     with TestPipeline() as p:
       output_dirs = cache_tasks_main.run_pipeline(
-          p, ["cached_task", "uncached_task"], cache_dir=self.test_data_dir)
+          p, ["cached_task", task_name], cache_dir=self.test_data_dir)
 
-    actual_task_dir = os.path.join(self.test_data_dir, "uncached_task")
+    actual_task_dir = os.path.join(self.test_data_dir, task_name)
     expected_task_dir = os.path.join(test_utils.TEST_DATA_DIR, "cached_task")
     expected_tfrecord_files = [
-        "train.tfrecord-00000-of-00002",
-        "train.tfrecord-00001-of-00002",
-        "validation.tfrecord-00000-of-00001"]
+        "train.tfrecord-00000-of-00002", "train.tfrecord-00001-of-00002",
+    ]
     expected_auxiliary_files = [
-        "stats.train.json", "stats.validation.json",
-        "info.train.json", "info.validation.json"]
+        "stats.train.json", "info.train.json"
+    ]
 
+    if "validation" in task.splits:
+      expected_tfrecord_files.append("validation.tfrecord-00000-of-00001")
+      expected_auxiliary_files.extend(
+          ["stats.validation.json", "info.validation.json"])
     self.assertEqual([actual_task_dir], output_dirs)
     self.assertCountEqual(
         expected_tfrecord_files + expected_auxiliary_files,
@@ -79,13 +83,23 @@ def test_pipeline(self):
               os.path.join(actual_task_dir, fname)).read().replace(", ", ","))
 
     # Add COMPLETED file so that we can load `uncached_task`.
-    test_utils.mark_completed(self.test_data_dir, "uncached_task")
+    test_utils.mark_completed(self.test_data_dir, task_name)
 
     # Load task.
-    uncached_task = TaskRegistry.get("uncached_task")
+    uncached_task = TaskRegistry.get(task_name)
 
     # Check datasets.
-    test_utils.verify_task_matches_fake_datasets(uncached_task, use_cached=True)
+    test_utils.verify_task_matches_fake_datasets(
+        uncached_task, use_cached=True, splits=task.splits)
+
+  def test_tfds_pipeline(self):
+    self.validate_pipeline("uncached_task")
+
+  def test_text_line_pipeline(self):
+    self.validate_pipeline("text_line_task")
+
+  def test_general_pipeline(self):
+    self.validate_pipeline("text_line_task")
 
   def test_overwrite(self):
     with TestPipeline() as p:
diff --git a/t5/data/test_utils.py b/t5/data/test_utils.py
--- a/t5/data/test_utils.py
+++ b/t5/data/test_utils.py
@@ -517,6 +517,22 @@ def _load_shard(shard_instruction):
         metric_fns=[])
     self.text_line_task = TaskRegistry.get("text_line_task")
 
+    # Prepare uncached Task.
+    def _dataset_fn(split, shuffle_files):
+      del split
+      del shuffle_files
+      filepattern = os.path.join(self.test_data_dir, "train.tsv*")
+      return tf.data.TextLineDataset(filepattern)
+    TaskRegistry.add(
+        "general_task",
+        dataset_utils.Task,
+        dataset_fn=_dataset_fn,
+        splits=["train"],
+        text_preprocessor=[_split_tsv_preprocessor, test_text_preprocessor],
+        output_features=dataset_utils.Feature(sentencepiece_vocab()),
+        metric_fns=[])
+    self.general_task = TaskRegistry.get("general_task")
+
   def tearDown(self):
     super().tearDown()
     self._tfds_patcher.stop()