minor changes.

PiperOrigin-RevId: 633691869

seqio/beam_utils.py

@@ -19,6 +19,7 @@
 import importlib
 import json
 import operator
+import os
 from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
 
 from absl import logging
@@ -207,6 +208,8 @@ def expand(self, pcoll):
     )
 
 
+
+
 class _ArrayRecordSink(beam.io.filebasedsink.FileBasedSink):
   """Sink Class for use in Arrayrecord PTransform."""
 

seqio/beam_utils_test.py

@@ -125,6 +125,7 @@ def test_write_example_tf_record(self):
     parsed_example = tf.train.Example.FromString(next(iter(ds)).numpy())
     self.assertEqual(parsed_example, seqio.dict_to_tfexample(example))
 
+
   def test_write_json(self):
     output_path = os.path.join(self.test_data_dir, "output.json")
     data = {

seqio/scripts/cache_tasks_main.py

@@ -145,7 +145,10 @@
 flags.DEFINE_enum(
     "output_format",
     "tfrecord",
-    ["arrayrecord", "tfrecord"],
+    [
+        "arrayrecord",
+        "tfrecord",
+    ],
     "Output format of the cached tasks.",
 )
 flags.DEFINE_boolean(
@@ -173,6 +176,8 @@ def run_pipeline(
     overwrite=False,
     ignore_other_caches=False,
     completed_file_contents="",
+    store_metadata_proto: bool = False,  # GOOGLE-INTERNAL,
+    output_format: str = "tfrecord",
 ):
   """Run preprocess pipeline."""
   output_dirs = []
@@ -304,7 +309,7 @@ def run_pipeline(
           | "%s_global_example_shuffle" % label >> beam.Reshuffle()
       )
 
-      if FLAGS.output_format == "arrayrecord":
+      if output_format == "arrayrecord":
         completion_values.append(
             examples
             | "%s_write_arrayrecord" % label
@@ -316,7 +321,7 @@ def run_pipeline(
                 preserve_random_access=FLAGS.preserve_random_access,
             )
         )
-      elif FLAGS.output_format == "tfrecord":
+      elif output_format == "tfrecord":
         completion_values.append(
             examples
             | "%s_write_tfrecord" % label
@@ -389,6 +394,7 @@ def main(_):
         FLAGS.module_import,
         FLAGS.overwrite,
         FLAGS.ignore_other_caches,
+        FLAGS.output_format,
     )
 
 

seqio/scripts/cache_tasks_test.py

@@ -24,6 +24,7 @@
 from seqio.scripts import cache_tasks_main
 import tensorflow.compat.v2 as tf
 
+
 tf.compat.v1.enable_eager_execution()
 
 flags.FLAGS.min_shards = 0