# Lint as: python3
# Copyright 2018 Google LLC
#
# 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
#
#     https://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.
"""Test for using the MetricsPlotsAndValidationsWriter API."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import string
import tempfile

from absl.testing import parameterized
import apache_beam as beam
from apache_beam.testing import util
import numpy as np
import tensorflow as tf
from tensorflow_model_analysis import constants
from tensorflow_model_analysis import types
from tensorflow_model_analysis.api import model_eval_lib
from tensorflow_model_analysis.eval_saved_model import testutil
from tensorflow_model_analysis.eval_saved_model.example_trainers import fixed_prediction_estimator
from tensorflow_model_analysis.evaluators import legacy_metrics_and_plots_evaluator
from tensorflow_model_analysis.evaluators import metrics_plots_and_validations_evaluator
from tensorflow_model_analysis.extractors import example_weights_extractor
from tensorflow_model_analysis.extractors import features_extractor
from tensorflow_model_analysis.extractors import labels_extractor
from tensorflow_model_analysis.extractors import legacy_predict_extractor
from tensorflow_model_analysis.extractors import predictions_extractor
from tensorflow_model_analysis.extractors import slice_key_extractor
from tensorflow_model_analysis.extractors import unbatch_extractor
from tensorflow_model_analysis.metrics import attributions
from tensorflow_model_analysis.metrics import binary_confusion_matrices
from tensorflow_model_analysis.metrics import metric_types
from tensorflow_model_analysis.post_export_metrics import metric_keys
from tensorflow_model_analysis.post_export_metrics import post_export_metrics
from tensorflow_model_analysis.proto import config_pb2
from tensorflow_model_analysis.proto import metrics_for_slice_pb2
from tensorflow_model_analysis.proto import validation_result_pb2
from tensorflow_model_analysis.slicer import slicer_lib as slicer
from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer
from tfx_bsl.tfxio import raw_tf_record
from tfx_bsl.tfxio import tensor_adapter
from tfx_bsl.tfxio import test_util

from google.protobuf import text_format
from tensorflow_metadata.proto.v0 import schema_pb2


def _make_slice_key(*args):
  if len(args) % 2 != 0:
    raise ValueError('number of arguments should be even')

  result = []
  for i in range(0, len(args), 2):
    result.append((args[i], args[i + 1]))
  result = tuple(result)
  return result


class MetricsPlotsAndValidationsWriterTest(testutil.TensorflowModelAnalysisTest,
                                           parameterized.TestCase):

  def setUp(self):
    super(MetricsPlotsAndValidationsWriterTest, self).setUp()
    self.longMessage = True  # pylint: disable=invalid-name

  def _getTempDir(self):
    return tempfile.mkdtemp()

  def _getExportDir(self):
    return os.path.join(self._getTempDir(), 'export_dir')

  def _getBaselineDir(self):
    return os.path.join(self._getTempDir(), 'baseline_export_dir')

  def _build_keras_model(self, model_dir, mul):
    input_layer = tf.keras.layers.Input(shape=(1,), name='input_1')
    output_layer = tf.keras.layers.Lambda(
        lambda x, mul: x * mul, output_shape=(1,), arguments={'mul': mul})(
            input_layer)
    model = tf.keras.models.Model([input_layer], output_layer)
    model.compile(
        optimizer=tf.keras.optimizers.Adam(lr=.001),
        loss=tf.keras.losses.BinaryCrossentropy(),
        metrics=['accuracy'])

    model.fit(x=[[0], [1]], y=[[0], [1]], steps_per_epoch=1)
    model.save(model_dir, save_format='tf')
    return self.createTestEvalSharedModel(
        eval_saved_model_path=model_dir, tags=[tf.saved_model.SERVING])

  def testConvertSlicePlotsToProto(self):
    slice_key = _make_slice_key('fruit', 'apple')
    plot_key = metric_types.PlotKey(
        name='calibration_plot', output_name='output_name')
    calibration_plot = text_format.Parse(
        """
        buckets {
          lower_threshold_inclusive: -inf
          upper_threshold_exclusive: 0.0
          num_weighted_examples { value: 0.0 }
          total_weighted_label { value: 0.0 }
          total_weighted_refined_prediction { value: 0.0 }
        }
        buckets {
          lower_threshold_inclusive: 0.0
          upper_threshold_exclusive: 0.5
          num_weighted_examples { value: 1.0 }
          total_weighted_label { value: 1.0 }
          total_weighted_refined_prediction { value: 0.3 }
        }
        buckets {
          lower_threshold_inclusive: 0.5
          upper_threshold_exclusive: 1.0
          num_weighted_examples { value: 1.0 }
          total_weighted_label { value: 0.0 }
          total_weighted_refined_prediction { value: 0.7 }
        }
        buckets {
          lower_threshold_inclusive: 1.0
          upper_threshold_exclusive: inf
          num_weighted_examples { value: 0.0 }
          total_weighted_label { value: 0.0 }
          total_weighted_refined_prediction { value: 0.0 }
        }
     """, metrics_for_slice_pb2.CalibrationHistogramBuckets())

    expected_plots_for_slice = text_format.Parse(
        """
      slice_key {
        single_slice_keys {
          column: 'fruit'
          bytes_value: 'apple'
        }
      }
      plot_keys_and_values {
        key {
          output_name: "output_name"
        }
        value {
          calibration_histogram_buckets {
            buckets {
              lower_threshold_inclusive: -inf
              upper_threshold_exclusive: 0.0
              num_weighted_examples { value: 0.0 }
              total_weighted_label { value: 0.0 }
              total_weighted_refined_prediction { value: 0.0 }
            }
            buckets {
              lower_threshold_inclusive: 0.0
              upper_threshold_exclusive: 0.5
              num_weighted_examples { value: 1.0 }
              total_weighted_label { value: 1.0 }
              total_weighted_refined_prediction { value: 0.3 }
            }
            buckets {
              lower_threshold_inclusive: 0.5
              upper_threshold_exclusive: 1.0
              num_weighted_examples { value: 1.0 }
              total_weighted_label { value: 0.0 }
              total_weighted_refined_prediction { value: 0.7 }
            }
            buckets {
              lower_threshold_inclusive: 1.0
              upper_threshold_exclusive: inf
              num_weighted_examples { value: 0.0 }
              total_weighted_label { value: 0.0 }
              total_weighted_refined_prediction { value: 0.0 }
            }
          }
        }
      }
    """, metrics_for_slice_pb2.PlotsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_plots_to_proto(
        (slice_key, {
            plot_key: calibration_plot
        }), None)
    self.assertProtoEquals(expected_plots_for_slice, got)

  def testConvertSlicePlotsToProtoLegacyStringKeys(self):
    slice_key = _make_slice_key('fruit', 'apple')
    tfma_plots = {
        metric_keys.CALIBRATION_PLOT_MATRICES:
            np.array([
                [0.0, 0.0, 0.0],
                [0.3, 1.0, 1.0],
                [0.7, 0.0, 1.0],
                [0.0, 0.0, 0.0],
            ]),
        metric_keys.CALIBRATION_PLOT_BOUNDARIES:
            np.array([0.0, 0.5, 1.0]),
    }
    expected_plot_data = """
      slice_key {
        single_slice_keys {
          column: 'fruit'
          bytes_value: 'apple'
        }
      }
      plots {
        key: "post_export_metrics"
        value {
          calibration_histogram_buckets {
            buckets {
              lower_threshold_inclusive: -inf
              upper_threshold_exclusive: 0.0
              num_weighted_examples { value: 0.0 }
              total_weighted_label { value: 0.0 }
              total_weighted_refined_prediction { value: 0.0 }
            }
            buckets {
              lower_threshold_inclusive: 0.0
              upper_threshold_exclusive: 0.5
              num_weighted_examples { value: 1.0 }
              total_weighted_label { value: 1.0 }
              total_weighted_refined_prediction { value: 0.3 }
            }
            buckets {
              lower_threshold_inclusive: 0.5
              upper_threshold_exclusive: 1.0
              num_weighted_examples { value: 1.0 }
              total_weighted_label { value: 0.0 }
              total_weighted_refined_prediction { value: 0.7 }
            }
            buckets {
              lower_threshold_inclusive: 1.0
              upper_threshold_exclusive: inf
              num_weighted_examples { value: 0.0 }
              total_weighted_label { value: 0.0 }
              total_weighted_refined_prediction { value: 0.0 }
            }
          }
        }
      }
    """
    calibration_plot = (
        post_export_metrics.calibration_plot_and_prediction_histogram())
    got = metrics_plots_and_validations_writer.convert_slice_plots_to_proto(
        (slice_key, tfma_plots), [calibration_plot])
    self.assertProtoEquals(expected_plot_data, got)

  def testConvertSlicePlotsToProtoEmptyPlot(self):
    slice_key = _make_slice_key('fruit', 'apple')
    tfma_plots = {metric_keys.ERROR_METRIC: 'error_message'}

    actual_plot = metrics_plots_and_validations_writer.convert_slice_plots_to_proto(
        (slice_key, tfma_plots), [])
    expected_plot = metrics_for_slice_pb2.PlotsForSlice()
    expected_plot.slice_key.CopyFrom(slicer.serialize_slice_key(slice_key))
    expected_plot.plots[
        metric_keys.ERROR_METRIC].debug_message = 'error_message'
    self.assertProtoEquals(expected_plot, actual_plot)

  def testConvertSliceMetricsToProto(self):
    slice_key = _make_slice_key('age', 5, 'language', 'english', 'price', 0.3)
    slice_metrics = {
        metric_types.MetricKey(name='accuracy', output_name='output_name'): 0.8
    }
    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {
          single_slice_keys {
            column: 'age'
            int64_value: 5
          }
          single_slice_keys {
            column: 'language'
            bytes_value: 'english'
          }
          single_slice_keys {
            column: 'price'
            float_value: 0.3
          }
        }
        metric_keys_and_values {
          key {
            name: "accuracy"
            output_name: "output_name"
          }
          value {
            double_value {
              value: 0.8
            }
          }
        }""", metrics_for_slice_pb2.MetricsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics), None)
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceMetricsToProtoConfusionMatrices(self):
    slice_key = _make_slice_key()
    slice_metrics = {
        metric_types.MetricKey(name='confusion_matrix_at_thresholds'):
            binary_confusion_matrices.Matrices(
                thresholds=[0.25, 0.75, 1.00],
                fn=[0.0, 1.0, 2.0],
                tn=[1.0, 1.0, 1.0],
                fp=[0.0, 0.0, 0.0],
                tp=[2.0, 1.0, 0.0])
    }
    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {}
        metric_keys_and_values {
          key: { name: "confusion_matrix_at_thresholds" }
          value {
            confusion_matrix_at_thresholds {
              matrices {
                threshold: 0.25
                false_negatives: 0.0
                true_negatives: 1.0
                false_positives: 0.0
                true_positives: 2.0
                precision: 1.0
                recall: 1.0
              }
              matrices {
                threshold: 0.75
                false_negatives: 1.0
                true_negatives: 1.0
                false_positives: 0.0
                true_positives: 1.0
                precision: 1.0
                recall: 0.5
              }
              matrices {
                threshold: 1.00
                false_negatives: 2.0
                true_negatives: 1.0
                false_positives: 0.0
                true_positives: 0.0
                precision: 1.0
                recall: 0.0
              }
            }
          }
        }
        """, metrics_for_slice_pb2.MetricsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics), add_metrics_callbacks=[])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceMetricsToProtoConfusionMatricesPostExport(self):
    slice_key = _make_slice_key()

    thresholds = [0.25, 0.75, 1.00]
    matrices = [[0.0, 1.0, 0.0, 2.0, 1.0, 1.0], [1.0, 1.0, 0.0, 1.0, 1.0, 0.5],
                [2.0, 1.0, 0.0, 0.0, float('nan'), 0.0]]

    slice_metrics = {
        metric_keys.CONFUSION_MATRIX_AT_THRESHOLDS_MATRICES: matrices,
        metric_keys.CONFUSION_MATRIX_AT_THRESHOLDS_THRESHOLDS: thresholds,
    }
    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {}
        metrics {
          key: "post_export_metrics/confusion_matrix_at_thresholds"
          value {
            confusion_matrix_at_thresholds {
              matrices {
                threshold: 0.25
                false_negatives: 0.0
                true_negatives: 1.0
                false_positives: 0.0
                true_positives: 2.0
                precision: 1.0
                recall: 1.0
                bounded_false_negatives {
                  value {
                    value: 0.0
                  }
                }
                bounded_true_negatives {
                  value {
                    value: 1.0
                  }
                }
                bounded_true_positives {
                  value {
                    value: 2.0
                  }
                }
                bounded_false_positives {
                  value {
                    value: 0.0
                  }
                }
                bounded_precision {
                  value {
                    value: 1.0
                  }
                }
                bounded_recall {
                  value {
                    value: 1.0
                  }
                }
                t_distribution_false_negatives {
                  unsampled_value {
                    value: 0.0
                  }
                }
                t_distribution_true_negatives {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_true_positives {
                  unsampled_value {
                    value: 2.0
                  }
                }
                t_distribution_false_positives {
                  unsampled_value {
                    value: 0.0
                  }
                }
                t_distribution_precision {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_recall {
                  unsampled_value {
                    value: 1.0
                  }
                }
              }
              matrices {
                threshold: 0.75
                false_negatives: 1.0
                true_negatives: 1.0
                false_positives: 0.0
                true_positives: 1.0
                precision: 1.0
                recall: 0.5
                bounded_false_negatives {
                  value {
                    value: 1.0
                  }
                }
                bounded_true_negatives {
                  value {
                    value: 1.0
                  }
                }
                bounded_true_positives {
                  value {
                    value: 1.0
                  }
                }
                bounded_false_positives {
                  value {
                    value: 0.0
                  }
                }
                bounded_precision {
                  value {
                    value: 1.0
                  }
                }
                bounded_recall {
                  value {
                    value: 0.5
                  }
                }
                t_distribution_false_negatives {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_true_negatives {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_true_positives {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_false_positives {
                  unsampled_value {
                    value: 0.0
                  }
                }
                t_distribution_precision {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_recall {
                  unsampled_value {
                    value: 0.5
                  }
                }
              }
              matrices {
                threshold: 1.00
                false_negatives: 2.0
                true_negatives: 1.0
                false_positives: 0.0
                true_positives: 0.0
                precision: nan
                recall: 0.0
                bounded_false_negatives {
                  value {
                    value: 2.0
                  }
                }
                bounded_true_negatives {
                  value {
                    value: 1.0
                  }
                }
                bounded_true_positives {
                  value {
                    value: 0.0
                  }
                }
                bounded_false_positives {
                  value {
                    value: 0.0
                  }
                }
                bounded_precision {
                  value {
                    value: nan
                  }
                }
                bounded_recall {
                  value {
                    value: 0.0
                  }
                }
                t_distribution_false_negatives {
                  unsampled_value {
                    value: 2.0
                  }
                }
                t_distribution_true_negatives {
                  unsampled_value {
                    value: 1.0
                  }
                }
                t_distribution_true_positives {
                  unsampled_value {
                    value: 0.0
                  }
                }
                t_distribution_false_positives {
                  unsampled_value {
                    value: 0.0
                  }
                }
                t_distribution_precision {
                  unsampled_value {
                    value: nan
                  }
                }
                t_distribution_recall {
                  unsampled_value {
                    value: 0.0
                  }
                }
              }
            }
          }
        }
        """, metrics_for_slice_pb2.MetricsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics),
        [post_export_metrics.confusion_matrix_at_thresholds(thresholds)])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceMetricsToProtoMetricsRanges(self):
    slice_key = _make_slice_key('age', 5, 'language', 'english', 'price', 0.3)
    slice_metrics = {
        'accuracy': types.ValueWithTDistribution(0.8, 0.1, 9, 0.8),
        metric_keys.AUPRC: 0.1,
        metric_keys.lower_bound_key(metric_keys.AUPRC): 0.05,
        metric_keys.upper_bound_key(metric_keys.AUPRC): 0.17,
        metric_keys.AUC: 0.2,
        metric_keys.lower_bound_key(metric_keys.AUC): 0.1,
        metric_keys.upper_bound_key(metric_keys.AUC): 0.3
    }
    expected_metrics_for_slice = text_format.Parse(
        string.Template("""
        slice_key {
          single_slice_keys {
            column: 'age'
            int64_value: 5
          }
          single_slice_keys {
            column: 'language'
            bytes_value: 'english'
          }
          single_slice_keys {
            column: 'price'
            float_value: 0.3
          }
        }
        metrics {
          key: "accuracy"
          value {
            bounded_value {
              value {
                value: 0.8
              }
              lower_bound { value: 0.5737843 }
              upper_bound { value: 1.0262157 }
              methodology: POISSON_BOOTSTRAP
            }
          }
        }
        metrics {
          key: "$auc"
          value {
            bounded_value {
              lower_bound {
                value: 0.1
              }
              upper_bound {
                value: 0.3
              }
              value {
                value: 0.2
              }
              methodology: RIEMANN_SUM
            }
          }
        }
        metrics {
          key: "$auprc"
          value {
            bounded_value {
              lower_bound {
                value: 0.05
              }
              upper_bound {
                value: 0.17
              }
              value {
                value: 0.1
              }
              methodology: RIEMANN_SUM
            }
          }
        }""").substitute(auc=metric_keys.AUC, auprc=metric_keys.AUPRC),
        metrics_for_slice_pb2.MetricsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics),
        [post_export_metrics.auc(),
         post_export_metrics.auc(curve='PR')])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceMetricsToProtoFromLegacyStrings(self):
    slice_key = _make_slice_key('age', 5, 'language', 'english', 'price', 0.3)
    slice_metrics = {
        'accuracy': 0.8,
        metric_keys.AUPRC: 0.1,
        metric_keys.lower_bound_key(metric_keys.AUPRC): 0.05,
        metric_keys.upper_bound_key(metric_keys.AUPRC): 0.17,
        metric_keys.AUC: 0.2,
        metric_keys.lower_bound_key(metric_keys.AUC): 0.1,
        metric_keys.upper_bound_key(metric_keys.AUC): 0.3
    }
    expected_metrics_for_slice = text_format.Parse(
        string.Template("""
        slice_key {
          single_slice_keys {
            column: 'age'
            int64_value: 5
          }
          single_slice_keys {
            column: 'language'
            bytes_value: 'english'
          }
          single_slice_keys {
            column: 'price'
            float_value: 0.3
          }
        }
        metrics {
          key: "accuracy"
          value {
            double_value {
              value: 0.8
            }
          }
        }
        metrics {
          key: "$auc"
          value {
            bounded_value {
              lower_bound {
                value: 0.1
              }
              upper_bound {
                value: 0.3
              }
              value {
                value: 0.2
              }
              methodology: RIEMANN_SUM
            }
          }
        }
        metrics {
          key: "$auprc"
          value {
            bounded_value {
              lower_bound {
                value: 0.05
              }
              upper_bound {
                value: 0.17
              }
              value {
                value: 0.1
              }
              methodology: RIEMANN_SUM
            }
          }
        }""").substitute(auc=metric_keys.AUC, auprc=metric_keys.AUPRC),
        metrics_for_slice_pb2.MetricsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics),
        [post_export_metrics.auc(),
         post_export_metrics.auc(curve='PR')])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceMetricsToProtoEmptyMetrics(self):
    slice_key = _make_slice_key('age', 5, 'language', 'english', 'price', 0.3)
    slice_metrics = {metric_keys.ERROR_METRIC: 'error_message'}

    actual_metrics = (
        metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
            (slice_key, slice_metrics),
            [post_export_metrics.auc(),
             post_export_metrics.auc(curve='PR')]))

    expected_metrics = metrics_for_slice_pb2.MetricsForSlice()
    expected_metrics.slice_key.CopyFrom(slicer.serialize_slice_key(slice_key))
    expected_metrics.metrics[
        metric_keys.ERROR_METRIC].debug_message = 'error_message'
    self.assertProtoEquals(expected_metrics, actual_metrics)

  def testConvertSliceMetricsToProtoStringMetrics(self):
    slice_key = _make_slice_key()
    slice_metrics = {
        'valid_ascii': b'test string',
        'valid_unicode': b'\xF0\x9F\x90\x84',  # U+1F404, Cow
        'invalid_unicode': b'\xE2\x28\xA1',
    }
    expected_metrics_for_slice = metrics_for_slice_pb2.MetricsForSlice()
    expected_metrics_for_slice.slice_key.SetInParent()
    expected_metrics_for_slice.metrics[
        'valid_ascii'].bytes_value = slice_metrics['valid_ascii']
    expected_metrics_for_slice.metrics[
        'valid_unicode'].bytes_value = slice_metrics['valid_unicode']
    expected_metrics_for_slice.metrics[
        'invalid_unicode'].bytes_value = slice_metrics['invalid_unicode']

    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics), [])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testCombineValidationsValidationOk(self):
    input_validations = [
        text_format.Parse(
            """
            validation_ok: true
            metric_validations_per_slice {
              slice_key  {
                single_slice_keys {
                  column: "x"
                  bytes_value: "x1"
                }
              }
            }
            validation_details {
              slicing_details {
                slicing_spec {}
                num_matching_slices: 1
              }
              slicing_details {
                slicing_spec {
                  feature_keys: ["x", "y"]
                }
                num_matching_slices: 1
              }
            }""", validation_result_pb2.ValidationResult()),
        text_format.Parse(
            """
            validation_ok: true
            validation_details {
              slicing_details {
                slicing_spec {
                  feature_keys: ["x"]
                }
                num_matching_slices: 1
              }
              slicing_details {
                slicing_spec {
                  feature_keys: ["x", "y"]
                }
                num_matching_slices: 2
              }
            }""", validation_result_pb2.ValidationResult())
    ]

    eval_config = config_pb2.EvalConfig(
        model_specs=[
            config_pb2.ModelSpec(name='candidate'),
            config_pb2.ModelSpec(name='baseline', is_baseline=True)
        ],
        slicing_specs=[config_pb2.SlicingSpec()],
        metrics_specs=[
            config_pb2.MetricsSpec(
                metrics=[
                    config_pb2.MetricConfig(
                        class_name='AUC',
                        per_slice_thresholds=[
                            config_pb2.PerSliceMetricThreshold(
                                slicing_specs=[config_pb2.SlicingSpec()],
                                threshold=config_pb2.MetricThreshold(
                                    value_threshold=config_pb2
                                    .GenericValueThreshold(
                                        lower_bound={'value': 0.7})))
                        ]),
                ],
                model_names=['candidate', 'baseline']),
        ])

    expected_validation = text_format.Parse(
        """
        validation_ok: true
        metric_validations_per_slice {
          slice_key  {
            single_slice_keys {
              column: "x"
              bytes_value: "x1"
            }
          }
        }
        validation_details {
          slicing_details {
            slicing_spec {}
            num_matching_slices: 1
          }
          slicing_details {
            slicing_spec {
              feature_keys: ["x", "y"]
            }
            num_matching_slices: 3
          }
          slicing_details {
            slicing_spec {
              feature_keys: ["x"]
            }
            num_matching_slices: 1
          }
        }""", validation_result_pb2.ValidationResult())

    def verify_fn(result):
      self.assertLen(result, 1)
      self.assertProtoEquals(expected_validation, result[0])

    with beam.Pipeline() as pipeline:
      result = (
          pipeline
          | 'Create' >> beam.Create(input_validations)
          | 'CombineValidations' >> beam.CombineGlobally(
              metrics_plots_and_validations_writer.CombineValidations(
                  eval_config)))
      util.assert_that(result, verify_fn)

  def testCombineValidationsMissingSlices(self):
    input_validations = [
        text_format.Parse(
            """
            validation_ok: false
            metric_validations_per_slice {
              slice_key  {
                single_slice_keys {
                  column: "x"
                  bytes_value: "x1"
                }
              }
              failures {
                metric_key {
                  name: "auc"
                  model_name: "candidate"
                  is_diff: true
                }
                metric_threshold {
                  value_threshold {
                    lower_bound { value: 0.7 }
                  }
                }
                metric_value {
                  double_value { value: 0.6 }
                }
              }
            }
            validation_details {
              slicing_details {
                slicing_spec {}
                num_matching_slices: 1
              }
              slicing_details {
                slicing_spec {
                  feature_keys: ["x", "y"]
                }
                num_matching_slices: 1
              }
            }""", validation_result_pb2.ValidationResult()),
        text_format.Parse(
            """
            validation_ok: true
            validation_details {
              slicing_details {
                slicing_spec {
                  feature_keys: ["x"]
                }
                num_matching_slices: 1
              }
              slicing_details {
                slicing_spec {
                  feature_keys: ["x", "y"]
                }
                num_matching_slices: 2
              }
            }""", validation_result_pb2.ValidationResult())
    ]

    slicing_specs = [
        config_pb2.SlicingSpec(),
        config_pb2.SlicingSpec(feature_keys=['x']),
        config_pb2.SlicingSpec(feature_keys=['x', 'y']),
        config_pb2.SlicingSpec(feature_keys=['z']),
    ]
    eval_config = config_pb2.EvalConfig(
        model_specs=[
            config_pb2.ModelSpec(name='candidate'),
            config_pb2.ModelSpec(name='baseline', is_baseline=True)
        ],
        slicing_specs=slicing_specs,
        metrics_specs=[
            config_pb2.MetricsSpec(
                metrics=[
                    config_pb2.MetricConfig(
                        class_name='AUC',
                        per_slice_thresholds=[
                            config_pb2.PerSliceMetricThreshold(
                                slicing_specs=slicing_specs,
                                threshold=config_pb2.MetricThreshold(
                                    value_threshold=config_pb2
                                    .GenericValueThreshold(
                                        lower_bound={'value': 0.7})))
                        ]),
                ],
                model_names=['candidate', 'baseline']),
        ])

    expected_validation = text_format.Parse(
        """
        validation_ok: false
        metric_validations_per_slice {
          slice_key  {
            single_slice_keys {
              column: "x"
              bytes_value: "x1"
            }
          }
          failures {
            metric_key {
              name: "auc"
              model_name: "candidate"
              is_diff: true
            }
            metric_threshold {
              value_threshold {
                lower_bound { value: 0.7 }
              }
            }
            metric_value {
              double_value { value: 0.6 }
            }
          }
        }
        missing_slices {
          feature_keys: "z"
        }
        validation_details {
          slicing_details {
            slicing_spec {}
            num_matching_slices: 1
          }
          slicing_details {
            slicing_spec {
              feature_keys: ["x", "y"]
            }
            num_matching_slices: 3
          }
          slicing_details {
            slicing_spec {
              feature_keys: ["x"]
            }
            num_matching_slices: 1
          }
        }""", validation_result_pb2.ValidationResult())

    def verify_fn(result):
      self.assertLen(result, 1)
      self.assertProtoEquals(expected_validation, result[0])

    with beam.Pipeline() as pipeline:
      result = (
          pipeline
          | 'Create' >> beam.Create(input_validations)
          | 'CombineValidations' >> beam.CombineGlobally(
              metrics_plots_and_validations_writer.CombineValidations(
                  eval_config)))
      util.assert_that(result, verify_fn)

  def testUncertaintyValuedMetrics(self):
    slice_key = _make_slice_key()
    slice_metrics = {
        'one_dim':
            types.ValueWithTDistribution(2.0, 1.0, 3, 2.0),
        'nans':
            types.ValueWithTDistribution(
                float('nan'), float('nan'), -1, float('nan')),
    }
    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {}
        metrics {
          key: "one_dim"
          value {
            bounded_value {
              value {
                value: 2.0
              }
              lower_bound {
                value: -1.1824463
              }
              upper_bound {
                value: 5.1824463
              }
              methodology: POISSON_BOOTSTRAP
            }
          }
        }
        metrics {
          key: "nans"
          value {
            bounded_value {
              value {
                value: nan
              }
              lower_bound {
                value: nan
              }
              upper_bound {
                value: nan
              }
              methodology: POISSON_BOOTSTRAP
            }
          }
        }
        """, metrics_for_slice_pb2.MetricsForSlice())
    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics), [])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceMetricsToProtoTensorValuedMetrics(self):
    slice_key = _make_slice_key()
    slice_metrics = {
        'one_dim':
            np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
        'two_dims':
            np.array([['two', 'dims', 'test'], ['TWO', 'DIMS', 'TEST']]),
        'three_dims':
            np.array([[[100, 200, 300]], [[500, 600, 700]]], dtype=np.int64),
    }
    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {}
        metrics {
          key: "one_dim"
          value {
            array_value {
              data_type: FLOAT32
              shape: 4
              float32_values: [1.0, 2.0, 3.0, 4.0]
            }
          }
        }
        metrics {
          key: "two_dims"
          value {
            array_value {
              data_type: BYTES
              shape: [2, 3]
              bytes_values: ["two", "dims", "test", "TWO", "DIMS", "TEST"]
            }
          }
        }
        metrics {
          key: "three_dims"
          value {
            array_value {
              data_type: INT64
              shape: [2, 1, 3]
              int64_values: [100, 200, 300, 500, 600, 700]
            }
          }
        }
        """, metrics_for_slice_pb2.MetricsForSlice())
    got = metrics_plots_and_validations_writer.convert_slice_metrics_to_proto(
        (slice_key, slice_metrics), [])
    self.assertProtoEquals(expected_metrics_for_slice, got)

  def testConvertSliceAttributionsToProto(self):
    slice_key = _make_slice_key('language', 'english', 'price', 0.3)
    slice_attributions = {
        metric_types.AttributionsKey(name='mean', output_name='output_name'): {
            'age': 0.8,
            'language': 1.2,
            'price': 2.3,
        },
    }
    expected_attributions_for_slice = text_format.Parse(
        """
        slice_key {
          single_slice_keys {
            column: 'language'
            bytes_value: 'english'
          }
          single_slice_keys {
            column: 'price'
            float_value: 0.3
          }
        }
        attributions_keys_and_values {
          key {
            name: "mean"
            output_name: "output_name"
          }
          values {
            key: "age"
            value: {
              double_value {
                value: 0.8
              }
            }
          }
          values {
            key: "language"
            value: {
              double_value {
                value: 1.2
              }
            }
          }
          values {
            key: "price"
            value: {
              double_value {
                value: 2.3
              }
            }
          }
        }""", metrics_for_slice_pb2.AttributionsForSlice())

    got = metrics_plots_and_validations_writer.convert_slice_attributions_to_proto(
        (slice_key, slice_attributions))
    self.assertProtoEquals(expected_attributions_for_slice, got)

  _OUTPUT_FORMAT_PARAMS = [('without_output_file_format', ''),
                           ('tfrecord_file_format', 'tfrecord'),
                           ('parquet_file_format', 'parquet')]

  @parameterized.named_parameters(_OUTPUT_FORMAT_PARAMS)
  def testWriteValidationResults(self, output_file_format):
    model_dir, baseline_dir = self._getExportDir(), self._getBaselineDir()
    eval_shared_model = self._build_keras_model(model_dir, mul=0)
    baseline_eval_shared_model = self._build_keras_model(baseline_dir, mul=1)
    validations_file = os.path.join(self._getTempDir(),
                                    constants.VALIDATIONS_KEY)
    schema = text_format.Parse(
        """
        tensor_representation_group {
          key: ""
          value {
            tensor_representation {
              key: "input_1"
              value {
                dense_tensor {
                  column_name: "input_1"
                  shape { dim { size: 1 } }
                }
              }
            }
          }
        }
        feature {
          name: "input_1"
          type: FLOAT
        }
        feature {
          name: "label"
          type: FLOAT
        }
        feature {
          name: "example_weight"
          type: FLOAT
        }
        feature {
          name: "extra_feature"
          type: BYTES
        }
        """, schema_pb2.Schema())
    tfx_io = test_util.InMemoryTFExampleRecord(
        schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN)
    tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
        arrow_schema=tfx_io.ArrowSchema(),
        tensor_representations=tfx_io.TensorRepresentations())
    examples = [
        self._makeExample(
            input_1=0.0,
            label=1.0,
            example_weight=1.0,
            extra_feature='non_model_feature'),
        self._makeExample(
            input_1=1.0,
            label=0.0,
            example_weight=0.5,
            extra_feature='non_model_feature'),
    ]

    slicing_specs = [
        config_pb2.SlicingSpec(),
        config_pb2.SlicingSpec(feature_keys=['slice_does_not_exist'])
    ]
    cross_slicing_specs = [
        config_pb2.CrossSlicingSpec(
            baseline_spec=config_pb2.SlicingSpec(
                feature_keys=['slice_does_not_exist']),
            slicing_specs=[
                config_pb2.SlicingSpec(feature_keys=['slice_does_not_exist'])
            ])
    ]
    eval_config = config_pb2.EvalConfig(
        model_specs=[
            config_pb2.ModelSpec(
                name='candidate',
                label_key='label',
                example_weight_key='example_weight'),
            config_pb2.ModelSpec(
                name='baseline',
                label_key='label',
                example_weight_key='example_weight',
                is_baseline=True)
        ],
        slicing_specs=slicing_specs,
        cross_slicing_specs=cross_slicing_specs,
        metrics_specs=[
            config_pb2.MetricsSpec(
                metrics=[
                    config_pb2.MetricConfig(
                        class_name='WeightedExampleCount',
                        per_slice_thresholds=[
                            config_pb2.PerSliceMetricThreshold(
                                slicing_specs=slicing_specs,
                                # 1.5 < 1, NOT OK.
                                threshold=config_pb2.MetricThreshold(
                                    value_threshold=config_pb2
                                    .GenericValueThreshold(
                                        upper_bound={'value': 1})))
                        ],
                        # missing cross slice
                        cross_slice_thresholds=[
                            config_pb2.CrossSliceMetricThreshold(
                                cross_slicing_specs=cross_slicing_specs,
                                threshold=config_pb2.MetricThreshold(
                                    value_threshold=config_pb2
                                    .GenericValueThreshold(
                                        upper_bound={'value': 1})))
                        ]),
                    config_pb2.MetricConfig(
                        class_name='ExampleCount',
                        # 2 > 10, NOT OK.
                        threshold=config_pb2.MetricThreshold(
                            value_threshold=config_pb2.GenericValueThreshold(
                                lower_bound={'value': 10}))),
                    config_pb2.MetricConfig(
                        class_name='MeanLabel',
                        # 0.5 > 1 and 0.5 > 1?: NOT OK.
                        threshold=config_pb2.MetricThreshold(
                            change_threshold=config_pb2.GenericChangeThreshold(
                                direction=config_pb2.MetricDirection
                                .HIGHER_IS_BETTER,
                                relative={'value': 1},
                                absolute={'value': 1}))),
                    config_pb2.MetricConfig(
                        # MeanPrediction = (0+0)/(1+0.5) = 0
                        class_name='MeanPrediction',
                        # -.01 < 0 < .01, OK.
                        # Diff% = -.333/.333 = -100% < -99%, OK.
                        # Diff = 0 - .333 = -.333 < 0, OK.
                        threshold=config_pb2.MetricThreshold(
                            value_threshold=config_pb2.GenericValueThreshold(
                                upper_bound={'value': .01},
                                lower_bound={'value': -.01}),
                            change_threshold=config_pb2.GenericChangeThreshold(
                                direction=config_pb2.MetricDirection
                                .LOWER_IS_BETTER,
                                relative={'value': -.99},
                                absolute={'value': 0})))
                ],
                model_names=['candidate', 'baseline']),
        ],
        options=config_pb2.Options(
            disabled_outputs={'values': ['eval_config_pb2.json']}),
    )
    slice_spec = [
        slicer.SingleSliceSpec(spec=s) for s in eval_config.slicing_specs
    ]
    eval_shared_models = {
        'candidate': eval_shared_model,
        'baseline': baseline_eval_shared_model
    }
    extractors = [
        features_extractor.FeaturesExtractor(eval_config),
        labels_extractor.LabelsExtractor(eval_config),
        example_weights_extractor.ExampleWeightsExtractor(eval_config),
        predictions_extractor.PredictionsExtractor(
            eval_shared_model=eval_shared_models,
            eval_config=eval_config,
            tensor_adapter_config=tensor_adapter_config),
        unbatch_extractor.UnbatchExtractor(),
        slice_key_extractor.SliceKeyExtractor(slice_spec=slice_spec)
    ]
    evaluators = [
        metrics_plots_and_validations_evaluator
        .MetricsPlotsAndValidationsEvaluator(
            eval_config=eval_config, eval_shared_model=eval_shared_models)
    ]
    output_paths = {
        constants.VALIDATIONS_KEY: validations_file,
    }
    writers = [
        metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(
            output_paths,
            eval_config=eval_config,
            add_metrics_callbacks=[],
            output_file_format=output_file_format)
    ]

    with beam.Pipeline() as pipeline:
      # pylint: disable=no-value-for-parameter
      _ = (
          pipeline
          | 'Create' >> beam.Create([e.SerializeToString() for e in examples])
          | 'BatchExamples' >> tfx_io.BeamSource()
          | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
          | 'ExtractEvaluate' >> model_eval_lib.ExtractAndEvaluate(
              extractors=extractors, evaluators=evaluators)
          | 'WriteResults' >> model_eval_lib.WriteResults(writers=writers))
      # pylint: enable=no-value-for-parameter

    validation_result = (
        metrics_plots_and_validations_writer
        .load_and_deserialize_validation_result(
            os.path.dirname(validations_file), output_file_format))

    expected_validations = [
        text_format.Parse(
            """
            metric_key {
              name: "weighted_example_count"
              model_name: "candidate"
            }
            metric_threshold {
              value_threshold {
                upper_bound {
                  value: 1.0
                }
              }
            }
            metric_value {
              double_value {
                value: 1.5
              }
            }
            """, validation_result_pb2.ValidationFailure()),
        text_format.Parse(
            """
            metric_key {
              name: "example_count"
              model_name: "candidate"
            }
            metric_threshold {
              value_threshold {
                lower_bound {
                  value: 10.0
                }
              }
            }
            metric_value {
              double_value {
                value: 2.0
              }
            }
            """, validation_result_pb2.ValidationFailure()),
        text_format.Parse(
            """
            metric_key {
              name: "mean_label"
              model_name: "candidate"
              is_diff: true
            }
            metric_threshold {
              change_threshold {
                absolute {
                  value: 1.0
                }
                relative {
                  value: 1.0
                }
                direction: HIGHER_IS_BETTER
              }
            }
            metric_value {
              double_value {
                value: 0.0
              }
            }
            """, validation_result_pb2.ValidationFailure()),
    ]
    self.assertFalse(validation_result.validation_ok)
    self.assertFalse(validation_result.missing_thresholds)
    self.assertLen(validation_result.metric_validations_per_slice, 1)
    self.assertCountEqual(
        expected_validations,
        validation_result.metric_validations_per_slice[0].failures)

    expected_missing_slices = [
        config_pb2.SlicingSpec(feature_keys=['slice_does_not_exist'])
    ]
    self.assertLen(validation_result.missing_slices, 1)
    self.assertCountEqual(expected_missing_slices,
                          validation_result.missing_slices)
    expected_missing_cross_slices = [
        config_pb2.CrossSlicingSpec(
            baseline_spec=config_pb2.SlicingSpec(
                feature_keys=['slice_does_not_exist']),
            slicing_specs=[
                config_pb2.SlicingSpec(feature_keys=['slice_does_not_exist'])
            ])
    ]
    self.assertLen(validation_result.missing_cross_slices, 1)
    self.assertCountEqual(expected_missing_cross_slices,
                          validation_result.missing_cross_slices)

    expected_slicing_details = [
        text_format.Parse(
            """
            slicing_spec {
            }
            num_matching_slices: 1
            """, validation_result_pb2.SlicingDetails()),
    ]
    self.assertLen(validation_result.validation_details.slicing_details, 1)
    self.assertCountEqual(expected_slicing_details,
                          validation_result.validation_details.slicing_details)

  @parameterized.named_parameters(_OUTPUT_FORMAT_PARAMS)
  def testWriteValidationResultsNoThresholds(self, output_file_format):
    model_dir, baseline_dir = self._getExportDir(), self._getBaselineDir()
    eval_shared_model = self._build_keras_model(model_dir, mul=0)
    baseline_eval_shared_model = self._build_keras_model(baseline_dir, mul=1)
    validations_file = os.path.join(self._getTempDir(),
                                    constants.VALIDATIONS_KEY)
    schema = text_format.Parse(
        """
        tensor_representation_group {
          key: ""
          value {
            tensor_representation {
              key: "input_1"
              value {
                dense_tensor {
                  column_name: "input_1"
                  shape { dim { size: 1 } }
                }
              }
            }
          }
        }
        feature {
          name: "input_1"
          type: FLOAT
        }
        feature {
          name: "label"
          type: FLOAT
        }
        feature {
          name: "example_weight"
          type: FLOAT
        }
        feature {
          name: "extra_feature"
          type: BYTES
        }
        """, schema_pb2.Schema())
    tfx_io = test_util.InMemoryTFExampleRecord(
        schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN)
    tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
        arrow_schema=tfx_io.ArrowSchema(),
        tensor_representations=tfx_io.TensorRepresentations())
    examples = [
        self._makeExample(
            input_1=0.0,
            label=1.0,
            example_weight=1.0,
            extra_feature='non_model_feature'),
        self._makeExample(
            input_1=1.0,
            label=0.0,
            example_weight=0.5,
            extra_feature='non_model_feature'),
    ]

    slicing_specs = [
        config_pb2.SlicingSpec(),
    ]
    eval_config = config_pb2.EvalConfig(
        model_specs=[
            config_pb2.ModelSpec(
                name='candidate',
                label_key='label',
                example_weight_key='example_weight'),
            config_pb2.ModelSpec(
                name='baseline',
                label_key='label',
                example_weight_key='example_weight',
                is_baseline=True)
        ],
        slicing_specs=slicing_specs,
        metrics_specs=[
            config_pb2.MetricsSpec(
                metrics=[
                    config_pb2.MetricConfig(class_name='WeightedExampleCount'),
                    config_pb2.MetricConfig(class_name='ExampleCount'),
                    config_pb2.MetricConfig(class_name='MeanLabel')
                ],
                model_names=['candidate', 'baseline']),
        ],
        options=config_pb2.Options(
            disabled_outputs={'values': ['eval_config_pb2.json']}),
    )
    slice_spec = [
        slicer.SingleSliceSpec(spec=s) for s in eval_config.slicing_specs
    ]
    eval_shared_models = {
        'candidate': eval_shared_model,
        'baseline': baseline_eval_shared_model
    }
    extractors = [
        features_extractor.FeaturesExtractor(eval_config),
        labels_extractor.LabelsExtractor(eval_config),
        example_weights_extractor.ExampleWeightsExtractor(eval_config),
        predictions_extractor.PredictionsExtractor(
            eval_shared_model=eval_shared_models,
            eval_config=eval_config,
            tensor_adapter_config=tensor_adapter_config),
        unbatch_extractor.UnbatchExtractor(),
        slice_key_extractor.SliceKeyExtractor(slice_spec=slice_spec)
    ]
    evaluators = [
        metrics_plots_and_validations_evaluator
        .MetricsPlotsAndValidationsEvaluator(
            eval_config=eval_config, eval_shared_model=eval_shared_models)
    ]
    output_paths = {
        constants.VALIDATIONS_KEY: validations_file,
    }
    writers = [
        metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(
            output_paths,
            eval_config=eval_config,
            add_metrics_callbacks=[],
            output_file_format=output_file_format)
    ]

    with beam.Pipeline() as pipeline:
      # pylint: disable=no-value-for-parameter
      _ = (
          pipeline
          | 'Create' >> beam.Create([e.SerializeToString() for e in examples])
          | 'BatchExamples' >> tfx_io.BeamSource()
          | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
          | 'ExtractEvaluate' >> model_eval_lib.ExtractAndEvaluate(
              extractors=extractors, evaluators=evaluators)
          | 'WriteResults' >> model_eval_lib.WriteResults(writers=writers))
      # pylint: enable=no-value-for-parameter

    validation_result = (
        metrics_plots_and_validations_writer
        .load_and_deserialize_validation_result(
            os.path.dirname(validations_file), output_file_format))

    self.assertFalse(validation_result.validation_ok)
    self.assertTrue(validation_result.missing_thresholds)
    self.assertEmpty(validation_result.metric_validations_per_slice)

    # Add rubber stamp would make validation ok.
    writers = [
        metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(
            output_paths,
            eval_config=eval_config,
            add_metrics_callbacks=[],
            output_file_format=output_file_format,
            rubber_stamp=True)
    ]

    with beam.Pipeline() as pipeline:
      # pylint: disable=no-value-for-parameter
      _ = (
          pipeline
          | 'Create' >> beam.Create([e.SerializeToString() for e in examples])
          | 'BatchExamples' >> tfx_io.BeamSource()
          | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
          | 'ExtractEvaluate' >> model_eval_lib.ExtractAndEvaluate(
              extractors=extractors, evaluators=evaluators)
          | 'WriteResults' >> model_eval_lib.WriteResults(writers=writers))
      # pylint: enable=no-value-for-parameter

    validation_result = (
        metrics_plots_and_validations_writer
        .load_and_deserialize_validation_result(
            os.path.dirname(validations_file), output_file_format))

    self.assertTrue(validation_result.validation_ok)
    self.assertFalse(validation_result.missing_thresholds)
    self.assertEmpty(validation_result.metric_validations_per_slice)
    self.assertTrue(validation_result.rubber_stamp)

  @parameterized.named_parameters(_OUTPUT_FORMAT_PARAMS)
  def testWriteMetricsAndPlots(self, output_file_format):
    metrics_file = os.path.join(self._getTempDir(), 'metrics')
    plots_file = os.path.join(self._getTempDir(), 'plots')
    temp_eval_export_dir = os.path.join(self._getTempDir(), 'eval_export_dir')

    _, eval_export_dir = (
        fixed_prediction_estimator.simple_fixed_prediction_estimator(
            None, temp_eval_export_dir))
    eval_config = config_pb2.EvalConfig(
        model_specs=[config_pb2.ModelSpec()],
        options=config_pb2.Options(
            disabled_outputs={'values': ['eval_config_pb2.json']}))
    eval_shared_model = self.createTestEvalSharedModel(
        eval_saved_model_path=eval_export_dir,
        add_metrics_callbacks=[
            post_export_metrics.example_count(),
            post_export_metrics.calibration_plot_and_prediction_histogram(
                num_buckets=2)
        ])
    extractors = [
        legacy_predict_extractor.PredictExtractor(
            eval_shared_model, eval_config=eval_config),
        unbatch_extractor.UnbatchExtractor(),
        slice_key_extractor.SliceKeyExtractor()
    ]
    evaluators = [
        legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator(
            eval_shared_model)
    ]
    output_paths = {
        constants.METRICS_KEY: metrics_file,
        constants.PLOTS_KEY: plots_file
    }
    writers = [
        metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(
            output_paths,
            eval_config=eval_config,
            add_metrics_callbacks=eval_shared_model.add_metrics_callbacks,
            output_file_format=output_file_format)
    ]

    tfx_io = raw_tf_record.RawBeamRecordTFXIO(
        physical_format='inmemory',
        raw_record_column_name=constants.ARROW_INPUT_COLUMN,
        telemetry_descriptors=['TFMATest'])
    with beam.Pipeline() as pipeline:
      example1 = self._makeExample(prediction=0.0, label=1.0)
      example2 = self._makeExample(prediction=1.0, label=1.0)

      # pylint: disable=no-value-for-parameter
      _ = (
          pipeline
          | 'Create' >> beam.Create([
              example1.SerializeToString(),
              example2.SerializeToString(),
          ])
          | 'BatchExamples' >> tfx_io.BeamSource()
          | 'ExtractEvaluateAndWriteResults' >>
          model_eval_lib.ExtractEvaluateAndWriteResults(
              eval_config=eval_config,
              eval_shared_model=eval_shared_model,
              extractors=extractors,
              evaluators=evaluators,
              writers=writers))
      # pylint: enable=no-value-for-parameter

    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {}
        metrics {
          key: "average_loss"
          value {
            double_value {
              value: 0.5
            }
          }
        }
        metrics {
          key: "post_export_metrics/example_count"
          value {
            double_value {
              value: 2.0
            }
          }
        }
        """, metrics_for_slice_pb2.MetricsForSlice())

    metric_records = list(
        metrics_plots_and_validations_writer.load_and_deserialize_metrics(
            metrics_file, output_file_format))
    self.assertLen(metric_records, 1, 'metrics: %s' % metric_records)
    self.assertProtoEquals(expected_metrics_for_slice, metric_records[0])

    expected_plots_for_slice = text_format.Parse(
        """
      slice_key {}
      plots {
        key: "post_export_metrics"
        value {
          calibration_histogram_buckets {
            buckets {
              lower_threshold_inclusive: -inf
              num_weighted_examples {}
              total_weighted_label {}
              total_weighted_refined_prediction {}
            }
            buckets {
              upper_threshold_exclusive: 0.5
              num_weighted_examples {
                value: 1.0
              }
              total_weighted_label {
                value: 1.0
              }
              total_weighted_refined_prediction {}
            }
            buckets {
              lower_threshold_inclusive: 0.5
              upper_threshold_exclusive: 1.0
              num_weighted_examples {
              }
              total_weighted_label {}
              total_weighted_refined_prediction {}
            }
            buckets {
              lower_threshold_inclusive: 1.0
              upper_threshold_exclusive: inf
              num_weighted_examples {
                value: 1.0
              }
              total_weighted_label {
                value: 1.0
              }
              total_weighted_refined_prediction {
                value: 1.0
              }
            }
         }
        }
      }
    """, metrics_for_slice_pb2.PlotsForSlice())

    plot_records = list(
        metrics_plots_and_validations_writer.load_and_deserialize_plots(
            plots_file, output_file_format))
    self.assertLen(plot_records, 1, 'plots: %s' % plot_records)
    self.assertProtoEquals(expected_plots_for_slice, plot_records[0])

  @parameterized.named_parameters(('parquet_file_format', 'parquet'))
  def testLoadAndDeserializeFilteredMetricsAndPlots(self, output_file_format):
    metrics_file = os.path.join(self._getTempDir(), 'metrics')
    plots_file = os.path.join(self._getTempDir(), 'plots')
    temp_eval_export_dir = os.path.join(self._getTempDir(), 'eval_export_dir')

    _, eval_export_dir = (
        fixed_prediction_estimator.simple_fixed_prediction_estimator(
            None, temp_eval_export_dir))
    eval_config = config_pb2.EvalConfig(
        model_specs=[config_pb2.ModelSpec()],
        slicing_specs=[
            config_pb2.SlicingSpec(),
            config_pb2.SlicingSpec(feature_keys=['prediction'])
        ],
        options=config_pb2.Options(
            disabled_outputs={'values': ['eval_config_pb2.json']}))
    eval_shared_model = self.createTestEvalSharedModel(
        eval_saved_model_path=eval_export_dir,
        add_metrics_callbacks=[
            post_export_metrics.example_count(),
            post_export_metrics.calibration_plot_and_prediction_histogram(
                num_buckets=2)
        ])
    extractors = [
        legacy_predict_extractor.PredictExtractor(
            eval_shared_model, eval_config=eval_config),
        unbatch_extractor.UnbatchExtractor(),
        slice_key_extractor.SliceKeyExtractor(
            eval_config=eval_config, materialize=False)
    ]
    evaluators = [
        legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator(
            eval_shared_model)
    ]
    output_paths = {
        constants.METRICS_KEY: metrics_file,
        constants.PLOTS_KEY: plots_file
    }
    writers = [
        metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(
            output_paths,
            eval_config=eval_config,
            add_metrics_callbacks=eval_shared_model.add_metrics_callbacks,
            output_file_format=output_file_format)
    ]

    tfx_io = raw_tf_record.RawBeamRecordTFXIO(
        physical_format='inmemory',
        raw_record_column_name=constants.ARROW_INPUT_COLUMN,
        telemetry_descriptors=['TFMATest'])
    with beam.Pipeline() as pipeline:
      example1 = self._makeExample(prediction=0.0, label=1.0, country='US')
      example2 = self._makeExample(prediction=1.0, label=1.0, country='CA')

      # pylint: disable=no-value-for-parameter
      _ = (
          pipeline
          | 'Create' >> beam.Create([
              example1.SerializeToString(),
              example2.SerializeToString(),
          ])
          | 'BatchExamples' >> tfx_io.BeamSource()
          | 'ExtractEvaluateAndWriteResults' >>
          model_eval_lib.ExtractEvaluateAndWriteResults(
              eval_config=eval_config,
              eval_shared_model=eval_shared_model,
              extractors=extractors,
              evaluators=evaluators,
              writers=writers))
      # pylint: enable=no-value-for-parameter

    # only read the metrics with slice keys that match the following spec
    slice_keys_filter = [slicer.SingleSliceSpec(features=[('prediction', 0)])]

    expected_metrics_for_slice = text_format.Parse(
        """
        slice_key {
          single_slice_keys {
            column: "prediction"
            float_value: 0
          }
        }
        metrics {
          key: "average_loss"
          value {
            double_value {
              value: 1.0
            }
          }
        }
        metrics {
          key: "post_export_metrics/example_count"
          value {
            double_value {
              value: 1.0
            }
          }
        }
        """, metrics_for_slice_pb2.MetricsForSlice())

    metric_records = list(
        metrics_plots_and_validations_writer.load_and_deserialize_metrics(
            metrics_file, output_file_format, slice_keys_filter))
    self.assertLen(metric_records, 1, 'metrics: %s' % metric_records)
    self.assertProtoEquals(expected_metrics_for_slice, metric_records[0])

    expected_plots_for_slice = text_format.Parse(
        """
      slice_key {
        single_slice_keys {
          column: "prediction"
          float_value: 0
        }
      }
      plots {
        key: "post_export_metrics"
        value {
          calibration_histogram_buckets {
            buckets {
              lower_threshold_inclusive: -inf
              num_weighted_examples {}
              total_weighted_label {}
              total_weighted_refined_prediction {}
            }
            buckets {
              upper_threshold_exclusive: 0.5
              num_weighted_examples {
                value: 1.0
              }
              total_weighted_label {
                value: 1.0
              }
              total_weighted_refined_prediction {}
            }
            buckets {
              lower_threshold_inclusive: 0.5
              upper_threshold_exclusive: 1.0
              num_weighted_examples {
              }
              total_weighted_label {}
              total_weighted_refined_prediction {}
            }
            buckets {
              lower_threshold_inclusive: 1.0
              upper_threshold_exclusive: inf
              num_weighted_examples {
                value: 0.0
              }
              total_weighted_label {
                value: 0.0
              }
              total_weighted_refined_prediction {
                value: 0.0
              }
            }
         }
        }
      }
    """, metrics_for_slice_pb2.PlotsForSlice())

    plot_records = list(
        metrics_plots_and_validations_writer.load_and_deserialize_plots(
            plots_file, output_file_format, slice_keys_filter))
    self.assertLen(plot_records, 1, 'plots: %s' % plot_records)
    self.assertProtoEquals(expected_plots_for_slice, plot_records[0])

  @parameterized.named_parameters(_OUTPUT_FORMAT_PARAMS)
  def testWriteAttributions(self, output_file_format):
    attributions_file = os.path.join(self._getTempDir(), 'attributions')
    eval_config = config_pb2.EvalConfig(
        model_specs=[config_pb2.ModelSpec()],
        metrics_specs=[
            config_pb2.MetricsSpec(metrics=[
                config_pb2.MetricConfig(class_name=attributions
                                        .TotalAttributions().__class__.__name__)
            ])
        ],
        options=config_pb2.Options(
            disabled_outputs={'values': ['eval_config.json']}))
    extractors = [slice_key_extractor.SliceKeyExtractor()]
    evaluators = [
        metrics_plots_and_validations_evaluator
        .MetricsPlotsAndValidationsEvaluator(eval_config=eval_config)
    ]
    output_paths = {
        constants.ATTRIBUTIONS_KEY: attributions_file,
    }
    writers = [
        metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(
            output_paths,
            eval_config=eval_config,
            output_file_format=output_file_format)
    ]

    example1 = {
        'features': {},
        'attributions': {
            'feature1': 1.1,
            'feature2': 1.2
        }
    }
    example2 = {
        'features': {},
        'attributions': {
            'feature1': 2.1,
            'feature2': 2.2
        }
    }
    example3 = {
        'features': {},
        'attributions': {
            'feature1': np.array([3.1]),
            'feature2': np.array([3.2])
        }
    }

    with beam.Pipeline() as pipeline:
      # pylint: disable=no-value-for-parameter
      _ = (
          pipeline
          | 'Create' >> beam.Create([example1, example2, example3])
          | 'ExtractEvaluate' >> model_eval_lib.ExtractAndEvaluate(
              extractors=extractors, evaluators=evaluators)
          | 'WriteResults' >> model_eval_lib.WriteResults(writers=writers))
      # pylint: enable=no-value-for-parameter

    expected_attributions_for_slice = text_format.Parse(
        """
        slice_key {}
        attributions_keys_and_values {
          key {
            name: "total_attributions"
          }
          values {
            key: "feature1"
            value: {
              double_value {
                value: 6.3
              }
            }
          }
          values {
            key: "feature2"
            value: {
              double_value {
                value: 6.6
              }
            }
          }
        }""", metrics_for_slice_pb2.AttributionsForSlice())

    attribution_records = list(
        metrics_plots_and_validations_writer.load_and_deserialize_attributions(
            attributions_file, output_file_format))
    self.assertLen(attribution_records, 1)
    self.assertProtoEquals(expected_attributions_for_slice,
                           attribution_records[0])


if __name__ == '__main__':
  tf.compat.v1.enable_v2_behavior()
  tf.test.main()

tensorflow.compat.v1.enable_v2_behavior tensorflow.keras.layers.Lambda tensorflow.keras.models.Model tensorflow.test.main tensorflow.keras.losses.BinaryCrossentropy tensorflow.keras.optimizers.Adam numpy.array tensorflow.keras.layers.Input
[(1207, 'absl.testing.parameterized.named_parameters', 'parameterized.named_parameters', (['_OUTPUT_FORMAT_PARAMS'], {}), False, 'from absl.testing import parameterized\n'), (1502, 'absl.testing.parameterized.named_parameters', 'parameterized.named_parameters', (['_OUTPUT_FORMAT_PARAMS'], {}), False, 'from absl.testing import parameterized\n'), (1675, 'absl.testing.parameterized.named_parameters', 'parameterized.named_parameters', (['_OUTPUT_FORMAT_PARAMS'], {}), False, 'from absl.testing import parameterized\n'), (1824, 'absl.testing.parameterized.named_parameters', 'parameterized.named_parameters', (["('parquet_file_format', 'parquet')"], {}), False, 'from absl.testing import parameterized\n'), (1991, 'absl.testing.parameterized.named_parameters', 'parameterized.named_parameters', (['_OUTPUT_FORMAT_PARAMS'], {}), False, 'from absl.testing import parameterized\n'), (2085, 'tensorflow.compat.v1.enable_v2_behavior', 'tf.compat.v1.enable_v2_behavior', ([], {}), True, 'import tensorflow as tf\n'), (2086, 'tensorflow.test.main', 'tf.test.main', ([], {}), True, 'import tensorflow as tf\n'), (81, 'tempfile.mkdtemp', 'tempfile.mkdtemp', ([], {}), False, 'import tempfile\n'), (90, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', ([], {'shape': '(1,)', 'name': '"""input_1"""'}), True, 'import tensorflow as tf\n'), (94, 'tensorflow.keras.models.Model', 'tf.keras.models.Model', (['[input_layer]', 'output_layer'], {}), True, 'import tensorflow as tf\n'), (107, 'tensorflow_model_analysis.metrics.metric_types.PlotKey', 'metric_types.PlotKey', ([], {'name': '"""calibration_plot"""', 'output_name': '"""output_name"""'}), False, 'from tensorflow_model_analysis.metrics import metric_types\n'), (188, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_plots_to_proto', 'metrics_plots_and_validations_writer.convert_slice_plots_to_proto', (['(slice_key, {plot_key: calibration_plot})', 'None'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (251, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.calibration_plot_and_prediction_histogram', 'post_export_metrics.calibration_plot_and_prediction_histogram', ([], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (252, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_plots_to_proto', 'metrics_plots_and_validations_writer.convert_slice_plots_to_proto', (['(slice_key, tfma_plots)', '[calibration_plot]'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (260, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_plots_to_proto', 'metrics_plots_and_validations_writer.convert_slice_plots_to_proto', (['(slice_key, tfma_plots)', '[]'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (262, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.PlotsForSlice', 'metrics_for_slice_pb2.PlotsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (301, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', 'metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', (['(slice_key, slice_metrics)', 'None'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (355, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', 'metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', (['(slice_key, slice_metrics)'], {'add_metrics_callbacks': '[]'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (763, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (776, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (785, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', 'metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', (['(slice_key, slice_metrics)', '[]'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1095, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', 'metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', (['(slice_key, slice_metrics)', '[]'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1143, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', 'metrics_plots_and_validations_writer.convert_slice_metrics_to_proto', (['(slice_key, slice_metrics)', '[]'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1199, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.convert_slice_attributions_to_proto', 'metrics_plots_and_validations_writer.convert_slice_attributions_to_proto', (['(slice_key, slice_attributions)'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1247, 'tfx_bsl.tfxio.test_util.InMemoryTFExampleRecord', 'test_util.InMemoryTFExampleRecord', ([], {'schema': 'schema', 'raw_record_column_name': 'constants.ARROW_INPUT_COLUMN'}), False, 'from tfx_bsl.tfxio import test_util\n'), (1542, 'tfx_bsl.tfxio.test_util.InMemoryTFExampleRecord', 'test_util.InMemoryTFExampleRecord', ([], {'schema': 'schema', 'raw_record_column_name': 'constants.ARROW_INPUT_COLUMN'}), False, 'from tfx_bsl.tfxio import test_util\n'), (1682, 'tensorflow_model_analysis.eval_saved_model.example_trainers.fixed_prediction_estimator.simple_fixed_prediction_estimator', 'fixed_prediction_estimator.simple_fixed_prediction_estimator', (['None', 'temp_eval_export_dir'], {}), False, 'from tensorflow_model_analysis.eval_saved_model.example_trainers import fixed_prediction_estimator\n'), (1717, 'tfx_bsl.tfxio.raw_tf_record.RawBeamRecordTFXIO', 'raw_tf_record.RawBeamRecordTFXIO', ([], {'physical_format': '"""inmemory"""', 'raw_record_column_name': 'constants.ARROW_INPUT_COLUMN', 'telemetry_descriptors': "['TFMATest']"}), False, 'from tfx_bsl.tfxio import raw_tf_record\n'), (1831, 'tensorflow_model_analysis.eval_saved_model.example_trainers.fixed_prediction_estimator.simple_fixed_prediction_estimator', 'fixed_prediction_estimator.simple_fixed_prediction_estimator', (['None', 'temp_eval_export_dir'], {}), False, 'from tensorflow_model_analysis.eval_saved_model.example_trainers import fixed_prediction_estimator\n'), (1871, 'tfx_bsl.tfxio.raw_tf_record.RawBeamRecordTFXIO', 'raw_tf_record.RawBeamRecordTFXIO', ([], {'physical_format': '"""inmemory"""', 'raw_record_column_name': 'constants.ARROW_INPUT_COLUMN', 'telemetry_descriptors': "['TFMATest']"}), False, 'from tfx_bsl.tfxio import raw_tf_record\n'), (91, 'tensorflow.keras.layers.Lambda', 'tf.keras.layers.Lambda', (['(lambda x, mul: x * mul)'], {'output_shape': '(1,)', 'arguments': "{'mul': mul}"}), True, 'import tensorflow as tf\n'), (139, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.CalibrationHistogramBuckets', 'metrics_for_slice_pb2.CalibrationHistogramBuckets', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (186, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.PlotsForSlice', 'metrics_for_slice_pb2.PlotsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (198, 'numpy.array', 'np.array', (['[[0.0, 0.0, 0.0], [0.3, 1.0, 1.0], [0.7, 0.0, 1.0], [0.0, 0.0, 0.0]]'], {}), True, 'import numpy as np\n'), (205, 'numpy.array', 'np.array', (['[0.0, 0.5, 1.0]'], {}), True, 'import numpy as np\n'), (263, 'tensorflow_model_analysis.slicer.slicer_lib.serialize_slice_key', 'slicer.serialize_slice_key', (['slice_key'], {}), True, 'from tensorflow_model_analysis.slicer import slicer_lib as slicer\n'), (271, 'tensorflow_model_analysis.metrics.metric_types.MetricKey', 'metric_types.MetricKey', ([], {'name': '"""accuracy"""', 'output_name': '"""output_name"""'}), False, 'from tensorflow_model_analysis.metrics import metric_types\n'), (299, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (308, 'tensorflow_model_analysis.metrics.metric_types.MetricKey', 'metric_types.MetricKey', ([], {'name': '"""confusion_matrix_at_thresholds"""'}), False, 'from tensorflow_model_analysis.metrics import metric_types\n'), (309, 'tensorflow_model_analysis.metrics.binary_confusion_matrices.Matrices', 'binary_confusion_matrices.Matrices', ([], {'thresholds': '[0.25, 0.75, 1.0]', 'fn': '[0.0, 1.0, 2.0]', 'tn': '[1.0, 1.0, 1.0]', 'fp': '[0.0, 0.0, 0.0]', 'tp': '[2.0, 1.0, 0.0]'}), False, 'from tensorflow_model_analysis.metrics import binary_confusion_matrices\n'), (353, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (587, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (599, 'tensorflow_model_analysis.post_export_metrics.metric_keys.lower_bound_key', 'metric_keys.lower_bound_key', (['metric_keys.AUPRC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (600, 'tensorflow_model_analysis.post_export_metrics.metric_keys.upper_bound_key', 'metric_keys.upper_bound_key', (['metric_keys.AUPRC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (602, 'tensorflow_model_analysis.post_export_metrics.metric_keys.lower_bound_key', 'metric_keys.lower_bound_key', (['metric_keys.AUC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (603, 'tensorflow_model_analysis.post_export_metrics.metric_keys.upper_bound_key', 'metric_keys.upper_bound_key', (['metric_keys.AUC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (597, 'tensorflow_model_analysis.types.ValueWithTDistribution', 'types.ValueWithTDistribution', (['(0.8)', '(0.1)', '(9)', '(0.8)'], {}), False, 'from tensorflow_model_analysis import types\n'), (668, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (681, 'tensorflow_model_analysis.post_export_metrics.metric_keys.lower_bound_key', 'metric_keys.lower_bound_key', (['metric_keys.AUPRC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (682, 'tensorflow_model_analysis.post_export_metrics.metric_keys.upper_bound_key', 'metric_keys.upper_bound_key', (['metric_keys.AUPRC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (684, 'tensorflow_model_analysis.post_export_metrics.metric_keys.lower_bound_key', 'metric_keys.lower_bound_key', (['metric_keys.AUC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (685, 'tensorflow_model_analysis.post_export_metrics.metric_keys.upper_bound_key', 'metric_keys.upper_bound_key', (['metric_keys.AUC'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import metric_keys\n'), (745, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (764, 'tensorflow_model_analysis.slicer.slicer_lib.serialize_slice_key', 'slicer.serialize_slice_key', (['slice_key'], {}), True, 'from tensorflow_model_analysis.slicer import slicer_lib as slicer\n'), (884, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationResult', 'validation_result_pb2.ValidationResult', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (890, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (897, 'apache_beam.testing.util.assert_that', 'util.assert_that', (['result', 'verify_fn'], {}), False, 'from apache_beam.testing import util\n'), (959, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (960, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['x']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (961, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['x', 'y']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (962, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['z']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1033, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationResult', 'validation_result_pb2.ValidationResult', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (1039, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (1046, 'apache_beam.testing.util.assert_that', 'util.assert_that', (['result', 'verify_fn'], {}), False, 'from apache_beam.testing import util\n'), (1052, 'tensorflow_model_analysis.types.ValueWithTDistribution', 'types.ValueWithTDistribution', (['(2.0)', '(1.0)', '(3)', '(2.0)'], {}), False, 'from tensorflow_model_analysis import types\n'), (1094, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1103, 'numpy.array', 'np.array', (['[1.0, 2.0, 3.0, 4.0]'], {'dtype': 'np.float32'}), True, 'import numpy as np\n'), (1105, 'numpy.array', 'np.array', (["[['two', 'dims', 'test'], ['TWO', 'DIMS', 'TEST']]"], {}), True, 'import numpy as np\n'), (1107, 'numpy.array', 'np.array', (['[[[100, 200, 300]], [[500, 600, 700]]]'], {'dtype': 'np.int64'}), True, 'import numpy as np\n'), (1142, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1150, 'tensorflow_model_analysis.metrics.metric_types.AttributionsKey', 'metric_types.AttributionsKey', ([], {'name': '"""mean"""', 'output_name': '"""output_name"""'}), False, 'from tensorflow_model_analysis.metrics import metric_types\n'), (1197, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.AttributionsForSlice', 'metrics_for_slice_pb2.AttributionsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1246, 'tensorflow_metadata.proto.v0.schema_pb2.Schema', 'schema_pb2.Schema', ([], {}), False, 'from tensorflow_metadata.proto.v0 import schema_pb2\n'), (1266, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1267, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['slice_does_not_exist']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1351, 'tensorflow_model_analysis.slicer.slicer_lib.SingleSliceSpec', 'slicer.SingleSliceSpec', ([], {'spec': 's'}), True, 'from tensorflow_model_analysis.slicer import slicer_lib as slicer\n'), (1358, 'tensorflow_model_analysis.extractors.features_extractor.FeaturesExtractor', 'features_extractor.FeaturesExtractor', (['eval_config'], {}), False, 'from tensorflow_model_analysis.extractors import features_extractor\n'), (1359, 'tensorflow_model_analysis.extractors.labels_extractor.LabelsExtractor', 'labels_extractor.LabelsExtractor', (['eval_config'], {}), False, 'from tensorflow_model_analysis.extractors import labels_extractor\n'), (1360, 'tensorflow_model_analysis.extractors.example_weights_extractor.ExampleWeightsExtractor', 'example_weights_extractor.ExampleWeightsExtractor', (['eval_config'], {}), False, 'from tensorflow_model_analysis.extractors import example_weights_extractor\n'), (1361, 'tensorflow_model_analysis.extractors.predictions_extractor.PredictionsExtractor', 'predictions_extractor.PredictionsExtractor', ([], {'eval_shared_model': 'eval_shared_models', 'eval_config': 'eval_config', 'tensor_adapter_config': 'tensor_adapter_config'}), False, 'from tensorflow_model_analysis.extractors import predictions_extractor\n'), (1365, 'tensorflow_model_analysis.extractors.unbatch_extractor.UnbatchExtractor', 'unbatch_extractor.UnbatchExtractor', ([], {}), False, 'from tensorflow_model_analysis.extractors import unbatch_extractor\n'), (1366, 'tensorflow_model_analysis.extractors.slice_key_extractor.SliceKeyExtractor', 'slice_key_extractor.SliceKeyExtractor', ([], {'slice_spec': 'slice_spec'}), False, 'from tensorflow_model_analysis.extractors import slice_key_extractor\n'), (1369, 'tensorflow_model_analysis.evaluators.metrics_plots_and_validations_evaluator.MetricsPlotsAndValidationsEvaluator', 'metrics_plots_and_validations_evaluator.MetricsPlotsAndValidationsEvaluator', ([], {'eval_config': 'eval_config', 'eval_shared_model': 'eval_shared_models'}), False, 'from tensorflow_model_analysis.evaluators import metrics_plots_and_validations_evaluator\n'), (1377, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', 'metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', (['output_paths'], {'eval_config': 'eval_config', 'add_metrics_callbacks': '[]', 'output_file_format': 'output_file_format'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1384, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (1399, 'os.path.dirname', 'os.path.dirname', (['validations_file'], {}), False, 'import os\n'), (1473, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['slice_does_not_exist']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1541, 'tensorflow_metadata.proto.v0.schema_pb2.Schema', 'schema_pb2.Schema', ([], {}), False, 'from tensorflow_metadata.proto.v0 import schema_pb2\n'), (1561, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1589, 'tensorflow_model_analysis.slicer.slicer_lib.SingleSliceSpec', 'slicer.SingleSliceSpec', ([], {'spec': 's'}), True, 'from tensorflow_model_analysis.slicer import slicer_lib as slicer\n'), (1596, 'tensorflow_model_analysis.extractors.features_extractor.FeaturesExtractor', 'features_extractor.FeaturesExtractor', (['eval_config'], {}), False, 'from tensorflow_model_analysis.extractors import features_extractor\n'), (1597, 'tensorflow_model_analysis.extractors.labels_extractor.LabelsExtractor', 'labels_extractor.LabelsExtractor', (['eval_config'], {}), False, 'from tensorflow_model_analysis.extractors import labels_extractor\n'), (1598, 'tensorflow_model_analysis.extractors.example_weights_extractor.ExampleWeightsExtractor', 'example_weights_extractor.ExampleWeightsExtractor', (['eval_config'], {}), False, 'from tensorflow_model_analysis.extractors import example_weights_extractor\n'), (1599, 'tensorflow_model_analysis.extractors.predictions_extractor.PredictionsExtractor', 'predictions_extractor.PredictionsExtractor', ([], {'eval_shared_model': 'eval_shared_models', 'eval_config': 'eval_config', 'tensor_adapter_config': 'tensor_adapter_config'}), False, 'from tensorflow_model_analysis.extractors import predictions_extractor\n'), (1603, 'tensorflow_model_analysis.extractors.unbatch_extractor.UnbatchExtractor', 'unbatch_extractor.UnbatchExtractor', ([], {}), False, 'from tensorflow_model_analysis.extractors import unbatch_extractor\n'), (1604, 'tensorflow_model_analysis.extractors.slice_key_extractor.SliceKeyExtractor', 'slice_key_extractor.SliceKeyExtractor', ([], {'slice_spec': 'slice_spec'}), False, 'from tensorflow_model_analysis.extractors import slice_key_extractor\n'), (1607, 'tensorflow_model_analysis.evaluators.metrics_plots_and_validations_evaluator.MetricsPlotsAndValidationsEvaluator', 'metrics_plots_and_validations_evaluator.MetricsPlotsAndValidationsEvaluator', ([], {'eval_config': 'eval_config', 'eval_shared_model': 'eval_shared_models'}), False, 'from tensorflow_model_analysis.evaluators import metrics_plots_and_validations_evaluator\n'), (1615, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', 'metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', (['output_paths'], {'eval_config': 'eval_config', 'add_metrics_callbacks': '[]', 'output_file_format': 'output_file_format'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1622, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (1637, 'os.path.dirname', 'os.path.dirname', (['validations_file'], {}), False, 'import os\n'), (1645, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', 'metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', (['output_paths'], {'eval_config': 'eval_config', 'add_metrics_callbacks': '[]', 'output_file_format': 'output_file_format', 'rubber_stamp': '(True)'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1653, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (1668, 'os.path.dirname', 'os.path.dirname', (['validations_file'], {}), False, 'import os\n'), (1696, 'tensorflow_model_analysis.extractors.legacy_predict_extractor.PredictExtractor', 'legacy_predict_extractor.PredictExtractor', (['eval_shared_model'], {'eval_config': 'eval_config'}), False, 'from tensorflow_model_analysis.extractors import legacy_predict_extractor\n'), (1698, 'tensorflow_model_analysis.extractors.unbatch_extractor.UnbatchExtractor', 'unbatch_extractor.UnbatchExtractor', ([], {}), False, 'from tensorflow_model_analysis.extractors import unbatch_extractor\n'), (1699, 'tensorflow_model_analysis.extractors.slice_key_extractor.SliceKeyExtractor', 'slice_key_extractor.SliceKeyExtractor', ([], {}), False, 'from tensorflow_model_analysis.extractors import slice_key_extractor\n'), (1702, 'tensorflow_model_analysis.evaluators.legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator', 'legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator', (['eval_shared_model'], {}), False, 'from tensorflow_model_analysis.evaluators import legacy_metrics_and_plots_evaluator\n'), (1710, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', 'metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', (['output_paths'], {'eval_config': 'eval_config', 'add_metrics_callbacks': 'eval_shared_model.add_metrics_callbacks', 'output_file_format': 'output_file_format'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1721, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (1761, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1764, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.load_and_deserialize_metrics', 'metrics_plots_and_validations_writer.load_and_deserialize_metrics', (['metrics_file', 'output_file_format'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1816, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.PlotsForSlice', 'metrics_for_slice_pb2.PlotsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1819, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.load_and_deserialize_plots', 'metrics_plots_and_validations_writer.load_and_deserialize_plots', (['plots_file', 'output_file_format'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1849, 'tensorflow_model_analysis.extractors.legacy_predict_extractor.PredictExtractor', 'legacy_predict_extractor.PredictExtractor', (['eval_shared_model'], {'eval_config': 'eval_config'}), False, 'from tensorflow_model_analysis.extractors import legacy_predict_extractor\n'), (1851, 'tensorflow_model_analysis.extractors.unbatch_extractor.UnbatchExtractor', 'unbatch_extractor.UnbatchExtractor', ([], {}), False, 'from tensorflow_model_analysis.extractors import unbatch_extractor\n'), (1852, 'tensorflow_model_analysis.extractors.slice_key_extractor.SliceKeyExtractor', 'slice_key_extractor.SliceKeyExtractor', ([], {'eval_config': 'eval_config', 'materialize': '(False)'}), False, 'from tensorflow_model_analysis.extractors import slice_key_extractor\n'), (1856, 'tensorflow_model_analysis.evaluators.legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator', 'legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator', (['eval_shared_model'], {}), False, 'from tensorflow_model_analysis.evaluators import legacy_metrics_and_plots_evaluator\n'), (1864, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', 'metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', (['output_paths'], {'eval_config': 'eval_config', 'add_metrics_callbacks': 'eval_shared_model.add_metrics_callbacks', 'output_file_format': 'output_file_format'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1875, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (1897, 'tensorflow_model_analysis.slicer.slicer_lib.SingleSliceSpec', 'slicer.SingleSliceSpec', ([], {'features': "[('prediction', 0)]"}), True, 'from tensorflow_model_analysis.slicer import slicer_lib as slicer\n'), (1923, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.MetricsForSlice', 'metrics_for_slice_pb2.MetricsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1926, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.load_and_deserialize_metrics', 'metrics_plots_and_validations_writer.load_and_deserialize_metrics', (['metrics_file', 'output_file_format', 'slice_keys_filter'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1983, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.PlotsForSlice', 'metrics_for_slice_pb2.PlotsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (1986, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.load_and_deserialize_plots', 'metrics_plots_and_validations_writer.load_and_deserialize_plots', (['plots_file', 'output_file_format', 'slice_keys_filter'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (2004, 'tensorflow_model_analysis.extractors.slice_key_extractor.SliceKeyExtractor', 'slice_key_extractor.SliceKeyExtractor', ([], {}), False, 'from tensorflow_model_analysis.extractors import slice_key_extractor\n'), (2006, 'tensorflow_model_analysis.evaluators.metrics_plots_and_validations_evaluator.MetricsPlotsAndValidationsEvaluator', 'metrics_plots_and_validations_evaluator.MetricsPlotsAndValidationsEvaluator', ([], {'eval_config': 'eval_config'}), False, 'from tensorflow_model_analysis.evaluators import metrics_plots_and_validations_evaluator\n'), (2013, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', 'metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter', (['output_paths'], {'eval_config': 'eval_config', 'output_file_format': 'output_file_format'}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (2041, 'apache_beam.Pipeline', 'beam.Pipeline', ([], {}), True, 'import apache_beam as beam\n'), (2074, 'tensorflow_model_analysis.proto.metrics_for_slice_pb2.AttributionsForSlice', 'metrics_for_slice_pb2.AttributionsForSlice', ([], {}), False, 'from tensorflow_model_analysis.proto import metrics_for_slice_pb2\n'), (2077, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.load_and_deserialize_attributions', 'metrics_plots_and_validations_writer.load_and_deserialize_attributions', (['attributions_file', 'output_file_format'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (96, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'lr': '(0.001)'}), True, 'import tensorflow as tf\n'), (97, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', ([], {}), True, 'import tensorflow as tf\n'), (591, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.confusion_matrix_at_thresholds', 'post_export_metrics.confusion_matrix_at_thresholds', (['thresholds'], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (672, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.auc', 'post_export_metrics.auc', ([], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (673, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.auc', 'post_export_metrics.auc', ([], {'curve': '"""PR"""'}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (749, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.auc', 'post_export_metrics.auc', ([], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (750, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.auc', 'post_export_metrics.auc', ([], {'curve': '"""PR"""'}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (760, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.auc', 'post_export_metrics.auc', ([], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (761, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.auc', 'post_export_metrics.auc', ([], {'curve': '"""PR"""'}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (813, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationResult', 'validation_result_pb2.ValidationResult', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (830, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationResult', 'validation_result_pb2.ValidationResult', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (938, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationResult', 'validation_result_pb2.ValidationResult', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (955, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationResult', 'validation_result_pb2.ValidationResult', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (1347, 'tensorflow_model_analysis.proto.config_pb2.Options', 'config_pb2.Options', ([], {'disabled_outputs': "{'values': ['eval_config_pb2.json']}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1420, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationFailure', 'validation_result_pb2.ValidationFailure', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (1439, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationFailure', 'validation_result_pb2.ValidationFailure', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (1463, 'tensorflow_model_analysis.proto.validation_result_pb2.ValidationFailure', 'validation_result_pb2.ValidationFailure', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (1496, 'tensorflow_model_analysis.proto.validation_result_pb2.SlicingDetails', 'validation_result_pb2.SlicingDetails', ([], {}), False, 'from tensorflow_model_analysis.proto import validation_result_pb2\n'), (1585, 'tensorflow_model_analysis.proto.config_pb2.Options', 'config_pb2.Options', ([], {'disabled_outputs': "{'values': ['eval_config_pb2.json']}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1686, 'tensorflow_model_analysis.proto.config_pb2.Options', 'config_pb2.Options', ([], {'disabled_outputs': "{'values': ['eval_config_pb2.json']}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1839, 'tensorflow_model_analysis.proto.config_pb2.Options', 'config_pb2.Options', ([], {'disabled_outputs': "{'values': ['eval_config_pb2.json']}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (2002, 'tensorflow_model_analysis.proto.config_pb2.Options', 'config_pb2.Options', ([], {'disabled_outputs': "{'values': ['eval_config.json']}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (2036, 'numpy.array', 'np.array', (['[3.1]'], {}), True, 'import numpy as np\n'), (2037, 'numpy.array', 'np.array', (['[3.2]'], {}), True, 'import numpy as np\n'), (606, 'string.Template', 'string.Template', (['"""\n        slice_key {\n          single_slice_keys {\n            column: \'age\'\n            int64_value: 5\n          }\n          single_slice_keys {\n            column: \'language\'\n            bytes_value: \'english\'\n          }\n          single_slice_keys {\n            column: \'price\'\n            float_value: 0.3\n          }\n        }\n        metrics {\n          key: "accuracy"\n          value {\n            bounded_value {\n              value {\n                value: 0.8\n              }\n              lower_bound { value: 0.5737843 }\n              upper_bound { value: 1.0262157 }\n              methodology: POISSON_BOOTSTRAP\n            }\n          }\n        }\n        metrics {\n          key: "$auc"\n          value {\n            bounded_value {\n              lower_bound {\n                value: 0.1\n              }\n              upper_bound {\n                value: 0.3\n              }\n              value {\n                value: 0.2\n              }\n              methodology: RIEMANN_SUM\n            }\n          }\n        }\n        metrics {\n          key: "$auprc"\n          value {\n            bounded_value {\n              lower_bound {\n                value: 0.05\n              }\n              upper_bound {\n                value: 0.17\n              }\n              value {\n                value: 0.1\n              }\n              methodology: RIEMANN_SUM\n            }\n          }\n        }"""'], {}), False, 'import string\n'), (688, 'string.Template', 'string.Template', (['"""\n        slice_key {\n          single_slice_keys {\n            column: \'age\'\n            int64_value: 5\n          }\n          single_slice_keys {\n            column: \'language\'\n            bytes_value: \'english\'\n          }\n          single_slice_keys {\n            column: \'price\'\n            float_value: 0.3\n          }\n        }\n        metrics {\n          key: "accuracy"\n          value {\n            double_value {\n              value: 0.8\n            }\n          }\n        }\n        metrics {\n          key: "$auc"\n          value {\n            bounded_value {\n              lower_bound {\n                value: 0.1\n              }\n              upper_bound {\n                value: 0.3\n              }\n              value {\n                value: 0.2\n              }\n              methodology: RIEMANN_SUM\n            }\n          }\n        }\n        metrics {\n          key: "$auprc"\n          value {\n            bounded_value {\n              lower_bound {\n                value: 0.05\n              }\n              upper_bound {\n                value: 0.17\n              }\n              value {\n                value: 0.1\n              }\n              methodology: RIEMANN_SUM\n            }\n          }\n        }"""'], {}), False, 'import string\n'), (835, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""candidate"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (836, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""baseline"""', 'is_baseline': '(True)'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (838, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (966, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""candidate"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (967, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""baseline"""', 'is_baseline': '(True)'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1271, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['slice_does_not_exist']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1279, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""candidate"""', 'label_key': '"""label"""', 'example_weight_key': '"""example_weight"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1283, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""baseline"""', 'label_key': '"""label"""', 'example_weight_key': '"""example_weight"""', 'is_baseline': '(True)'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1393, 'tensorflow_model_analysis.api.model_eval_lib.WriteResults', 'model_eval_lib.WriteResults', ([], {'writers': 'writers'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1480, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['slice_does_not_exist']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1565, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""candidate"""', 'label_key': '"""label"""', 'example_weight_key': '"""example_weight"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1569, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {'name': '"""baseline"""', 'label_key': '"""label"""', 'example_weight_key': '"""example_weight"""', 'is_baseline': '(True)'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1631, 'tensorflow_model_analysis.api.model_eval_lib.WriteResults', 'model_eval_lib.WriteResults', ([], {'writers': 'writers'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1662, 'tensorflow_model_analysis.api.model_eval_lib.WriteResults', 'model_eval_lib.WriteResults', ([], {'writers': 'writers'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1685, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1691, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.example_count', 'post_export_metrics.example_count', ([], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (1692, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.calibration_plot_and_prediction_histogram', 'post_export_metrics.calibration_plot_and_prediction_histogram', ([], {'num_buckets': '(2)'}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (1734, 'tensorflow_model_analysis.api.model_eval_lib.ExtractEvaluateAndWriteResults', 'model_eval_lib.ExtractEvaluateAndWriteResults', ([], {'eval_config': 'eval_config', 'eval_shared_model': 'eval_shared_model', 'extractors': 'extractors', 'evaluators': 'evaluators', 'writers': 'writers'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1834, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1836, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1837, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['prediction']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1844, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.example_count', 'post_export_metrics.example_count', ([], {}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (1845, 'tensorflow_model_analysis.post_export_metrics.post_export_metrics.calibration_plot_and_prediction_histogram', 'post_export_metrics.calibration_plot_and_prediction_histogram', ([], {'num_buckets': '(2)'}), False, 'from tensorflow_model_analysis.post_export_metrics import post_export_metrics\n'), (1888, 'tensorflow_model_analysis.api.model_eval_lib.ExtractEvaluateAndWriteResults', 'model_eval_lib.ExtractEvaluateAndWriteResults', ([], {'eval_config': 'eval_config', 'eval_shared_model': 'eval_shared_model', 'extractors': 'extractors', 'evaluators': 'evaluators', 'writers': 'writers'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1995, 'tensorflow_model_analysis.proto.config_pb2.ModelSpec', 'config_pb2.ModelSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (2048, 'tensorflow_model_analysis.api.model_eval_lib.WriteResults', 'model_eval_lib.WriteResults', ([], {'writers': 'writers'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (893, 'apache_beam.Create', 'beam.Create', (['input_validations'], {}), True, 'import apache_beam as beam\n'), (895, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.CombineValidations', 'metrics_plots_and_validations_writer.CombineValidations', (['eval_config'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1042, 'apache_beam.Create', 'beam.Create', (['input_validations'], {}), True, 'import apache_beam as beam\n'), (1044, 'tensorflow_model_analysis.writers.metrics_plots_and_validations_writer.CombineValidations', 'metrics_plots_and_validations_writer.CombineValidations', (['eval_config'], {}), False, 'from tensorflow_model_analysis.writers import metrics_plots_and_validations_writer\n'), (1274, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['slice_does_not_exist']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1391, 'tensorflow_model_analysis.api.model_eval_lib.ExtractAndEvaluate', 'model_eval_lib.ExtractAndEvaluate', ([], {'extractors': 'extractors', 'evaluators': 'evaluators'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1483, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {'feature_keys': "['slice_does_not_exist']"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1629, 'tensorflow_model_analysis.api.model_eval_lib.ExtractAndEvaluate', 'model_eval_lib.ExtractAndEvaluate', ([], {'extractors': 'extractors', 'evaluators': 'evaluators'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1660, 'tensorflow_model_analysis.api.model_eval_lib.ExtractAndEvaluate', 'model_eval_lib.ExtractAndEvaluate', ([], {'extractors': 'extractors', 'evaluators': 'evaluators'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (2046, 'tensorflow_model_analysis.api.model_eval_lib.ExtractAndEvaluate', 'model_eval_lib.ExtractAndEvaluate', ([], {'extractors': 'extractors', 'evaluators': 'evaluators'}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1390, 'tensorflow_model_analysis.api.model_eval_lib.BatchedInputsToExtracts', 'model_eval_lib.BatchedInputsToExtracts', ([], {}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1628, 'tensorflow_model_analysis.api.model_eval_lib.BatchedInputsToExtracts', 'model_eval_lib.BatchedInputsToExtracts', ([], {}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (1659, 'tensorflow_model_analysis.api.model_eval_lib.BatchedInputsToExtracts', 'model_eval_lib.BatchedInputsToExtracts', ([], {}), False, 'from tensorflow_model_analysis.api import model_eval_lib\n'), (2045, 'apache_beam.Create', 'beam.Create', (['[example1, example2, example3]'], {}), True, 'import apache_beam as beam\n'), (1579, 'tensorflow_model_analysis.proto.config_pb2.MetricConfig', 'config_pb2.MetricConfig', ([], {'class_name': '"""WeightedExampleCount"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1580, 'tensorflow_model_analysis.proto.config_pb2.MetricConfig', 'config_pb2.MetricConfig', ([], {'class_name': '"""ExampleCount"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1581, 'tensorflow_model_analysis.proto.config_pb2.MetricConfig', 'config_pb2.MetricConfig', ([], {'class_name': '"""MeanLabel"""'}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1318, 'tensorflow_model_analysis.proto.config_pb2.GenericValueThreshold', 'config_pb2.GenericValueThreshold', ([], {'lower_bound': "{'value': 10}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1324, 'tensorflow_model_analysis.proto.config_pb2.GenericChangeThreshold', 'config_pb2.GenericChangeThreshold', ([], {'direction': 'config_pb2.MetricDirection.HIGHER_IS_BETTER', 'relative': "{'value': 1}", 'absolute': "{'value': 1}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1336, 'tensorflow_model_analysis.proto.config_pb2.GenericValueThreshold', 'config_pb2.GenericValueThreshold', ([], {'upper_bound': "{'value': 0.01}", 'lower_bound': "{'value': -0.01}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1339, 'tensorflow_model_analysis.proto.config_pb2.GenericChangeThreshold', 'config_pb2.GenericChangeThreshold', ([], {'direction': 'config_pb2.MetricDirection.LOWER_IS_BETTER', 'relative': "{'value': -0.99}", 'absolute': "{'value': 0}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1998, 'tensorflow_model_analysis.metrics.attributions.TotalAttributions', 'attributions.TotalAttributions', ([], {}), False, 'from tensorflow_model_analysis.metrics import attributions\n'), (846, 'tensorflow_model_analysis.proto.config_pb2.SlicingSpec', 'config_pb2.SlicingSpec', ([], {}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (848, 'tensorflow_model_analysis.proto.config_pb2.GenericValueThreshold', 'config_pb2.GenericValueThreshold', ([], {'lower_bound': "{'value': 0.7}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (979, 'tensorflow_model_analysis.proto.config_pb2.GenericValueThreshold', 'config_pb2.GenericValueThreshold', ([], {'lower_bound': "{'value': 0.7}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1301, 'tensorflow_model_analysis.proto.config_pb2.GenericValueThreshold', 'config_pb2.GenericValueThreshold', ([], {'upper_bound': "{'value': 1}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n'), (1310, 'tensorflow_model_analysis.proto.config_pb2.GenericValueThreshold', 'config_pb2.GenericValueThreshold', ([], {'upper_bound': "{'value': 1}"}), False, 'from tensorflow_model_analysis.proto import config_pb2\n')]