def shuffle_and_batch_data(dataset, target_names, features_info, training):
  """Shuffle and batch the given dataset."""
  def append_targets(example):
    """Append targets to the example dictionary. Needed for Keras."""
    if len(target_names) == 1:
      return (example, example[target_names[0]])
    targets = {}
    for name in target_names:
      targets[name] = example[name]
    return (example, targets)
  dataset = dataset.map(append_targets)
  if training:
    dataset = dataset.repeat()
  shapes = {k: features_info[k].shape for k in features_info}
  shapes = (shapes, shapes[target_names[0]])
  dataset = dataset.shuffle(128)
  dataset = preprocess_fn(dataset, training)
  dataset = batch_fn(dataset, training, shapes, target_names)
  return dataset.prefetch(8)
Shuffle and batch the given dataset.