def _make_tfrecord_vocabulary_dataset(vocab_path,
                                      key_dtype=tf.string,
                                      value_dtype=tf.int64,
                                      return_indicator_as_value=False,
                                      has_indicator=False):
  """Makes a (key, value) dataset from a compressed tfrecord file."""
  if not (value_dtype.is_floating or value_dtype.is_integer):
    raise ValueError('value_dtype must be numeric. Got: %s' % value_dtype)
  dataset = tf.data.TFRecordDataset(vocab_path, compression_type='GZIP')
  key_dtype_fn = _make_vocab_entry_to_dtype_fn(key_dtype)
  value_dtype_fn = _make_vocab_entry_to_dtype_fn(value_dtype)

  if return_indicator_as_value:
    assert has_indicator

    def convert_dtype(k, v):
      return key_dtype_fn(k), value_dtype_fn(v)

    return dataset.map(
        _split_vocabulary_entries,
        num_parallel_calls=tf.data.experimental.AUTOTUNE).map(convert_dtype)

  else:
    if has_indicator:
      drop_indicator = lambda k, v: k
      dataset = dataset.map(
          _split_vocabulary_entries,
          num_parallel_calls=tf.data.experimental.AUTOTUNE).map(drop_indicator)

    def convert_dtype_and_swap(v, k):
      return key_dtype_fn(k), tf.cast(v, value_dtype)

    return dataset.enumerate().map(convert_dtype_and_swap)
Makes a (key, value) dataset from a compressed tfrecord file.
tensorflow_transform/tf_utils.py