Add support for custom prediction function in colab and jupyter notebook modes. (#1842)

* Add custom_predict_fn support

* predict_fn model name fix

* Update README.md

* minor fixes

* minor fixes

Author: tolga-b

tensorboard/plugins/interactive_inference/README.md

@@ -57,6 +57,11 @@ You can use the What-If Tool to analyze a classification or regression
 that takes TensorFlow Example or SequenceExample protos
 (data points) as inputs directly in a jupyter or colab notebook.
 
+You can also use What-If-Tool with a custom prediction function that takes
+Tensorflow examples and produces predictions. In this mode, you can load any model
+(including non-TensorFlow models) as long as your custom function's input and output
+specifications are correct.
+
 If you want to train an ML model from a dataset and explore the dataset and
 model, check out the [What_If_Tool_Notebook_Usage.ipynb notebook](https://colab.research.google.com/github/tensorflow/tensorboard/blob/master/tensorboard/plugins/interactive_inference/What_If_Tool_Notebook_Usage.ipynb) in colab, which starts from a CSV file,
 converts the data to tf.Example protos, trains a classifier, and then uses the
@@ -254,11 +259,13 @@ The WitConfigBuilder object takes a list of tf.Example or tf.SequenceExample
 protos as a constructor argument. These protos will be shown in the tool and
 inferred in the specified model.
 
-The model to be used for inference by the tool can be specified one of two ways:
+The model to be used for inference by the tool can be specified one of three ways:
 - As a TensorFlow [Estimator](https://www.tensorflow.org/guide/estimators)
   object that is provided through the `set_estimator_and_feature_spec` method.
   In this case the inference will be done inside the notebook using the
   provided estimator.
+- As a custom prediction function provided through `set_custom_predict_fn` method.
+  In this case WIT will directly call the function for inference.
 - As an endpoint for a model being served by [TensorFlow Serving](https://github.com/tensorflow/serving),
   through the `set_inference_address` and `set_model_name` methods. In this case
   the inference will be done on the model server specified. To query a model served

tensorboard/plugins/interactive_inference/utils/inference_utils.py

@@ -181,14 +181,16 @@ class ServingBundle(object):
       Predict API.
     estimator: An estimator to use instead of calling an external model.
     feature_spec: A feature spec for use with the estimator.
+    custom_predict_fn: A custom prediction function.
 
   Raises:
     ValueError: If ServingBundle fails init validation.
   """
 
   def __init__(self, inference_address, model_name, model_type, model_version,
                signature, use_predict, predict_input_tensor,
-               predict_output_tensor, estimator=None, feature_spec=None):
+               predict_output_tensor, estimator=None, feature_spec=None,
+               custom_predict_fn=None):
     """Inits ServingBundle."""
     if not isinstance(inference_address, string_types):
       raise ValueError('Invalid inference_address has type: {}'.format(
@@ -215,6 +217,7 @@ def __init__(self, inference_address, model_name, model_type, model_version,
     self.predict_output_tensor = predict_output_tensor
     self.estimator = estimator
     self.feature_spec = feature_spec
+    self.custom_predict_fn = custom_predict_fn
 
 
 def proto_value_for_feature(example, feature_name):
@@ -750,5 +753,10 @@ def run_inference(examples, serving_bundle):
     for pred in preds:
       values.append(pred[preds_key])
     return common_utils.convert_prediction_values(values, serving_bundle)
+  elif serving_bundle.custom_predict_fn:
+    # If custom_predict_fn is provided, pass examples directly for local
+    # inference.
+    values = serving_bundle.custom_predict_fn(examples)
+    return common_utils.convert_prediction_values(values, serving_bundle)
   else:
     return platform_utils.call_servo(examples, serving_bundle)

tensorboard/plugins/interactive_inference/witwidget/notebook/colab/wit.py

@@ -195,6 +195,18 @@ def __init__(self, config_builder, height=1000):
     if 'compare_estimator_and_spec' in copied_config:
       del copied_config['compare_estimator_and_spec']
 
+    self.custom_predict_fn = (
+      config.get('custom_predict_fn')
+      if 'custom_predict_fn' in config else None)
+    self.compare_custom_predict_fn = (
+      config.get('compare_custom_predict_fn')
+      if 'compare_custom_predict_fn' in config else None)
+    if 'custom_predict_fn' in copied_config:
+      del copied_config['custom_predict_fn']
+    if 'compare_custom_predict_fn' in copied_config:
+      del copied_config['compare_custom_predict_fn']
+
+
     self._set_examples(config['examples'])
     del copied_config['examples']
 
@@ -248,11 +260,13 @@ def infer(self):
       self.config.get('predict_input_tensor'),
       self.config.get('predict_output_tensor'),
       self.estimator_and_spec.get('estimator'),
-      self.estimator_and_spec.get('feature_spec'))
+      self.estimator_and_spec.get('feature_spec'),
+      self.custom_predict_fn)
     infer_objs.append(inference_utils.run_inference_for_inference_results(
         examples_to_infer, serving_bundle))
     if ('inference_address_2' in self.config or
-        self.compare_estimator_and_spec.get('estimator')):
+        self.compare_estimator_and_spec.get('estimator') or
+        self.compare_custom_predict_fn):
       serving_bundle = inference_utils.ServingBundle(
         self.config.get('inference_address_2'),
         self.config.get('model_name_2'),
@@ -263,7 +277,8 @@ def infer(self):
         self.config.get('predict_input_tensor'),
         self.config.get('predict_output_tensor'),
         self.compare_estimator_and_spec.get('estimator'),
-        self.compare_estimator_and_spec.get('feature_spec'))
+        self.compare_estimator_and_spec.get('feature_spec'),
+        self.compare_custom_predict_fn)
       infer_objs.append(inference_utils.run_inference_for_inference_results(
           examples_to_infer, serving_bundle))
     self.updated_example_indices = set()
@@ -314,9 +329,11 @@ def infer_mutants(self, info):
       self.config.get('predict_input_tensor'),
       self.config.get('predict_output_tensor'),
       self.estimator_and_spec.get('estimator'),
-      self.estimator_and_spec.get('feature_spec')))
+      self.estimator_and_spec.get('feature_spec'),
+      self.custom_predict_fn))
     if ('inference_address_2' in self.config or
-        self.compare_estimator_and_spec.get('estimator')):
+        self.compare_estimator_and_spec.get('estimator') or
+        self.compare_custom_predict_fn):
       serving_bundles.append(inference_utils.ServingBundle(
         self.config.get('inference_address_2'),
         self.config.get('model_name_2'),
@@ -327,7 +344,8 @@ def infer_mutants(self, info):
         self.config.get('predict_input_tensor'),
         self.config.get('predict_output_tensor'),
         self.compare_estimator_and_spec.get('estimator'),
-        self.compare_estimator_and_spec.get('feature_spec')))
+        self.compare_estimator_and_spec.get('feature_spec'),
+        self.compare_custom_predict_fn))
     viz_params = inference_utils.VizParams(
       info['x_min'], info['x_max'],
       scan_examples, 10,

tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/wit.py

@@ -73,6 +73,17 @@ def __init__(self, config_builder, height=1000):
     if 'compare_estimator_and_spec' in copied_config:
       del copied_config['compare_estimator_and_spec']
 
+    self.custom_predict_fn = (
+      config.get('custom_predict_fn')
+      if 'custom_predict_fn' in config else None)
+    self.compare_custom_predict_fn = (
+      config.get('compare_custom_predict_fn')
+      if 'compare_custom_predict_fn' in config else None)
+    if 'custom_predict_fn' in copied_config:
+      del copied_config['custom_predict_fn']
+    if 'compare_custom_predict_fn' in copied_config:
+      del copied_config['compare_custom_predict_fn']
+
     self._set_examples(config['examples'])
     del copied_config['examples']
 
@@ -109,11 +120,13 @@ def _infer(self, change):
       self.config.get('predict_input_tensor'),
       self.config.get('predict_output_tensor'),
       self.estimator_and_spec.get('estimator'),
-      self.estimator_and_spec.get('feature_spec'))
+      self.estimator_and_spec.get('feature_spec'),
+      self.custom_predict_fn)
     infer_objs.append(inference_utils.run_inference_for_inference_results(
       examples_to_infer, serving_bundle))
     if ('inference_address_2' in self.config or
-        self.compare_estimator_and_spec.get('estimator')):
+        self.compare_estimator_and_spec.get('estimator') or
+        self.compare_custom_predict_fn):
       serving_bundle = inference_utils.ServingBundle(
         self.config.get('inference_address_2'),
         self.config.get('model_name_2'),
@@ -124,7 +137,8 @@ def _infer(self, change):
         self.config.get('predict_input_tensor'),
         self.config.get('predict_output_tensor'),
         self.compare_estimator_and_spec.get('estimator'),
-        self.compare_estimator_and_spec.get('feature_spec'))
+        self.compare_estimator_and_spec.get('feature_spec'),
+        self.compare_custom_predict_fn)
       infer_objs.append(inference_utils.run_inference_for_inference_results(
         examples_to_infer, serving_bundle))
     self.updated_example_indices = set()
@@ -160,9 +174,11 @@ def _infer_mutants(self, change):
       self.config.get('predict_input_tensor'),
       self.config.get('predict_output_tensor'),
       self.estimator_and_spec.get('estimator'),
-      self.estimator_and_spec.get('feature_spec')))
+      self.estimator_and_spec.get('feature_spec'),
+      self.custom_predict_fn))
     if ('inference_address_2' in self.config or
-        self.compare_estimator_and_spec.get('estimator')):
+        self.compare_estimator_and_spec.get('estimator') or
+        self.compare_custom_predict_fn):
       serving_bundles.append(inference_utils.ServingBundle(
         self.config.get('inference_address_2'),
         self.config.get('model_name_2'),
@@ -173,7 +189,8 @@ def _infer_mutants(self, change):
         self.config.get('predict_input_tensor'),
         self.config.get('predict_output_tensor'),
         self.compare_estimator_and_spec.get('estimator'),
-        self.compare_estimator_and_spec.get('feature_spec')))
+        self.compare_estimator_and_spec.get('feature_spec'),
+        self.compare_custom_predict_fn))
     viz_params = inference_utils.VizParams(
       info['x_min'], info['x_max'],
       scan_examples, 10,