def filter_for_generation_pipeline(input_pattern,
                                   output_pattern,
                                   vocab_model_file,
                                   max_article_length=512,
                                   excessive_length_strategy="truncate",
                                   max_mention_length=256,
                                   max_mentions=256,
                                   excessive_mention_strategy="truncate",
                                   use_source_mentions=True,
                                   include_new_articles=False,
                                   dry_run=False):
  """Returns filter for generation beam pipeline."""

  def pipeline(root):
    filtered = (
        root
        | "Read" >> beam.io.ReadFromText(input_pattern)
        | "ParseJSONL" >> beam.Map(json.loads)
        | "FilterForGeneration" >> beam.ParDo(
            FilterForGeneration(vocab_model_file, max_article_length,
                                excessive_length_strategy, max_mention_length,
                                max_mentions, excessive_mention_strategy,
                                use_source_mentions, include_new_articles)))
    if not dry_run:
      _ = (
          filtered
          | "ToJSONL" >> beam.Map(json.dumps)
          | "Write" >> beam.io.WriteToText(output_pattern))

  return pipeline
Returns filter for generation beam pipeline.
language/fruit/beam_pipelines.py