WIT notebook mode cleanup/updates (#1788)

* Add flags to TBContext

* Add flags to TBContext in application

* Add flags to tbcontext

* gitignore changes

* hide images, fix bug

* doc updates

* update

* typo

* merge thing

tensorboard/plugins/interactive_inference/README.md

@@ -242,7 +242,36 @@ We imagine WIT to be useful for a wide variety of users.
 * Lay users - Learn about machine learning by interactively playing with
   datasets and models.
 
-## How do I use it in a Jupyter notebook?
+## Notebook mode details
+
+As seen in the [example notebook](https://colab.research.google.com/github/tensorflow/tensorboard/blob/master/tensorboard/plugins/interactive_inference/What_If_Tool_Notebook_Usage.ipynb),
+creating the `WitWidget` object is what causes the What-If Tool to be displayed
+in an output cell. The `WitWidget` object takes a `WitConfigBuilder` object as a
+constructor argument. The `WitConfigBuilder` object specifies the data and model
+information that the What-If Tool will use.
+
+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:
+- 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 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
+  on host "localhost" on port 8888, named "my_model", you would set on your
+  builder
+  `builder.set_inference_address('localhost:8888').set_model_name('my_model')`.
+
+See the documentation of [WitConfigBuilder]https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/interactive_inference/witwidget/notebook/visualization.py)
+for all options you can provide, including how to specify other model types
+(defaults to binary classification) and how to specify an optional second model
+to compare to the first model.
+
+### How do I enable it for use in a Jupyter notebook?
 First, install and enable WIT for Jupyter through the following commands:
 ```sh
 pip install witwidget
@@ -253,7 +282,7 @@ jupyter nbextension enable --py --sys-prefix witwidget
 Then, use it as seen at the bottom of the
 [What_If_Tool_Notebook_Usage.ipynb notebook](./What_If_Tool_Notebook_Usage.ipynb).
 
-## How do I use it in a Colab notebook?
+### How do I enable it for use in a Colab notebook?
 Install the widget into the runtime of the notebook kernel by running a cell
 containing:
 ```
@@ -262,3 +291,4 @@ containing:
 
 Then, use it as seen at the bottom of 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).
+

tensorboard/plugins/interactive_inference/tf_interactive_inference_dashboard/tf-interactive-inference-dashboard.html

@@ -1099,28 +1099,36 @@ <h2>Create a distance feature</h2>
                           all the features and values associated with that example. Some of the things you can
                           do in the datapoint editor are:</div>
                         <div class="datapoint-info-bullet">
-                          <img class="doc-image" src="editedexample.png">
+                          <template is="dom-if" if="[[!hideImages]]">
+                            <img class="doc-image" src="editedexample.png">
+                          </template>
                           <div class="datapoint-info-block">
                             <div class="bold">Test inference on edited values</div>
                             <div>Edit features and run inference to see how your model performs</div>
                           </div>
                         </div>
                         <div class="datapoint-info-bullet">
-                          <img class="doc-image" src="distance.png">
+                          <template is="dom-if" if="[[!hideImages]]">
+                            <img class="doc-image" src="distance.png">
+                          </template>
                           <div class="datapoint-info-block">
                             <div class="bold">Compute distances from a selected datapoint</div>
                             <div>Have the selected example be an anchor and create a new distance feature for all loaded examples</div>
                           </div>
                         </div>
                         <div class="datapoint-info-bullet">
-                          <img class="doc-image" src="explorecounterfactuals.png">
+                          <template is="dom-if" if="[[!hideImages]]">
+                            <img class="doc-image" src="explorecounterfactuals.png">
+                          </template>
                           <div class="datapoint-info-block">
                             <div class="bold">Find closest counterfactuals</div>
                             <div>See the closest example with a different classification</div>
                           </div>
                         </div>
                         <div class="datapoint-info-bullet">
-                          <img class="doc-image" src="pdplots.png">
+                          <template is="dom-if" if="[[!hideImages]]">
+                            <img class="doc-image" src="pdplots.png">
+                          </template>
                           <div class="datapoint-info-block">
                             <div class="bold">Partial Dependence Plots</div>
                             <div>Explore plots for every feature that show the change in inference results across different valid values for that feature</div>
@@ -1314,7 +1322,7 @@ <h2>Create a distance feature</h2>
                             <div class="flex-wrap">
                               <template is="dom-repeat" items="{{featureValueThreshold.threshold}}">
                                 <div>
-                                  <div hidden$=[[showModelNumbers_(featureValueThreshold)]]>Model: [[index]]</div>
+                                  <div hidden$=[[showModelNumbers_(featureValueThreshold)]]>Model: [[getPrintableModelNumber_(index)]]</div>
                                   <div class="flex">
                                     <paper-slider class="slider" editable=true min="0" max="1" step="0.01" immediate-value="{{item.threshold}}"
                                       value="[[item.threshold]]" on-value-changed="refreshInferencesNoRegen_">
@@ -1729,6 +1737,12 @@ <h2>Create a distance feature</h2>
           value: false,
           observer: 'localChanged_',
         },
+        // Boolean for if documentation images should be hidden, needed when
+        // running in an environment where the images aren't available.
+        hideImages: {
+          type: Boolean,
+          value: false,
+        },
         // Atlas URL for "local" demo mode.
         localAtlasUrl: String,
         // List of features for which we can make partial dep plots.
@@ -4373,6 +4387,10 @@ <h2>Create a distance feature</h2>
         return this.isBinaryClassification_(modelType, multiClass) ?
             'Compare Slices + Fairness Metrics' : 'Compare Slices';
       },
+
+      getPrintableModelNumber_: function(index) {
+        return index + 1;
+      },
     });
 
     tf_tensorboard.registerDashboard({

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

@@ -54,7 +54,7 @@ def infer_mutants(wit_id, details):
 
 # HTML/javascript for the WIT frontend.
 WIT_HTML = """
-  <tf-interactive-inference-dashboard id="wit" local>
+  <tf-interactive-inference-dashboard id="wit" local hide-images>
   </tf-interactive-inference-dashboard>
   <script>
     const examples = {examples};
@@ -175,6 +175,12 @@ class WitWidget(object):
   index = 0
 
   def __init__(self, config_builder, height=1000):
+    """Constructor for colab notebook WitWidget.
+
+    Args:
+      config_builder: WitConfigBuilder object containing settings for WIT.
+      height: Optional height in pixels for WIT to occupy. Defaults to 1000.
+    """
     tf.logging.set_verbosity(tf.logging.WARN)
     config = config_builder.build()
     copied_config = dict(config)

tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/js/lib/wit.js

@@ -63,6 +63,7 @@ var WITView = widgets.DOMWidgetView.extend({
     this.view_ = document.createElement(
       'tf-interactive-inference-dashboard');
     this.view_.local = true;
+    this.view_.hideImages = true;
     this.el.appendChild(this.view_);
 
     // Add listeners for changes from WIT Polymer element. Passes changes

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

@@ -52,6 +52,12 @@ class WitWidget(widgets.DOMWidget):
   sprite = Unicode('').tag(sync=True)
 
   def __init__(self, config_builder, height=1000):
+    """Constructor for Jupyter notebook WitWidget.
+
+    Args:
+      config_builder: WitConfigBuilder object containing settings for WIT.
+      height: Optional height in pixels for WIT to occupy. Defaults to 1000.
+    """
     super(WitWidget, self).__init__(layout=Layout(height='%ipx' % height))
     tf.logging.set_verbosity(tf.logging.WARN)
     config = config_builder.build()