def _train_and_eval_dataset(dataset_name,
                            data_dir,
                            eval_holdout_size,
                            train_shuffle_files=True,
                            eval_shuffle_files=False,
                            use_alt_eval=False,
                            subsplit=None):
  """Return train and evaluation datasets, feature info and supervised keys.

  Args:
    dataset_name: a string, the name of the dataset; if it starts with 't2t_'
      then we'll search T2T Problem registry for it, otherwise we assume it is a
      dataset from TFDS and load it from there.
    data_dir: directory where the data is located.
    eval_holdout_size: float from 0 to <1; if >0 use this much of training data
      for evaluation (instead of looking for a pre-specified VALIDATION split).
    train_shuffle_files: Boolean determining whether or not to shuffle the train
      files at startup. Set to False if you want data determinism.
    eval_shuffle_files: Boolean determining whether or not to shuffle the test
      files at startup. Set to False if you want data determinism.
    use_alt_eval: If True, use the dataset's alternate/secondary eval split;
      else use the dataset's default/only eval split. Currently, only the
      `glue/mnli` dataset provides an alternate eval split, and this arg is
      ignored for other datasets.
    subsplit: a pair of floats (x, y), both in [0, 1], saying which part of the
      full training dataset we should return (default: all of it, [0, 1]).

  Returns:
    a 4-tuple consisting of:
     * the train tf.Dataset
     * the eval tf.Dataset
     * information about features: a python dictionary with feature names
         as keys and an object as value that provides .shape and .n_classes.
     * supervised_keys: information what's the input and what's the target,
         ie., a pair of lists with input and target feature names.
  """
  logging.info('Building TF data pipeline for %s', dataset_name)
  if dataset_name.startswith('t2t_'):
    return _train_and_eval_dataset_v1(dataset_name[4:], data_dir,
                                      train_shuffle_files, eval_shuffle_files)
  dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
  info = dataset_builder.info
  splits = dataset_builder.info.splits
  if dataset_name != 'c4/multilingual' and tfds.Split.TRAIN not in splits:
    raise ValueError('To train we require a train split in the dataset.')
  train_split = tfds.Split.TRAIN if dataset_name != 'c4/multilingual' else 'en'
  eval_split = None
  train_examples = info.splits[train_split].num_examples
  eval_holdout_examples = int(train_examples * eval_holdout_size)
  if eval_holdout_examples > 0 or subsplit is not None:
    if subsplit is None:
      subsplit = (0, 1)
    n_train = train_examples - eval_holdout_examples
    train_start = int(n_train * subsplit[0])
    train_end = int(n_train * subsplit[1])
    if train_end - train_start < 1:
      raise ValueError('Requested train subsplit has no examples: '
                       'n_train %d subsplit %s' % (n_train, subsplit))
    # Eval holdout examples from the end of the training set.
    if eval_holdout_examples > 0:
      eval_split = f'{train_split}[-{eval_holdout_examples}:]'
    # Shard the training set for this host.
    train_split = f'{train_split}[{train_start}:{train_end}]'

  if dataset_name == 'glue/mnli':
    eval_split = (
        'validation_mismatched' if use_alt_eval else 'validation_matched')
  elif dataset_name == 'c4/multilingual':
    eval_split = 'en-validation'
  elif eval_split is None:
    if tfds.Split.VALIDATION not in splits and 'test' not in splits:
      raise ValueError('We require a validation or test split in the dataset.')
    eval_split = tfds.Split.VALIDATION
    if tfds.Split.VALIDATION not in splits:
      eval_split = tfds.Split.TEST

  train = tfds.load(
      name=dataset_name,
      split=train_split,
      data_dir=data_dir,
      shuffle_files=train_shuffle_files)
  valid = tfds.load(
      name=dataset_name,
      split=eval_split,
      data_dir=data_dir,
      shuffle_files=eval_shuffle_files)
  keys = None
  if info.supervised_keys:
    keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
  return train, valid, keys
Return train and evaluation datasets, feature info and supervised keys.

  Args:
    dataset_name: a string, the name of the dataset; if it starts with 't2t_'
      then we'll search T2T Problem registry for it, otherwise we assume it is a
      dataset from TFDS and load it from there.
    data_dir: directory where the data is located.
    eval_holdout_size: float from 0 to <1; if >0 use this much of training data
      for evaluation (instead of looking for a pre-specified VALIDATION split).
    train_shuffle_files: Boolean determining whether or not to shuffle the train
      files at startup. Set to False if you want data determinism.
    eval_shuffle_files: Boolean determining whether or not to shuffle the test
      files at startup. Set to False if you want data determinism.
    use_alt_eval: If True, use the dataset's alternate/secondary eval split;
      else use the dataset's default/only eval split. Currently, only the
      `glue/mnli` dataset provides an alternate eval split, and this arg is
      ignored for other datasets.
    subsplit: a pair of floats (x, y), both in [0, 1], saying which part of the
      full training dataset we should return (default: all of it, [0, 1]).

  Returns:
    a 4-tuple consisting of:
     * the train tf.Dataset
     * the eval tf.Dataset
     * information about features: a python dictionary with feature names
         as keys and an object as value that provides .shape and .n_classes.
     * supervised_keys: information what's the input and what's the target,
         ie., a pair of lists with input and target feature names.
  
trax/data/tf_inputs.py