text: # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Visualization API."""
  import sys
  import tensorflow as tf
  
  
  def _is_colab():
    return "google.colab" in sys.modules
  
  
  if _is_colab():
    from witwidget.notebook.colab.wit import *  # pylint: disable=wildcard-import,g-import-not-at-top
  else:
    from witwidget.notebook.jupyter.wit import *  # pylint: disable=wildcard-import,g-import-not-at-top
  
  
  class WitConfigBuilder(object):
    """Configuration builder for WitWidget settings."""
  
    def __init__(self, examples, feature_names=None):
      """Constructs the WitConfigBuilder object.
  
      Args:
        examples: A list of tf.Example or tf.SequenceExample proto objects, or
        raw JSON objects. JSON is allowed only for AI Platform-hosted models (see
        'set_ai_platform_model' and 'set_compare_ai_platform_model methods).
        These are the examples that will be displayed in WIT. If no 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.
        feature_names: Optional, defaults to None. If examples are provided as
        JSON lists of numbers (not as feature dictionaries), then this array
        maps indices in the feature value lists to human-readable names of those
        features, used for display purposes.
      """
      self.config = {}
      self.set_model_type('classification')
      self.set_label_vocab([])
      self.set_examples(examples, feature_names)
  
    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 delete(self, key):
      if key in self.config:
        del self.config[key]
  
    def set_examples(self, examples, feature_names=None):
      """Sets the examples to be displayed in WIT.
  
      Args:
        examples: List of example protos or JSON objects.
        feature_names: Optional, defaults to None. If examples are provided as
        JSON lists of numbers (not as feature dictionaries), then this array
        maps indices in the feature value lists to human-readable names of those
        features, used just for display purposes.
  
      Returns:
        self, in order to enabled method chaining.
      """
      self.store('examples', examples)
      if feature_names:
        self.store('feature_names', feature_names)
      if len(examples) > 0 and not (
        isinstance(examples[0], tf.train.Example) or
        isinstance(examples[0], tf.train.SequenceExample)):
        self._set_uses_json_input(True)
        if isinstance(examples[0], list):
          self._set_uses_json_list(True)
      elif len(examples) > 0:
        self.store('are_sequence_examples',
                   isinstance(examples[0], tf.train.SequenceExample))
      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 has_model_name(self):
      return 'model_name' in self.config
  
    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 has_compare_model_name(self):
      return 'model_name_2' in self.config
  
    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 need output parsing.
  
      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.
  
      If using an AI Platform model which returns multiple prediction
      results in a dictionary, this method specifies the key corresponding to
      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 example
        parsing.
  
      Returns:
        self, in order to enabled method chaining.
      """
      # If custom function is set, remove it before setting estimator
      self.delete('custom_predict_fn')
  
      self.store('estimator_and_spec', {
        'estimator': estimator, 'feature_spec': feature_spec})
      self.set_inference_address('estimator')
      # If no model name has been set, give a default
      if not self.has_model_name():
        self.set_model_name('1')
      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 example
        parsing.
  
      Returns:
        self, in order to enabled method chaining.
      """
      # If custom function is set, remove it before setting estimator
      self.delete('compare_custom_predict_fn')
  
      self.store('compare_estimator_and_spec', {
        'estimator': estimator, 'feature_spec': feature_spec})
      self.set_compare_inference_address('estimator')
      # If no model name has been set, give a default
      if not self.has_compare_model_name():
        self.set_compare_model_name('2')
      return self
  
    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.
  
      Optionally, if attributions or other prediction-time information
      can be returned by the model with each prediction, then this method
      can return a dict with the key 'predictions' containing the predictions
      result list described above, and with the key 'attributions' containing
      a list of attributions for each example that was predicted.
  
      For each example, the attributions list should contain a dict mapping
      input feature names to attribution values for that feature on that example.
      The attribution value can be one of these things:
        - A single number representing the attribution for the entire feature
        - A list of numbers representing the attribution to each value in the
          feature for multivalent features - such as attributions to individual
          pixels in an image or numbers in a list of numbers.
  
      This dict can contain any other keys, with their values being a list of
      prediction-time strings or numbers for each example being predicted. These
      values will be displayed in WIT as extra information for each example,
      usable in the same ways by WIT as normal input features (such as for
      creating plots and slicing performance data).
  
      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
  
    def set_compare_custom_predict_fn(self, predict_fn):
      """Sets a second custom function for inference.
  
      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 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.
  
      Optionally, if attributions or other prediction-time information
      can be returned by the model with each prediction, then this method
      can return a dict with the key 'predictions' containing the predictions
      result list described above, and with the key 'attributions' containing
      a list of attributions for each example that was predicted.
  
      For each example, the attributions list should contain a dict mapping
      input feature names to attribution values for that feature on that example.
      The attribution value can be one of these things:
        - A single number representing the attribution for the entire feature
        - A list of numbers representing the attribution to each value in the
          feature for multivalent features - such as attributions to individual
          pixels in an image or numbers in a list of numbers.
        - A 2D list for sparse feature attribution. Index 0 contains a list of
          feature values that there are attribution scores for. Index 1 contains
          a list of attribution values for the corresponding feature values in
          the first list.
  
      This dict can contain any other keys, with their values being a list of
      prediction-time strings or numbers for each example being predicted. These
      values will be displayed in WIT as extra information for each example,
      usable in the same ways by WIT as normal input features (such as for
      creating plots and slicing performance data).
  
      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('compare_estimator_and_spec')
  
      self.store('compare_custom_predict_fn', predict_fn)
      self.set_compare_inference_address('custom_predict_fn')
      # If no model name has been set, give a default
      if not self.has_compare_model_name():
        self.set_compare_model_name('2')
      return self
  
    def set_custom_distance_fn(self, distance_fn):
      """Sets a custom function for distance computation.
  
      WIT can directly use a custom function for all distance computations within
      the tool. In this case, the provided function should accept a query example
      proto and a list of example protos to compute the distance against and
      return a 1D list of numbers containing the distances.
  
      Args:
        distance_fn: The python function which will be used for distance
        computation.
  
      Returns:
        self, in order to enabled method chaining.
      """
      if distance_fn is None:
        self.delete('custom_distance_fn')
      else:
        self.store('custom_distance_fn', distance_fn)
      return self
  
    def set_ai_platform_model(
      self, project, model, version=None, force_json_input=None,
      adjust_prediction=None, adjust_example=None, adjust_attribution=None,
      service_name='ml', service_version='v1', get_explanations=True,
      batch_size=500, api_key=None):
      """Sets the model information for a model served by AI Platform.
  
      AI Platform Prediction a Google Cloud serving platform.
  
      Args:
        project: The name of the AI Platform Prediction project.
        model: The name of the AI Platform Prediction model.
        version: Optional, the version of the AI Platform Prediction model.
        force_json_input: Optional. If True and examples are provided as
        tf.Example protos, convert them to raw JSON objects before sending them
        for inference to this model.
        adjust_prediction: Optional. If not None then this function takes the
        prediction output from the model for a single example and converts it to
        the appopriate format - a regression score or a list of class scores. Only
        necessary if the model doesn't already abide by this format.
        adjust_example: Optional. If not None then this function takes an example
        to run prediction on and converts it to the format expected by the model.
        Necessary for example if the served model expects a single data value to
        run inference on instead of a list or dict of values.
        adjust_attribution: Optional. If not None and the model returns attribution
        information, then this function takes the attribution information for an
        example and converts it to the format expected by the tool, which is a
        dictionary of input feature names to attribution scores. Usually necessary
        if making use of adjust_example and the model returns attribution results.
        service_name: Optional. Name of the AI Platform Prediction service. Defaults
        to 'ml'.
        service_version: Optional. Version of the AI Platform Prediction service. Defaults
        to 'v1'.
        get_explanations: Optional. If a model is deployed with explanations,
        then this specifies if explainations will be calculated and displayed.
        Defaults to True.
        batch_size: Optional. Sets the individual batch size to send for
        prediction. Defaults to 500.
        api_key: Optional. A generated API key to send with the requests to AI
        Platform.
  
      Returns:
        self, in order to enabled method chaining.
      """
      self.set_inference_address(project)
      self.set_model_name(model)
      self.store('use_aip', True)
      self.store('aip_service_name', service_name)
      self.store('aip_service_version', service_version)
      self.store('aip_batch_size', batch_size)
      self.store('get_explanations', get_explanations)
      if version is not None:
        self.set_model_signature(version)
      if force_json_input:
        self.store('force_json_input', True)
      if adjust_prediction:
        self.store('adjust_prediction', adjust_prediction)
      if adjust_example:
        self.store('adjust_example', adjust_example)
      if adjust_attribution:
        self.store('adjust_attribution', adjust_attribution)
      if api_key:
        self.store('aip_api_key', api_key)
      return self
  
    def set_compare_ai_platform_model(
      self, project, model, version=None, force_json_input=None,
      adjust_prediction=None, adjust_example=None, adjust_attribution=None,
      service_name='ml', service_version='v1', get_explanations=True,
      batch_size=500, api_key=None):
      """Sets the model information for a second model served by AI Platform.
  
      AI Platform Prediction a Google Cloud serving platform.
  
      Args:
        project: The name of the AI Platform Prediction project.
        model: The name of the AI Platform Prediction model.
        version: Optional, the version of the AI Platform Prediction model.
        force_json_input: Optional. If True and examples are provided as
        tf.Example protos, convert them to raw JSON objects before sending them
        for inference to this model.
        adjust_prediction: Optional. If not None then this function takes the
        prediction output from the model for a single example and converts it to
        the appopriate format - a regression score or a list of class scores. Only
        necessary if the model doesn't already abide by this format.
        adjust_example: Optional. If not None then this function takes an example
        to run prediction on and converts it to the format expected by the model.
        Necessary for example if the served model expects a single data value to
        run inference on instead of a list or dict of values.
        adjust_attribution: Optional. If not None and the model returns attribution
        information, then this function takes the attribution information for an
        example and converts it to the format expected by the tool, which is a
        dictionary of input feature names to attribution scores. Usually necessary
        if making use of adjust_example and the model returns attribution results.
        service_name: Optional. Name of the AI Platform Prediction service. Defaults
        to 'ml'.
        service_version: Optional. Version of the AI Platform Prediction service. Defaults
        to 'v1'.
        get_explanations: Optional. If a model is deployed with explanations,
        then this specifies if explainations will be calculated and displayed.
        Defaults to True.
        batch_size: Optional. Sets the individual batch size to send for
        prediction. Defaults to 500.
        api_key: Optional. A generated API key to send with the requests to AI
        Platform.
  
      Returns:
        self, in order to enabled method chaining.
      """
      self.set_compare_inference_address(project)
      self.set_compare_model_name(model)
      self.store('compare_use_aip', True)
      self.store('compare_aip_service_name', service_name)
      self.store('compare_aip_service_version', service_version)
      self.store('compare_aip_batch_size', batch_size)
      self.store('compare_get_explanations', get_explanations)
      if version is not None:
        self.set_compare_model_signature(version)
      if force_json_input:
        self.store('compare_force_json_input', True)
      if adjust_prediction:
        self.store('compare_adjust_prediction', adjust_prediction)
      if adjust_example:
        self.store('compare_adjust_example', adjust_example)
      if adjust_attribution:
        self.store('compare_adjust_attribution', adjust_attribution)
      if api_key:
        self.store('compare_aip_api_key', api_key)
      return self
  
    def set_target_feature(self, target):
      """Sets the name of the target feature in the provided examples.
  
      If the provided examples contain a feature that represents the target
      that the model is trying to predict, it can be specified by this method.
      This is necessary for AI Platform models so that the target feature isn't
      sent to the model for prediction, which can cause model inference errors.
  
      Args:
        target: The name of the feature in the examples that represents the value
        that the model is trying to predict.
  
      Returns:
        self, in order to enabled method chaining.
      """
      self.store('target_feature', target)
      return self
  
    def _set_uses_json_input(self, is_json):
      self.store('uses_json_input', is_json)
      return self
  
    def _set_uses_json_list(self, is_list):
      self.store('uses_json_list', is_list)
      return self
  