def create_training_dataset(input_file, batch_size, config):
  """Returns `tf.data.Dataset` for training."""
  name_to_features = {}

  name_to_features["wordpiece_ids"] = _create_int_feature(
      config["max_num_wordpieces"])
  name_to_features["num_wordpieces"] = _create_int_feature(1)

  name_to_features["application_span_begin"] = _create_int_feature(
      config["max_num_applications"])
  name_to_features["application_span_end"] = _create_int_feature(
      config["max_num_applications"])
  name_to_features["application_rule_idx"] = _create_int_feature(
      config["max_num_applications"])

  name_to_features["nu_node_type"] = _create_int_feature(
      config["max_num_numerator_nodes"])
  name_to_features["nu_node_1_idx"] = _create_int_feature(
      config["max_num_numerator_nodes"])
  name_to_features["nu_node_2_idx"] = _create_int_feature(
      config["max_num_numerator_nodes"])
  name_to_features["nu_application_idx"] = _create_int_feature(
      config["max_num_numerator_nodes"])
  name_to_features["nu_num_nodes"] = _create_int_feature(1)

  name_to_features["de_node_type"] = _create_int_feature(
      config["max_num_denominator_nodes"])
  name_to_features["de_node_1_idx"] = _create_int_feature(
      config["max_num_denominator_nodes"])
  name_to_features["de_node_2_idx"] = _create_int_feature(
      config["max_num_denominator_nodes"])
  name_to_features["de_application_idx"] = _create_int_feature(
      config["max_num_denominator_nodes"])
  name_to_features["de_num_nodes"] = _create_int_feature(1)

  if "*" in input_file:
    # Potentially match multiple input files.
    files = tf.io.matching_files(input_file)
    files = tf.random.shuffle(files)
    shards = tf.data.Dataset.from_tensor_slices(files)
    dataset = shards.interleave(tf.data.TFRecordDataset)
  else:
    # Only using single input file.
    dataset = tf.data.TFRecordDataset(input_file)

  dataset = dataset.repeat()
  dataset = dataset.shuffle(buffer_size=1000)

  decode_fn = lambda record: _decode_record(record, name_to_features)
  dataset = dataset.map(
      decode_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)

  # Send the single file to all workers.
  options = tf.data.Options()
  options.experimental_distribute.auto_shard_policy = (
      tf.data.experimental.AutoShardPolicy.OFF)
  dataset = dataset.with_options(options)

  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(1024)
  return dataset
Returns `tf.data.Dataset` for training.
language/compgen/nqg/model/parser/training/input_utils.py