generate_toy_data.py 文件源码

def main():
  """Main function"""

  if ARGS.type == "copy":
    generate_fn = make_copy
  elif ARGS.type == "reverse":
    generate_fn = make_reverse

  # Generate dataset
  examples = list(generate_fn(ARGS.num_examples, ARGS.min_len, ARGS.max_len))
  try:
    os.makedirs(ARGS.output_dir)
  except OSError:
    if not os.path.isdir(ARGS.output_dir):
      raise

  # Write train data
  train_sources, train_targets = zip(*examples)
  write_parallel_text(train_sources, train_targets, ARGS.output_dir)