docs: def input_fn(data_file, num_epochs, shuffle, batch_size):
    """Generate an input function for the Estimator."""
    assert tf.gfile.Exists(data_file), (
        '%s not found. Please make sure you have run census_dataset.py and '
        'set the --data_dir argument to the correct path.' % data_file)
  
    def parse_csv(value):
      tf.logging.info('Parsing {}'.format(data_file))
      columns = tf.decode_csv(value, record_defaults=_CSV_COLUMN_DEFAULTS)
      features = dict(zip(_CSV_COLUMNS, columns))
      labels = features.pop('income_bracket')
      classes = tf.equal(labels, '>50K')  # binary classification
      return features, classes
  
    # Extract lines from input files using the Dataset API.
    dataset = tf.data.TextLineDataset(data_file)
  
    if shuffle:
      dataset = dataset.shuffle(buffer_size=_NUM_EXAMPLES['train'])
  
    dataset = dataset.map(parse_csv, num_parallel_calls=5)
  
    # We call repeat after shuffling, rather than before, to prevent separate
    # epochs from blending together.
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)
    return dataset