Return train and evaluation datasets, feature info and supervised keys.

def _train_and_eval_dataset_v1(problem_name, data_dir):
  """Return train and evaluation datasets, feature info and supervised keys."""
  problem = problems.problem(problem_name)
  train_dataset = problem.dataset(tf.estimator.ModeKeys.TRAIN, data_dir)
  train_dataset = train_dataset.map(_select_features)
  eval_dataset = problem.dataset(tf.estimator.ModeKeys.EVAL, data_dir)
  eval_dataset = eval_dataset.map(_select_features)
  supervised_keys = (["inputs"], ["targets"])
  hparams = problem.get_hparams()
  # We take a few training examples to guess the shapes.
  input_shapes, target_shapes = [], []
  for example in train_dataset.take(3):
    input_shapes.append(example["inputs"].shape.as_list())
    target_shapes.append(example["targets"].shape.as_list())
  input_vocab_size = hparams.vocab_size["inputs"]
  target_vocab_size = hparams.vocab_size["targets"]
  input_info = _make_info(input_shapes, input_vocab_size)
  target_info = _make_info(target_shapes, target_vocab_size)
  info = {"inputs": input_info, "targets": target_info}
  return train_dataset, eval_dataset, info, supervised_keys