internal

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
@@ -205,6 +206,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

@@ -142,10 +142,11 @@
 )
 
 
+_OUTPUT_FORMATS = ["arrayrecord", "tfrecord"]
 flags.DEFINE_enum(
     "output_format",
     "tfrecord",
-    ["arrayrecord", "tfrecord"],
+    _OUTPUT_FORMATS,
     "Output format of the cached tasks.",
 )
 flags.DEFINE_boolean(
@@ -173,6 +174,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 +307,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 +319,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 +392,7 @@ def main(_):
         FLAGS.module_import,
         FLAGS.overwrite,
         FLAGS.ignore_other_caches,
+        FLAGS.output_format,
     )
 
 

seqio/scripts/cache_tasks_test.py

@@ -20,10 +20,12 @@
 from absl import flags
 from absl.testing import absltest
 from apache_beam.testing.test_pipeline import TestPipeline
+import bagz
 import seqio
 from seqio.scripts import cache_tasks_main
 import tensorflow.compat.v2 as tf
 
+
 tf.compat.v1.enable_eager_execution()
 
 flags.FLAGS.min_shards = 0