extracts: beam.pvalue.PCollection, extractors: List[extractor.Extractor]):

for x in extractors:

extracts = (extracts | x.stage_name >> x.ptransform)

"""Represents a bounded value for a metric for the TFMA unit test."""

def __init__(self,

lower_bound: float = float('-inf'),

upper_bound: float = float('inf')):

self.lower_bound = lower_bound

"""Test class with extra methods for unit-testing TFMA models / metrics."""

def makeExample(self, **kwargs):

"""Returns a serialized TF.Example with the fields set accordingly.

Example usage:

makeExample(name="Tom", age=20, weight=132.5, colors=["red", "green"],

label=1.0)

Note that field types will be inferred from the type of the arguments, so

1.0 goes into a float_list, 1 goes into an int64_list, and "one" goes into

a bytes_list. As the example illustrates, both singleton values and lists

of values are accepted.

Args:

**kwargs: key=value pairs specifying the field values for the example.

Returns:

Serialized TF.Example with the fields set accordingly.

"""

return self._makeExample(**kwargs).SerializeToString()

def assertDictElementsWithinBounds(self, got_values_dict: Dict[str, Any],

expected_values_dict: Dict[str, Any]):

"""Checks the elements for two dictionaries.

It asserts all values in `expected_values_dict` are close to values with the

same key in `got_values_dict`.

Args:

got_values_dict: The actual dictionary.

expected_values_dict: The expected dictionary. The values in can be either

`BoundedValue` or any type accepted by

`tf.test.TestCase.assertAllClose()`. When the type is `BoundedValue`, it

expects the corresponding value from `got_values_dict` falls into the

boundaries provided in the `BoundedValue`.

"""

for key, value in expected_values_dict.items():

self.assertIn(key, got_values_dict)

got_value = got_values_dict[key]

if isinstance(value, BoundedValue):

if got_value < value.lower_bound or got_value > value.upper_bound:

self.fail('expecting key %s to have value between %f and %f '

'(both ends inclusive), but value was %f instead' %

(key, value.lower_bound, value.upper_bound, got_value))

else:

self.assertAllClose(got_value, value, msg='key = %s' % key)

def assertMetricsComputedWithoutBeamAre(self, eval_saved_model_path: str,

serialized_examples: List[bytes],

expected_metrics: Dict[str, Any]):

"""Checks metrics in-memory using the low-level APIs without Beam.

Example usage:

self.assertMetricsComputedWithoutBeamAre(

eval_saved_model_path=path,

serialized_examples=[self.makeExample(age=5, label=1.0),

self.makeExample(age=10, label=0.0)],

expected_metrics={'average_loss': 0.1})

Args:

eval_saved_model_path: Path to the directory containing the

EvalSavedModel.

serialized_examples: List of serialized example bytes.

expected_metrics: Dictionary of expected metric values.

"""

self.assertDictElementsWithinBounds(

got_values_dict=self._computeMetricsWithoutBeam(eval_saved_model_path,

serialized_examples),

expected_values_dict=expected_metrics)

def assertMetricsComputedWithoutBeamNoBatchingAre(

self, eval_saved_model_path: str, serialized_examples: List[bytes],

expected_metrics: Dict[str, Any]):

"""Checks metrics in-memory using the low-level APIs without Beam.

This is the non-batched version of assertMetricsComputedWithoutBeamAre.

This can be useful for debugging batching issues with TFMA or with your

model (e.g. your model or metrics only works with a fixed-batch size - TFMA

requires that your model can accept batches of any size).

Args:

eval_saved_model_path: Path to the directory containing the

EvalSavedModel.

serialized_examples: List of serialized example bytes.

expected_metrics: Dictionary of expected metric values.

"""

self.assertDictElementsWithinBounds(

got_values_dict=self._computeMetricsWithoutBeamNoBatching(

eval_saved_model_path, serialized_examples),

expected_values_dict=expected_metrics)

def _computeMetricsWithoutBeam(

self, eval_saved_model_path: str,

serialized_examples: List[bytes]) -> Dict[str, Any]:

"""Computes metrics in-memory using the low-level APIs without Beam.

Args:

eval_saved_model_path: Path to the directory containing the

EvalSavedModel.

serialized_examples: List of serialized example bytes.

Returns:

Metrics computed by TFMA using your model on the given examples.

"""

eval_saved_model = load.EvalSavedModel(eval_saved_model_path)

eval_saved_model.metrics_reset_update_get_list(serialized_examples)

return eval_saved_model.get_metric_values()

def _computeMetricsWithoutBeamNoBatching(

self, eval_saved_model_path: str,

serialized_examples: List[bytes]) -> Dict[str, Any]:

"""Computes metrics in-memory using the low-level APIs without Beam.

This is the non-batched version of computeMetricsWithoutBeam. This can be

useful for debugging batching issues with TFMA or with your model

(e.g. your model or metrics only works with a fixed-batch size - TFMA

requires that your model can accept batches of any size)

Args:

eval_saved_model_path: Path to the directory containing the

EvalSavedModel.

serialized_examples: List of serialized example bytes.

Returns:

Metrics computed by TFMA using your model on the given examples.

"""

eval_saved_model = load.EvalSavedModel(eval_saved_model_path)

for example in serialized_examples:

eval_saved_model.metrics_reset_update_get_list([example])

return eval_saved_model.get_metric_values()

def assertMetricsComputedWithBeamAre(

self,

eval_saved_model_path: str,

serialized_examples: List[bytes],

expected_metrics: Dict[str, Any],

add_metrics_callbacks: Optional[List[

types.AddMetricsCallbackType]] = None):

"""Checks metrics computed using Beam.

Metrics will be computed over all examples, without any slicing. If you

want to provide your own PCollection (e.g. read a large number of examples

from a file), if you want to check metrics over certain slices, or if you

want to add additional post-export metrics, use the more general

assertGeneralMetricsComputedWithBeamAre.

Example usage:

self.assertMetricsComputedWithBeamAre(

eval_saved_model_path=path,

serialized_examples=[self.makeExample(age=5, label=1.0),

self.makeExample(age=10, label=0.0)],

expected_metrics={'average_loss': 0.1})

Args:

eval_saved_model_path: Path to the directory containing the

EvalSavedModel.

serialized_examples: List of serialized example bytes.

expected_metrics: Dictionary of expected metric values.

add_metrics_callbacks: Optional. Callbacks for adding additional metrics.

"""

def check_metrics(got):

"""Check metrics callback."""

try:

self.assertEqual(

1, len(got), 'expecting metrics for exactly one slice, but got %d '

'slices instead. metrics were: %s' % (len(got), got))

(slice_key, value) = got[0]

self.assertEqual((), slice_key)

self.assertDictElementsWithinBounds(

got_values_dict=value, expected_values_dict=expected_metrics)

except AssertionError as err:

raise beam_util.BeamAssertException(err)

eval_config = config_pb2.EvalConfig()

eval_shared_model = model_eval_lib.default_eval_shared_model(

eval_saved_model_path=eval_saved_model_path,

add_metrics_callbacks=add_metrics_callbacks)

extractors = model_eval_lib.default_extractors(

eval_config=eval_config, eval_shared_model=eval_shared_model)

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:

# pylint: disable=no-value-for-parameter

(metrics, _), _ = (

pipeline

| 'CreateExamples' >> beam.Create(serialized_examples)

| 'BatchExamples' >> tfx_io.BeamSource()

| 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()

| 'Extract' >> Extract(extractors=extractors)

| 'ComputeMetricsAndPlots' >>

legacy_metrics_and_plots_evaluator._ComputeMetricsAndPlots( # pylint: disable=protected-access

eval_shared_model=eval_shared_model))

# pylint: enable=no-value-for-parameter

beam_util.assert_that(metrics, check_metrics)

def assertGeneralMetricsComputedWithBeamAre(

self, eval_saved_model_path: str,

examples_pcollection: beam.pvalue.PCollection,

slice_spec: List[slicer.SingleSliceSpec],

add_metrics_callbacks: List[types.AddMetricsCallbackType],

expected_slice_metrics: Dict[Any, Dict[str, Any]]):

"""Checks metrics computed using Beam.

A more general version of assertMetricsComputedWithBeamAre. Note that the

caller is responsible for setting up and running the Beam pipeline.

Example usage:

def add_metrics(features, predictions, labels):

metric_ops = {

'mse': tf.metrics.mean_squared_error(labels, predictions['logits']),

'mae': tf.metrics.mean_absolute_error(labels, predictions['logits']),

}

return metric_ops

with beam.Pipeline() as pipeline:

expected_slice_metrics = {

(): {

'mae': 0.1,

'mse': 0.2,

tfma.post_export_metrics.metric_keys.AUC:

tfma.test.BoundedValue(lower_bound=0.5)

},

(('age', 10),): {

'mae': 0.2,

'mse': 0.3,

tfma.post_export_metrics.metric_keys.AUC:

tfma.test.BoundedValue(lower_bound=0.5)

},

}

examples = pipeline | 'ReadExamples' >> beam.io.ReadFromTFRecord(path)

self.assertGeneralMetricsComputedWithBeamAre(

eval_saved_model_path=path,

examples_pcollection=examples,

slice_spec=[tfma.slicer.SingleSliceSpec(),

tfma.slicer.SingleSliceSpec(columns=['age'])],

add_metrics_callbacks=[

add_metrics, tfma.post_export_metrics.auc()],

expected_slice_metrics=expected_slice_metrics)

Args:

eval_saved_model_path: Path to the directory containing the

EvalSavedModel.

examples_pcollection: A PCollection of serialized example bytes.

slice_spec: List of slice specifications.

add_metrics_callbacks: Callbacks for adding additional metrics.

expected_slice_metrics: Dictionary of dictionaries describing the expected

metrics for each slice. The outer dictionary map slice keys to the

expected metrics for that slice.

"""

def check_metrics(got):

"""Check metrics callback."""

try:

slices = {}

for slice_key, value in got:

slices[slice_key] = value

self.assertCountEqual(list(slices), list(expected_slice_metrics))

for slice_key, expected_metrics in expected_slice_metrics.items():

self.assertDictElementsWithinBounds(

got_values_dict=slices[slice_key],

expected_values_dict=expected_metrics)

except AssertionError as err:

raise beam_util.BeamAssertException(err)

slicing_specs = None

if slice_spec:

slicing_specs = [s.to_proto() for s in slice_spec]

eval_config = config_pb2.EvalConfig(slicing_specs=slicing_specs)

eval_shared_model = self.createTestEvalSharedModel(

eval_saved_model_path=eval_saved_model_path,

add_metrics_callbacks=add_metrics_callbacks)

extractors = model_eval_lib.default_extractors(

eval_config=eval_config, eval_shared_model=eval_shared_model)

tfx_io = raw_tf_record.RawBeamRecordTFXIO(

physical_format='inmemory',

raw_record_column_name=constants.ARROW_INPUT_COLUMN,

telemetry_descriptors=['TFMATest'])

# pylint: disable=no-value-for-parameter

(metrics, _), _ = (

examples_pcollection

| 'BatchExamples' >> tfx_io.BeamSource()

| 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()

| 'Extract' >> Extract(extractors=extractors)

| 'ComputeMetricsAndPlots' >>

legacy_metrics_and_plots_evaluator._ComputeMetricsAndPlots( # pylint: disable=protected-access

eval_shared_model=eval_shared_model))

# pylint: enable=no-value-for-parameter