code: def set_custom_predict_fn(self, predict_fn):
      
      # 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
docs: 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.