def _get_input_output_names(
    patterns: Union[str, Sequence[str]], decode_fn: Callable[[bytes],
                                                             Dict[str,
                                                                  tf.Tensor]],
    preprocess_fn: Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]]
) -> Tuple[Sequence[str], Dict[str, Any]]:
  """Records input and output features of decode_fn and preprocess_fn.

  Args:
    patterns: patterns for the input files.
    decode_fn: a function that decodes TF examples from binary format.
    preprocess_fn: a function that preprocess TF examples.

  Returns:
    A tuple of `input_names` and `output_types`.
      `input_names` is a list of features in the original TF records. This is
      determined by running `decode_fn` on the first sample from the data and
      recording all features it decoded.
      `output_types` is a dictionary containing all features in the preprocessed
      TF example and their types. This is determined by running `preprocess_fn`
      on the first TF Example from the data and recording
      all features and their types returned by the function.
  """
  if isinstance(patterns, list):
    input_files = tf.io.gfile.glob(patterns[0])
  else:
    input_files = tf.io.gfile.glob(patterns)
  first_record = next(iter(tf.data.TFRecordDataset(input_files[0])))
  first_example = decode_fn(first_record)
  input_names = list(first_example.keys())
  preprocessed_example = preprocess_fn(
      {k: v.numpy()[0] for k, v in first_example.items()})
  output_types = {k: v.dtype for k, v in preprocessed_example.items()}
  return input_names, output_types
Records input and output features of decode_fn and preprocess_fn.

  Args:
    patterns: patterns for the input files.
    decode_fn: a function that decodes TF examples from binary format.
    preprocess_fn: a function that preprocess TF examples.

  Returns:
    A tuple of `input_names` and `output_types`.
      `input_names` is a list of features in the original TF records. This is
      determined by running `decode_fn` on the first sample from the data and
      recording all features it decoded.
      `output_types` is a dictionary containing all features in the preprocessed
      TF example and their types. This is determined by running `preprocess_fn`
      on the first TF Example from the data and recording
      all features and their types returned by the function.
  
language/mentionmemory/utils/data_utils.py