def maybe_format_vocabulary_input(
    x: common_types.ConsistentTensorType,
) -> common_types.ConsistentTensorType:
  """Formats string vocabulary input.

  Args:
    x: a tensor containing the vocabulary.

  Returns:
    A similarly typed tensor.
  """

  if x.dtype == tf.string:
    # b/62379925: This is a workaround to allow tokens to contain spaces when
    # store_frequency=True, which should eventaully be removed.
    def map_spaces(t: Union[tf.Tensor, tf.RaggedTensor]
                   ) -> Union[tf.Tensor, tf.RaggedTensor]:
      return tf.strings.regex_replace(t, ' ', '__SPACE__')

    return _map_values(map_spaces, x)
  return x
Formats string vocabulary input.

  Args:
    x: a tensor containing the vocabulary.

  Returns:
    A similarly typed tensor.
  
tensorflow_transform/tf_utils.py