code: def train_and_eval_dataset(dataset_name, data_dir):
    
    if dataset_name.startswith("v1_"):
      return _train_and_eval_dataset_v1(dataset_name[3:], data_dir)
    dataset_builder = tfds.builder(dataset_name, data_dir=data_dir)
    info = dataset_builder.info
    splits = dataset_builder.info.splits
    if tfds.Split.TRAIN not in splits:
      raise ValueError("To train we require a train split in the dataset.")
    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, valid = tfds.load(
        name=dataset_name, split=[tfds.Split.TRAIN, eval_split])
    keys = None
    if info.supervised_keys:
      keys = ([info.supervised_keys[0]], [info.supervised_keys[1]])
    return train, valid, info.features, keys
docs: Return train and evaluation datasets, feature info and supervised keys.
  
    Args:
      dataset_name: a string, the name of the dataset; if it starts with "v1_"
        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.
  
    Returns:
      a 4-tuple consisting of:
       * the train tf.data.Dataset
       * the eval tf.data.Dataset
       * information about features: a python dictionary with feature names
           as keys and an object as value that provides .shape and .num_classes.
       * supervised_keys: information what's the input and what's the target,
           ie., a pair of lists with input and target feature names.