diff --git a/tensorboard/plugins/interactive_inference/README.md b/tensorboard/plugins/interactive_inference/README.md
index 0bb75303d0..26515d525a 100644
--- a/tensorboard/plugins/interactive_inference/README.md
+++ b/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).
+
diff --git a/tensorboard/plugins/interactive_inference/tf_interactive_inference_dashboard/tf-interactive-inference-dashboard.html b/tensorboard/plugins/interactive_inference/tf_interactive_inference_dashboard/tf-interactive-inference-dashboard.html
index 3f1c5384a7..0524f0bc73 100644
--- a/tensorboard/plugins/interactive_inference/tf_interactive_inference_dashboard/tf-interactive-inference-dashboard.html
+++ b/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({
diff --git a/tensorboard/plugins/interactive_inference/witwidget/notebook/colab/wit.py b/tensorboard/plugins/interactive_inference/witwidget/notebook/colab/wit.py
index 9aee2b672e..45c6437723 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/notebook/colab/wit.py
+++ b/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)
diff --git a/tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/js/lib/wit.js b/tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/js/lib/wit.js
index 29406d203b..83ef0d5649 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/js/lib/wit.js
+++ b/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
diff --git a/tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/wit.py b/tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/wit.py
index f673b2ccf9..47900f39c9 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/notebook/jupyter/wit.py
+++ b/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()
diff --git a/tensorboard/plugins/interactive_inference/witwidget/notebook/visualization.py b/tensorboard/plugins/interactive_inference/witwidget/notebook/visualization.py
index 82ebfcc9a6..834ca1da42 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/notebook/visualization.py
+++ b/tensorboard/plugins/interactive_inference/witwidget/notebook/visualization.py
@@ -31,18 +31,39 @@ class WitConfigBuilder(object):
   """Configuration builder for WitWidget settings."""
 
   def __init__(self, examples):
+    """Constructs the WitConfigBuilder object.
+
+    Args:
+      examples: A list of tf.Example or tf.SequenceExample proto objects.
+      These are the examples that will be displayed in WIT. If not model to
+      infer these examples with is specified through the methods on this class,
+      then WIT will display the examples for exploration, but no model inference
+      will be performed by the tool.
+    """
     self.config = {}
     self.set_examples(examples)
     self.set_model_type('classification')
     self.set_label_vocab([])
   
   def build(self):
+    """Returns the configuration set through use of this builder object.
+
+    Used by WitWidget to set the settings on an instance of the What-If Tool.
+    """
     return self.config
   
   def store(self, key, value):
     self.config[key] = value
-  
+
   def set_examples(self, examples):
+    """Sets the examples to be displayed in WIT.
+    
+    Args:
+      examples: List of example protos.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('examples', examples)
     if len(examples) > 0:
       self.store('are_sequence_examples',
@@ -50,66 +71,265 @@ def set_examples(self, examples):
     return self
       
   def set_model_type(self, model):
+    """Sets the type of the model being used for inference.
+
+    Args:
+      model: The model type, such as "classification" or "regression".
+      The model type defaults to "classification".
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_type', model)
     return self
   
   def set_inference_address(self, address):
+    """Sets the inference address for model inference through TF Serving.
+
+    Args:
+      address: The address of the served model, including port, such as
+      "localhost:8888".
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('inference_address', address)
     return self
     
   def set_model_name(self, name):
+    """Sets the model name for model inference through TF Serving.
+
+    Setting a model name is required if inferring through a model hosted by
+    TF Serving.
+
+    Args:
+      name: The name of the model to be queried through TF Serving at the
+      address provided by set_inference_address.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_name', name)
     return self
   
   def set_model_version(self, version):
+    """Sets the optional model version for model inference through TF Serving.
+
+    Args:
+      version: The string version number of the model to be queried through TF
+      Serving. This is optional, as TF Serving will use the latest model version
+      if none is provided.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_version', version)
     return self
       
   def set_model_signature(self, signature):
+    """Sets the optional model signature for model inference through TF Serving.
+
+    Args:
+      signature: The string signature of the model to be queried through TF
+      Serving. This is optional, as TF Serving will use the default model
+      signature if none is provided.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_signature', signature)
     return self
       
   def set_compare_inference_address(self, address):
+    """Sets the inference address for model inference for a second model hosted
+    by TF Serving.
+
+    If you wish to compare the results of two models in WIT, use this method
+    to setup the details of the second model.
+
+    Args:
+      address: The address of the served model, including port, such as
+      "localhost:8888".
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('inference_address_2', address)
     return self
       
   def set_compare_model_name(self, name):
+    """Sets the model name for a second model hosted by TF Serving.
+
+    If you wish to compare the results of two models in WIT, use this method
+    to setup the details of the second model.
+
+    Setting a model name is required if inferring through a model hosted by
+    TF Serving.
+
+    Args:
+      name: The name of the model to be queried through TF Serving at the
+      address provided by set_compare_inference_address.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_name_2', name)
     return self
   
   def set_compare_model_version(self, version):
+    """Sets the optional model version for a second model hosted by TF Serving.
+
+    If you wish to compare the results of two models in WIT, use this method
+    to setup the details of the second model.
+
+    Args:
+      version: The string version number of the model to be queried through TF
+      Serving. This is optional, as TF Serving will use the latest model version
+      if none is provided.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_version_2', version)
     return self
       
   def set_compare_model_signature(self, signature):
+    """Sets the optional model signature for a second model hosted by TF
+    Serving.
+
+    If you wish to compare the results of two models in WIT, use this method
+    to setup the details of the second model.
+
+    Args:
+      signature: The string signature of the model to be queried through TF
+      Serving. This is optional, as TF Serving will use the default model
+      signature if none is provided.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('model_signature_2', signature)
     return self
       
   def set_uses_predict_api(self, predict):
+    """Indicates that the model uses the Predict API, as opposed to the
+    Classification or Regression API.
+
+    If the model doesn't use the standard Classification or Regression APIs
+    provided through TF Serving, but instead uses the more flexible Predict API,
+    then use this method to indicate that. If this is true, then use the
+    set_predict_input_tensor and set_predict_output_tensor methods to indicate
+    the names of the tensors that are used as the input and output for the
+    models provided in order to perform the appropriate inference request.
+
+    Args:
+      predict: True if the model or models use the Predict API.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('uses_predict_api', predict)
     return self
   
   def set_max_classes_to_display(self, max_classes):
+    """Sets the maximum number of class results to display for multiclass
+    classification models.
+
+    When using WIT with a multiclass model with a large number of possible
+    classes, it can be helpful to restrict WIT to only display some smaller
+    number of the highest-scoring classes as inference results for any given
+    example. This method sets that limit.
+
+    Args:
+      max_classes: The maximum number of classes to display for inference
+      results for multiclass classification models.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('max_classes', max_classes)
     return self
   
   def set_multi_class(self, multiclass):
+    """Sets if the model(s) to query are mutliclass classification models.
+
+    Args:
+      multiclass: True if the model or models are multiclass classififcation
+      models. Defaults to false.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('multiclass', multiclass)
     return self
       
   def set_predict_input_tensor(self, tensor):
+    """Sets the name of the input tensor for models that use the Predict API.
+
+    If using WIT with set_uses_predict_api(True), then call this to specify
+    the name of the input tensor of the model or models that accepts the
+    example proto for inference.
+
+    Args:
+      tensor: The name of the input tensor.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('predict_input_tensor', tensor)
     return self
 
   def set_predict_output_tensor(self, tensor):
+    """Sets the name of the output tensor for models that use the Predict API.
+
+    If using WIT with set_uses_predict_api(True), then call this to specify
+    the name of the output tensor of the model or models that returns the
+    inference results to be explored by WIT.
+
+    Args:
+      tensor: The name of the output tensor.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('predict_output_tensor', tensor)
     return self
 
   def set_label_vocab(self, vocab):
+    """Sets the string value of numeric labels for classification models.
+
+    For classification models, the model returns scores for each class ID
+    number (classes 0 and 1 for binary classification models). In order for
+    WIT to visually display the results in a more-readable way, you can specify
+    string labels for each class ID.
+
+    Args:
+      vocab: A list of strings, where the string at each index corresponds to
+      the label for that class ID. For example ['<=50K', '>50K'] for the UCI
+      census binary classification task.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('label_vocab', vocab)
     return self
 
   def set_estimator_and_feature_spec(self, estimator, feature_spec):
+    """Sets the model for inference as a TF Estimator.
+
+    Instead of using TF Serving to host a model for WIT to query, WIT can
+    directly use a TF Estimator object as the model to query. In order to
+    accomplish this, a feature_spec must also be provided to parse the
+    example protos for input into the estimator.
+
+    Args:
+      estimator: The TF Estimator which will be used for model inference.
+      feature_spec: The feature_spec object which will be used for exmaple
+      parsing.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('estimator_and_spec', {
       'estimator': estimator, 'feature_spec': feature_spec})
     self.set_inference_address('estimator')
@@ -117,6 +337,24 @@ def set_estimator_and_feature_spec(self, estimator, feature_spec):
     return self
   
   def set_compare_estimator_and_feature_spec(self, estimator, feature_spec):
+    """Sets a second model for inference as a TF Estimator.
+
+    If you wish to compare the results of two models in WIT, use this method
+    to setup the details of the second model.
+
+    Instead of using TF Serving to host a model for WIT to query, WIT can
+    directly use a TF Estimator object as the model to query. In order to
+    accomplish this, a feature_spec must also be provided to parse the
+    example protos for input into the estimator.
+
+    Args:
+      estimator: The TF Estimator which will be used for model inference.
+      feature_spec: The feature_spec object which will be used for exmaple
+      parsing.
+
+    Returns:
+      self, in order to enabled method chaining.
+    """
     self.store('compare_estimator_and_spec', {
       'estimator': estimator, 'feature_spec': feature_spec})
     self.set_compare_inference_address('estimator')
diff --git a/tensorboard/plugins/interactive_inference/witwidget/version.py b/tensorboard/plugins/interactive_inference/witwidget/version.py
index 145d5c9da4..8a8afb622b 100644
--- a/tensorboard/plugins/interactive_inference/witwidget/version.py
+++ b/tensorboard/plugins/interactive_inference/witwidget/version.py
@@ -14,4 +14,4 @@
 
 """Contains the version string."""
 
-VERSION = '0.1.6a1'
+VERSION = '0.1.6a4'