Sets a custom function for inference.

    Instead of using TF Serving to host a model for WIT to query, WIT can
    directly use a custom function as the model to query. In this case, the
    provided function should accept example protos and return:
      - For classification: A 2D list of numbers. The first dimension is for
        each example being predicted. The second dimension are the probabilities
        for each class ID in the prediction.
      - For regression: A 1D list of numbers, with a regression score for each
        example being predicted.

    Args:
      predict_fn: The custom python function which will be used for model
      inference.

    Returns:
      self, in order to enabled method chaining.

def set_custom_predict_fn(self, predict_fn):
    """Sets a custom function for inference.

    Instead of using TF Serving to host a model for WIT to query, WIT can
    directly use a custom function as the model to query. In this case, the
    provided function should accept example protos and return:
      - For classification: A 2D list of numbers. The first dimension is for
        each example being predicted. The second dimension are the probabilities
        for each class ID in the prediction.
      - For regression: A 1D list of numbers, with a regression score for each
        example being predicted.

    Args:
      predict_fn: The custom python function which will be used for model
      inference.

    Returns:
      self, in order to enabled method chaining.
    """
    # If estimator is set, remove it before setting predict_fn
    self.delete('estimator_and_spec')

    self.store('custom_predict_fn', predict_fn)
    self.set_inference_address('custom_predict_fn')
    # If no model name has been set, give a default
    if not self.has_model_name():
      self.set_model_name('1')
    return self