def to_tfrecords_pipeline(
    input_pattern,
    output_pattern,
    vocab_model_file,
    task,
    delimiter_type,
    include_source = False,
    include_evidence = True,
    include_distractors = True,
    evidence_marker_type = EvidenceMarkerType.empty,
    max_input_length = 1024,
    filter_no_diff = True,
    plot_lengths = False,
):
  """Beam pipeline for converting article pairs to tfrecords."""

  def pipeline(root):
    seq2seq_inputs = (
        root
        | "Read" >> beam.io.ReadFromText(input_pattern)
        | "ParseJSONL" >> beam.Map(json.loads)
        | "ToSeq2SeqInput" >> beam.ParDo(
            ToSeq2SeqInput(
                vocab_model_file=vocab_model_file,
                task=task,
                delimiter_type=delimiter_type,
                include_source=include_source,
                include_evidence=include_evidence,
                include_distractors=include_distractors,
                evidence_marker_type=evidence_marker_type,
                max_input_length=max_input_length,
                filter_no_diff=filter_no_diff)))
    _ = (
        seq2seq_inputs
        | "ReShuffle" >> beam.Reshuffle()
        | "ToExample" >> beam.Map(tf_utils.to_example)
        | "Write" >> beam.io.tfrecordio.WriteToTFRecord(
            output_pattern,
            coder=beam.coders.ProtoCoder(tf.train.Example),
            num_shards=10))
    if plot_lengths:
      _ = (
          seq2seq_inputs
          | "GetLengths" >> beam.ParDo(GetLengths(vocab_model_file))
          | "Combine" >> beam.CombineGlobally(combine_lengths)
          | "PlotHistogram" >> beam.Map(
              plot_histogram, path=output_pattern + ".hist.png"))

  return pipeline
Beam pipeline for converting article pairs to tfrecords.
language/fruit/beam_pipelines.py