def pipeline(root):
  """Maps all dev-set query ids to their episodes."""

  completed, failed = (
      root
      | 'Read JSONL' >> beam.io.ReadFromText(file_pattern=_INPUT_JSONL.value)
      | 'RunEpisodes' >> beam.ParDo(EpisodeFn()).with_outputs(
          'failed', main='completed'))
  _ = (completed | 'SaveCompleted' >> beam.io.WriteToText(_OUTPUT_PATH.value))
  _ = (
      failed
      | 'SaveFailed' >> beam.io.WriteToText(f'{_OUTPUT_PATH.value}_failed'))
Maps all dev-set query ids to their episodes.
language/search_agents/muzero/run_inference_beam.py