def wrap_numpy_function(
    preprocess_input_names, preprocess_fn, preprocess_output_types
) -> Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]]:
  """Wraps python function into Tensorflow op.

  Args:
    preprocess_input_names: names of all input features to the `preprocess_fn`
      function.
    preprocess_fn: a function that preprocesses a dictionary of features using
      numpy library.
    preprocess_output_types: names of all output features and their types
      returned to the `preprocess_fn` by function.

  Returns:
    A function that preprocesses features using tensorflow ops.
  """

  def preprocess_with_list_input_fn(*args) -> Tuple[tf.Tensor]:
    input_dict = {k: v for k, v in zip(preprocess_input_names, args)}
    output_dict = preprocess_fn(input_dict)
    output_list = tuple(output_dict[k] for k in preprocess_output_types)
    return output_list

  def f(example: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
    input_list = [example[k] for k in preprocess_input_names]
    output_types = [v for k, v in preprocess_output_types.items()]
    output_list = tf.numpy_function(preprocess_with_list_input_fn, input_list,
                                    output_types)
    output_dict = {
        k: v for k, v in zip(preprocess_output_types.keys(), output_list)
    }
    return output_dict

  return f
Wraps python function into Tensorflow op.

  Args:
    preprocess_input_names: names of all input features to the `preprocess_fn`
      function.
    preprocess_fn: a function that preprocesses a dictionary of features using
      numpy library.
    preprocess_output_types: names of all output features and their types
      returned to the `preprocess_fn` by function.

  Returns:
    A function that preprocesses features using tensorflow ops.
  
language/mentionmemory/utils/data_utils.py