def write_examples(examples: Sequence[AugmentedExample],
                   base_path: str,
                   file_format: str = "tfr",
                   num_shards: Optional[int] = None) -> None:
  """Writes examples to sharded TSV or TFRecord files."""
  if not num_shards:
    # Automatically compute the number of shards
    num_shards = max(1, len(examples) // 5000)
  shard_size = len(examples) // num_shards + 1
  for i in range(num_shards):
    filename = "{}.{}-{:05d}-of-{:05d}".format(base_path, file_format, i,
                                               num_shards)
    shard_examples = examples[i * shard_size:(i + 1) * shard_size]
    logging.info("Writing to %s", filename)
    if file_format == "tsv":
      with tf.io.gfile.GFile(filename, "w") as writer:
        for example in shard_examples:
          writer.write("{}\t{}\n".format(example.inputs, example.targets))
    elif file_format == "tfr":
      with tf.io.TFRecordWriter(filename) as writer:
        for example in shard_examples:
          writer.write(_to_tf_example(example))
    else:
      raise ValueError("Unknown file format: {}".format(file_format))
Writes examples to sharded TSV or TFRecord files.
language/casper/augment/cached_retrieval_to_dataset_lib.py