def generate_dataset(orig_train: Iterable[RawExample],
                     orig_dev: Iterable[RawExample],
                     orig_test: Iterable[RawExample],
                     converter: casper_converters.BaseExampleConverter,
                     output_dir: str,
                     seed: int = 42,
                     log_every: int = 1000,
                     train_filename: Optional[str] = "train",
                     dev_filename: Optional[str] = "dev",
                     test_filename: Optional[str] = "test",
                     file_format: str = "tfr") -> None:
  """Generates and writes retrieval-augmented dataset files.

  Args:
    orig_train: Iterable of deserialized JSONs for training data.
    orig_dev: Iterable of deserialized JSONs for dev data.
    orig_test: Iterable of deserialized JSONs for test data.
    converter: A subclass of BaseExampleConverter.
    output_dir: Output directory.
    seed: Random seed.
    log_every: Logging frequency.
    train_filename: Training data filename prefix.
    dev_filename: Dev data filename prefix.
    test_filename: Test data filename prefix.
    file_format: Output file format.
  """
  random.seed(seed)
  tf.io.gfile.makedirs(output_dir)

  converter.stats.clear()
  examples = create_augmented_examples(
      orig_train, converter, split="train", log_every=log_every)
  if examples:
    # Shuffle the training data.
    random.shuffle(examples)
    base_path = os.path.join(output_dir, train_filename)
    write_examples(examples, base_path, file_format=file_format)
    logging.info("Train data stats: %s", dict(converter.stats))
  else:
    logging.warn("No train examples generated.")

  converter.stats.clear()
  examples = create_augmented_examples(
      orig_dev, converter, split="dev", log_every=log_every)
  if examples:
    # The dev data is not shuffled for easier error analysis.
    base_path = os.path.join(output_dir, dev_filename)
    write_examples(examples, base_path, file_format=file_format)
    logging.info("Dev data stats: %s", dict(converter.stats))
  else:
    logging.warn("No dev examples generated.")

  converter.stats.clear()
  examples = create_augmented_examples(
      orig_test, converter, split="test", log_every=log_every)
  if examples:
    # The test data is not shuffled for easier error analysis.
    base_path = os.path.join(output_dir, test_filename)
    write_examples(examples, base_path, file_format=file_format)
    logging.info("Test data stats: %s", dict(converter.stats))
  else:
    logging.warn("No test examples generated.")
Generates and writes retrieval-augmented dataset files.

  Args:
    orig_train: Iterable of deserialized JSONs for training data.
    orig_dev: Iterable of deserialized JSONs for dev data.
    orig_test: Iterable of deserialized JSONs for test data.
    converter: A subclass of BaseExampleConverter.
    output_dir: Output directory.
    seed: Random seed.
    log_every: Logging frequency.
    train_filename: Training data filename prefix.
    dev_filename: Dev data filename prefix.
    test_filename: Test data filename prefix.
    file_format: Output file format.
  
language/casper/augment/cached_retrieval_to_dataset_lib.py