"""Check that we're using a compatible TF version."""

# Fail with a clear error in case we are not using a compatible TF version.

major, minor, _ = tf.version.VERSION.split('.')

if (int(major) not in (1, 2)) or (int(major) == 1 and int(minor) < 15):

raise RuntimeError(

'Tensorflow version >= 1.15, < 3 is required. Found (%s). Please '

'install the latest 1.x or 2.x version from '

'https://github.com/tensorflow/tensorflow. ' % tf.version.VERSION)

if int(major) == 2:

logging.warning(

'Tensorflow version (%s) found. Note that TFMA support for TF 2.0 '

config_version: Optional[int],

eval_shared_models: Optional[List[types.EvalSharedModel]],

eval_config: Optional[config_pb2.EvalConfig],

):

"""Returns True if legacy evaluation is being used.

A legacy evaluation is an evalution that uses only a single EvalSharedModel,

has no tags (or uses "eval" as its tag), and does not specify an eval_config

The legacy evaluation is based on using add_metrics_callbacks to create a

modified version of the graph saved with an EvalSavedModel. The newer version

of evaluation supports both add_metrics_callbacks as well as metrics defined

in MetricsSpecs inside of EvalConfig. The newer version works with both "eval"

and serving models and also supports multi-model evaluation. This function is

used by code to support backwards compatibility for callers that have not

updated to use the new EvalConfig.

Args:

config_version: Optionally, An explicit version of the config determined

elsewhere. This is used to handle cases where the provided eval_config was

generated internally, and thus not a reliable indicator of user intent.

eval_shared_models: Optionally, the model(s) to be evaluated.

eval_config: Optionally, an EvalConfig specifying v2 config.

Returns:

Whether the user inputs should trigger a legacy evaluation.

"""

return (config_version is not None and config_version == 1) or (

eval_shared_models

and len(eval_shared_models) == 1

and eval_shared_models[0].model_type == constants.TFMA_EVAL

and not eval_config

slice_spec: Optional[List[slicer.SingleSliceSpec]],

write_config: Optional[bool],

compute_confidence_intervals: Optional[bool],

min_slice_size: int):

"""Creates default EvalConfig (for use in legacy evaluations)."""

model_specs = []

for shared_model in eval_shared_models:

example_weight_key = shared_model.example_weight_key

example_weight_keys = {}

if example_weight_key and isinstance(example_weight_key, dict):

example_weight_keys = example_weight_key

example_weight_key = ''

model_specs.append(

config_pb2.ModelSpec(

name=shared_model.model_name,

example_weight_key=example_weight_key,

example_weight_keys=example_weight_keys))

slicing_specs = None

if slice_spec:

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

options = config_pb2.Options()

options.compute_confidence_intervals.value = compute_confidence_intervals

options.min_slice_size.value = min_slice_size

if not write_config:

options.disabled_outputs.values.append(eval_config_writer.EVAL_CONFIG_FILE)

return config_pb2.EvalConfig(

eval_shared_models: Optional[List[types.EvalSharedModel]],

) -> Optional[Set[str]]:

"""Returns model types associated with given EvalSharedModels."""

if not eval_shared_models:

return None

else:

eval_config: config_pb2.EvalConfig,

eval_shared_model: Optional[types.MaybeMultipleEvalSharedModels]

) -> config_pb2.EvalConfig:

"""Returns updated eval config with default values."""

eval_shared_models = model_util.verify_and_update_eval_shared_models(

eval_shared_model)

has_baseline = eval_shared_models and len(eval_shared_models) > 1

return config_util.update_eval_config_with_defaults(

eval_config=eval_config,

has_baseline=has_baseline,

"""Returns the number of bytes in the input."""

if constants.ARROW_RECORD_BATCH_KEY in extract:

return extract[constants.ARROW_RECORD_BATCH_KEY].nbytes

if constants.INPUT_KEY in extract:

if isinstance(extract[constants.INPUT_KEY], bytes):

return len(extract[constants.INPUT_KEY])

logging.warning('Failed to extract number of input bytes.')

"""Increments the specified counter by the value."""

counter = beam.metrics.Metrics.counter(constants.METRICS_NAMESPACE,

counter_name)

counter.inc(value)

dataset: beam.PCollection[types.Extracts]) -> beam.pvalue.PCollection[int]:

"""Gathers telemetry on input Extracts."""

return (dataset

| 'GetExtractSize' >> beam.Map(_get_extract_num_bytes)

| 'SumTotalBytes' >> beam.CombineGlobally(sum)

| 'IncrementCounter' >>

output_path: str,

output_file_format: str = 'tfrecord') -> Iterator[MetricsForSlice]:

"""Read and deserialize the MetricsForSlice records."""

for m in metrics_plots_and_validations_writer.load_and_deserialize_metrics(

output_path, output_file_format):

output_file_format: str = 'tfrecord') -> Iterator[PlotsForSlice]:

"""Read and deserialize the PlotsForSlice records."""

for p in metrics_plots_and_validations_writer.load_and_deserialize_plots(

output_path, output_file_format):

output_path: str,

output_file_format: str = 'tfrecord') -> Iterator[AttributionsForSlice]:

"""Read and deserialize the AttributionsForSlice records."""

for a in (

metrics_plots_and_validations_writer.load_and_deserialize_attributions(

output_path, output_file_format)):

output_file_format: str = '') -> ValidationResult:

"""Read and deserialize the ValidationResult."""

return metrics_plots_and_validations_writer.load_and_deserialize_validation_result(

mode: str) -> view_types.EvalResults:

"""Run model analysis for a single model on multiple data sets.

Args:

results: A list of TFMA evaluation results.

mode: The mode of the evaluation. Currently, tfma.DATA_CENTRIC_MODE and

tfma.MODEL_CENTRIC_MODE are supported.

Returns:

An `tfma.view.EvalResults` object containing all evaluation results. This

can be used to construct a time series view.

"""

output_paths: Union[str, List[str]],

output_file_format: Optional[str] = 'tfrecord',

mode: str = constants.MODEL_CENTRIC_MODE,

model_name: Optional[str] = None) -> view_types.EvalResults:

"""Loads results for multiple models or multiple data sets.

Args:

output_paths: A single path or list of output paths of completed tfma runs.

output_file_format: Optional file extension to filter files by.

mode: The mode of the evaluation. Currently, tfma.DATA_CENTRIC_MODE and

tfma.MODEL_CENTRIC_MODE are supported.

model_name: Filters to only return results for given model. If unset all

models are returned.

Returns:

An EvalResults containing the evaluation results serialized at output_paths.

This can be used to construct a time series view.

"""

results = []

if not isinstance(output_paths, list):

output_paths = [output_paths]

for output_path in output_paths:

if model_name is None:

_, _, _, model_locations = eval_config_writer.load_eval_run(output_path)

model_names = list(model_locations)

else:

model_names = [model_name]

for model_name in model_names:

results.append(

load_eval_result(

output_path, output_file_format, model_name=model_name))

output_file_format: Optional[str] = 'tfrecord',

model_name: Optional[str] = None) -> view_types.EvalResult:

"""Loads EvalResult object for use with the visualization functions.

Args:

output_path: Output directory containing config, metrics, plots, etc.

output_file_format: Optional file extension to filter files by.

model_name: Optional model name. Required if multi-model evaluation was run.

Returns:

EvalResult object for use with the visualization functions.

"""

# Config, metrics, and plots files should all exist under the given output

# directory, but fairness plugin has a use-case where only the metrics are

# provided so we support all files as being optional (the EvalResult will have

# corresponding None values for files that are not present).

eval_config, data_location, file_format, model_locations = (

eval_config_writer.load_eval_run(output_path)

)

metrics_list = []

for p in metrics_plots_and_validations_writer.load_and_deserialize_metrics(

output_path, output_file_format

):

metrics = view_util.convert_metrics_proto_to_dict(p, model_name=model_name)

if metrics is not None:

metrics_list.append(metrics)

plots_list = []

for p in metrics_plots_and_validations_writer.load_and_deserialize_plots(

output_path, output_file_format

):

plots = view_util.convert_plots_proto_to_dict(p, model_name=model_name)

if plots is not None:

plots_list.append(plots)

attributions_list = []

for (

a

) in metrics_plots_and_validations_writer.load_and_deserialize_attributions(

output_path, output_file_format

):

attributions = view_util.convert_attributions_proto_to_dict(

a, model_name=model_name

)

if attributions is not None:

attributions_list.append(attributions)

if not model_locations:

model_location = ''

elif model_name is None:

model_location = list(model_locations.values())[0]

else:

model_location = model_locations[model_name]

return view_types.EvalResult( # pytype: disable=wrong-arg-types

slicing_metrics=metrics_list,

plots=plots_list,

attributions=attributions_list,

config=eval_config,

data_location=data_location,

file_format=file_format,

eval_saved_model_path: str,

add_metrics_callbacks: Optional[List[types.AddMetricsCallbackType]] = None,

include_default_metrics: Optional[bool] = True,

example_weight_key: Optional[Union[str, Dict[str, str]]] = None,

additional_fetches: Optional[List[str]] = None,

blacklist_feature_fetches: Optional[List[str]] = None,

tags: Optional[List[str]] = None,

model_name: str = '',

eval_config: Optional[config_pb2.EvalConfig] = None,

custom_model_loader: Optional[types.ModelLoader] = None,

rubber_stamp: Optional[bool] = False,

resource_hints: Optional[Dict[str, Any]] = None,

backend_config: Optional[Any] = None,

) -> types.EvalSharedModel:

"""Returns default EvalSharedModel.

Args:

eval_saved_model_path: Path to EvalSavedModel.

add_metrics_callbacks: Optional list of callbacks for adding additional

metrics to the graph (see EvalSharedModel for more information on how to

configure additional metrics). Metrics for example count and example

weights will be added automatically. Only used if EvalSavedModel used.

include_default_metrics: DEPRECATED. Use

eval_config.options.include_default_metrics.

example_weight_key: DEPRECATED. Use

eval_config.model_specs.example_weight_key or

eval_config.model_specs.example_weight_keys.

additional_fetches: Optional prefixes of additional tensors stored in

signature_def.inputs that should be fetched at prediction time. The

"features" and "labels" tensors are handled automatically and should not

be included. Only used if EvalSavedModel used.

blacklist_feature_fetches: Optional list of tensor names in the features

dictionary which should be excluded from the fetches request. This is

useful in scenarios where features are large (e.g. images) and can lead to

excessive memory use if stored. Only used if EvalSavedModel used.

tags: Optional model tags (e.g. 'serve' for serving or 'eval' for

EvalSavedModel).

model_name: Optional name of the model being created (should match

ModelSpecs.name). The name should only be provided if multiple models are

being evaluated.

eval_config: Eval config.

custom_model_loader: Optional custom model loader for non-TF models.

rubber_stamp: True when this run is a first run without a baseline model

while a baseline is configured, the diff thresholds will be ignored.

resource_hints: The beam resource hints to apply to the PTransform which

runs inference for this model.

backend_config: Optional configuration of backend running model inference

with *some* prediction extractors.

"""

if not eval_config:

# Default to tfma eval model unless eval

is_baseline = False

if tags and eval_constants.EVAL_TAG in tags:

model_type = constants.TFMA_EVAL

elif tags and tf.saved_model.SERVING in tags:

model_type = constants.TF_ESTIMATOR

else:

model_type = constants.TFMA_EVAL

if tags is None:

tags = [eval_constants.EVAL_TAG]

else:

model_spec = model_util.get_model_spec(eval_config, model_name)

if not model_spec:

raise ValueError('ModelSpec for model name {} not found in EvalConfig: '

'config={}'.format(model_name, eval_config))

is_baseline = model_spec.is_baseline

model_type = model_util.get_model_type(model_spec, eval_saved_model_path,

tags)

if tags is None:

# Default to serving unless tfma_eval is used.

if model_type == constants.TFMA_EVAL:

tags = [eval_constants.EVAL_TAG]

else:

tags = [tf.saved_model.SERVING]

if model_spec.example_weight_key or model_spec.example_weight_keys:

example_weight_key = (

model_spec.example_weight_key or model_spec.example_weight_keys)

if eval_config.options.HasField('include_default_metrics'):

include_default_metrics = (

eval_config.options.include_default_metrics.value)

# Backwards compatibility for legacy add_metrics_callbacks implementation.

if model_type == constants.TFMA_EVAL and eval_constants.EVAL_TAG in tags:

# PyType doesn't know about the magic exports we do in post_export_metrics.

# Additionally, the lines seem to get reordered in compilation, so we can't

# just put the disable-attr on the add_metrics_callbacks lines.

# pytype: disable=module-attr

if not add_metrics_callbacks:

add_metrics_callbacks = []

if include_default_metrics:

# Always compute example weight and example count if default metrics are

# enabled.

example_count_callback = post_export_metrics.example_count()

add_metrics_callbacks.append(example_count_callback)

if example_weight_key:

if isinstance(example_weight_key, dict):

for output_name, key in example_weight_key.items():

example_weight_callback = post_export_metrics.example_weight(

key, metric_tag=output_name)

add_metrics_callbacks.append(example_weight_callback)

else:

example_weight_callback = post_export_metrics.example_weight(

example_weight_key)

add_metrics_callbacks.append(example_weight_callback)

# pytype: enable=module-attr

model_loader = custom_model_loader

if not model_loader and model_type in constants.VALID_TF_MODEL_TYPES:

model_loader = types.ModelLoader(

construct_fn=model_util.model_construct_fn(

eval_saved_model_path=eval_saved_model_path,

add_metrics_callbacks=add_metrics_callbacks,

include_default_metrics=include_default_metrics,

additional_fetches=additional_fetches,

blacklist_feature_fetches=blacklist_feature_fetches,

model_type=model_type,

tags=tags),

tags=tags)

return types.EvalSharedModel(

model_name=model_name,

model_type=model_type,

model_path=eval_saved_model_path,

add_metrics_callbacks=add_metrics_callbacks,

include_default_metrics=include_default_metrics,

example_weight_key=example_weight_key,

additional_fetches=additional_fetches,

model_loader=model_loader,

rubber_stamp=rubber_stamp,

is_baseline=is_baseline,

resource_hints=resource_hints,

backend_config=backend_config,

if eval_config:

for spec in eval_config.slicing_specs:

if spec.slice_keys_sql:

return True

eval_shared_model: Optional[types.MaybeMultipleEvalSharedModels] = None,

eval_config: Optional[config_pb2.EvalConfig] = None,

slice_spec: Optional[List[slicer.SingleSliceSpec]] = None,

materialize: Optional[bool] = None,

tensor_adapter_config: Optional[tensor_adapter.TensorAdapterConfig] = None,

custom_predict_extractor: Optional[extractor.Extractor] = None,

config_version: Optional[int] = None) -> List[extractor.Extractor]:

"""Returns the default extractors for use in ExtractAndEvaluate.

Args:

eval_shared_model: Shared model (single-model evaluation) or list of shared

models (multi-model evaluation). Required unless the predictions are

provided alongside of the features (i.e. model-agnostic evaluations).

eval_config: Eval config.

slice_spec: Deprecated (use EvalConfig).

materialize: True to have extractors create materialized output.

tensor_adapter_config: Tensor adapter config which specifies how to obtain

tensors from the Arrow RecordBatch. If None, an attempt will be made to

create the tensors using default TensorRepresentations.

custom_predict_extractor: Optional custom predict extractor for non-TF

models.

config_version: Optional config version for this evaluation. This should not

be explicitly set by users. It is only intended to be used in cases where

the provided eval_config was generated internally, and thus not a reliable

indicator of user intent.

Raises:

NotImplementedError: If eval_config contains mixed serving and eval models.

"""

if materialize is None:

# TODO(b/172969312): Once analysis table is supported, remove defaulting

# to false unless 'analysis' is in disabled_outputs.

materialize = False

if slice_spec and eval_config:

raise ValueError('slice_spec is deprecated, only use eval_config')

if eval_config is not None:

eval_config = _update_eval_config_with_defaults(eval_config,

eval_shared_model)

tensor_representations = None

if tensor_adapter_config:

tensor_representations = tensor_adapter_config.tensor_representations

eval_shared_models = model_util.verify_and_update_eval_shared_models(

eval_shared_model

)

if _is_legacy_eval(config_version, eval_shared_models, eval_config):

# Backwards compatibility for previous add_metrics_callbacks implementation.

if not eval_config and slice_spec:

eval_config = config_pb2.EvalConfig(

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

return [

custom_predict_extractor or legacy_predict_extractor.PredictExtractor(

eval_shared_model, materialize=materialize),

slice_key_extractor.SliceKeyExtractor(

eval_config=eval_config, materialize=materialize)

]

slicing_extractors = []

if _has_sql_slices(eval_config):

slicing_extractors.append(

sql_slice_key_extractor.SqlSliceKeyExtractor(eval_config))

slicing_extractors.extend([

unbatch_extractor.UnbatchExtractor(),

slice_key_extractor.SliceKeyExtractor(

eval_config=eval_config, materialize=materialize)

])

extract_features = features_extractor.FeaturesExtractor(

eval_config=eval_config, tensor_representations=tensor_representations

)

extract_labels = labels_extractor.LabelsExtractor(eval_config=eval_config)

extract_example_weights = example_weights_extractor.ExampleWeightsExtractor(

eval_config=eval_config

)

extract_materialized_predictions = (

materialized_predictions_extractor.MaterializedPredictionsExtractor(

eval_config=eval_config

)

)

if eval_shared_model:

model_types = _model_types(eval_shared_models)

logging.info('eval_shared_models have model_types: %s', model_types)

assert model_types is not None

if (not model_types.issubset(constants.VALID_TF_MODEL_TYPES) and

not custom_predict_extractor):

raise NotImplementedError(

'either a custom_predict_extractor must be used or model type must '

'be one of: {}. evalconfig={}'.format(

str(constants.VALID_TF_MODEL_TYPES), eval_config))

if model_types == {constants.MATERIALIZED_PREDICTION}:

return [

extract_features,

extract_labels,

extract_example_weights,

extract_materialized_predictions,

] + slicing_extractors

elif model_types == {constants.TF_LITE}:

# TODO(b/163889779): Convert TFLite extractor to operate on batched

# extracts. Then we can remove the input extractor.

return [

extract_features,

transformed_features_extractor.TransformedFeaturesExtractor(

eval_config=eval_config, eval_shared_model=eval_shared_model

),

extract_labels,

extract_example_weights,

(

custom_predict_extractor

or tflite_predict_extractor.TFLitePredictExtractor(

eval_config=eval_config, eval_shared_model=eval_shared_model

)

),

] + slicing_extractors

elif constants.TF_LITE in model_types:

raise NotImplementedError(

'support for mixing tf_lite and non-tf_lite models is not '

'implemented: eval_config={}'.format(eval_config))

if model_types == {constants.TF_JS}:

return [

extract_features,

extract_labels,

extract_example_weights,

(

custom_predict_extractor

or tfjs_predict_extractor.TFJSPredictExtractor(

eval_config=eval_config, eval_shared_model=eval_shared_model

)

),

] + slicing_extractors

elif constants.TF_JS in model_types:

raise NotImplementedError(

'support for mixing tf_js and non-tf_js models is not '

'implemented: eval_config={}'.format(eval_config))

elif (

eval_config

and model_types == {constants.TFMA_EVAL}

and all(

eval_constants.EVAL_TAG in m.model_loader.tags

for m in eval_shared_models

)

):

return [

custom_predict_extractor or legacy_predict_extractor.PredictExtractor(

eval_shared_model,

materialize=materialize,

eval_config=eval_config)

] + slicing_extractors

elif (

eval_config

and constants.TFMA_EVAL in model_types

and any(

eval_constants.EVAL_TAG in m.model_loader.tags

for m in eval_shared_models

)

):

raise NotImplementedError(

'support for mixing tfma_eval and non-tfma_eval models is not '

'implemented: eval_config={}'.format(eval_config)

)

else:

extractors = [extract_features]

if not custom_predict_extractor:

extractors.append(

transformed_features_extractor.TransformedFeaturesExtractor(

eval_config=eval_config, eval_shared_model=eval_shared_model))

extractors.extend([

extract_labels,

extract_example_weights,

(

custom_predict_extractor

or predictions_extractor.PredictionsExtractor(

eval_config=eval_config, eval_shared_model=eval_shared_model

)

),

])

extractors.extend(slicing_extractors)

return extractors

else:

return [

extract_features,

extract_labels,

extract_example_weights,

extract_materialized_predictions,

eval_shared_model: Optional[types.MaybeMultipleEvalSharedModels] = None,

eval_config: Optional[config_pb2.EvalConfig] = None,

schema: Optional[schema_pb2.Schema] = None,

compute_confidence_intervals: Optional[bool] = False,

min_slice_size: int = 1,

serialize: bool = False,

random_seed_for_testing: Optional[int] = None,

config_version: Optional[int] = None) -> List[evaluator.Evaluator]:

"""Returns the default evaluators for use in ExtractAndEvaluate.

Args:

eval_shared_model: Optional shared model (single-model evaluation) or list

of shared models (multi-model evaluation). Only required if there are

metrics to be computed in-graph using the model.

eval_config: Eval config.

schema: A schema to use for customizing default evaluators.

compute_confidence_intervals: Deprecated (use eval_config).

min_slice_size: Deprecated (use eval_config).

serialize: Deprecated.

random_seed_for_testing: Provide for deterministic tests only.

config_version: Optional config version for this evaluation. This should not

be explicitly set by users. It is only intended to be used in cases where

the provided eval_config was generated internally, and thus not a reliable

indicator of user intent.

"""

disabled_outputs = []

eval_shared_models = model_util.verify_and_update_eval_shared_models(

eval_shared_model

)

if eval_config:

eval_config = _update_eval_config_with_defaults(eval_config,

eval_shared_model)

disabled_outputs = eval_config.options.disabled_outputs.values

if _model_types(eval_shared_models) == {constants.TF_LITE} or _model_types(

eval_shared_models

) == {constants.TF_JS}:

# no in-graph metrics present when tflite or tfjs is used.

if eval_shared_models:

eval_shared_models = [

v._replace(include_default_metrics=False)

for v in eval_shared_models

]

if (constants.METRICS_KEY in disabled_outputs and

constants.PLOTS_KEY in disabled_outputs and

constants.ATTRIBUTIONS_KEY in disabled_outputs):

return []

if _is_legacy_eval(config_version, eval_shared_models, eval_config):

# Backwards compatibility for previous add_metrics_callbacks implementation.

if eval_config is not None:

if eval_config.options.HasField('compute_confidence_intervals'):

compute_confidence_intervals = (

eval_config.options.compute_confidence_intervals.value)

if eval_config.options.HasField('min_slice_size'):

min_slice_size = eval_config.options.min_slice_size.value

return [

legacy_metrics_and_plots_evaluator.MetricsAndPlotsEvaluator(

eval_shared_model,

compute_confidence_intervals=compute_confidence_intervals,

min_slice_size=min_slice_size,

serialize=serialize,

random_seed_for_testing=random_seed_for_testing)

]

else:

return [

metrics_plots_and_validations_evaluator

.MetricsPlotsAndValidationsEvaluator(

eval_config=eval_config,

eval_shared_model=eval_shared_model,

schema=schema,

random_seed_for_testing=random_seed_for_testing)

output_path: Optional[str],

eval_shared_model: Optional[types.MaybeMultipleEvalSharedModels] = None,

eval_config: Optional[config_pb2.EvalConfig] = None,

display_only_data_location: Optional[str] = None,

display_only_data_file_format: Optional[str] = None,

output_file_format: str = 'tfrecord',

add_metric_callbacks: Optional[List[types.AddMetricsCallbackType]] = None

) -> List[writer.Writer]: # pylint: disable=invalid-name

"""Returns the default writers for use in WriteResults.

Note, sharding will be enabled by default if an output_file_format is

provided. Filenames will be <output_path>-SSSSS-of-NNNNN.<output_file_format>

where SSSSS is the shard number and NNNNN is the number of shards.

Args:

output_path: Output path.

eval_shared_model: Optional shared model (single-model evaluation) or list

of shared models (multi-model evaluation). Required unless the predictions

are provided alongside of the features (i.e. model-agnostic evaluations).

eval_config: Eval config for writing out config along with results. Also

used for to check for missing slices.

display_only_data_location: Optional path indicating where the examples were

read from. This is used only for display purposes - data will not actually

be read from this path.

display_only_data_file_format: Optional format of the input examples. This

is used only for display purposes.

output_file_format: File format to use when saving files. Currently only

'tfrecord' is supported.

add_metric_callbacks: Optional list of metric callbacks (if used).

"""

writers = []

if not add_metric_callbacks:

add_metric_callbacks = []

# The add_metric_callbacks are used in the metrics and plots serialization

# code to post process the metric data by calling populate_stats_and_pop.

# While both the legacy (V1) and new (V2) evaluation implementations support

# EvalSavedModels using add_metric_callbacks, this particular code is only

# required for the legacy evaluation based on the MetricsAndPlotsEvaluator.

# The V2 MetricsAndPlotsEvaluator output requires no additional processing.

# Since the V1 code only supports a single EvalSharedModel, we only set the

# add_metrics_callbacks if a dict is not passed.

if (eval_shared_model and not isinstance(eval_shared_model, dict) and

not isinstance(eval_shared_model, list)):

add_metric_callbacks = eval_shared_model.add_metrics_callbacks

eval_shared_models = model_util.verify_and_update_eval_shared_models(

eval_shared_model)

if eval_config:

model_locations = {}

for v in (eval_shared_models or [None]):

k = '' if v is None else v.model_name

model_locations[k] = ('<unknown>' if v is None or v.model_path is None

else v.model_path)

writers.append(

eval_config_writer.EvalConfigWriter(

output_path,

eval_config=eval_config,

data_location=display_only_data_location,

data_file_format=display_only_data_file_format,

model_locations=model_locations))

output_paths = {

constants.METRICS_KEY:

os.path.join(output_path, constants.METRICS_KEY),

constants.PLOTS_KEY:

os.path.join(output_path, constants.PLOTS_KEY),

constants.ATTRIBUTIONS_KEY:

os.path.join(output_path, constants.ATTRIBUTIONS_KEY),

constants.VALIDATIONS_KEY:

os.path.join(output_path, constants.VALIDATIONS_KEY)

}

writers.append(

metrics_plots_and_validations_writer.MetricsPlotsAndValidationsWriter(

output_paths=output_paths,

# Empty EvalConfig supported for backwards compatibility.

eval_config=eval_config or config_pb2.EvalConfig(),

add_metrics_callbacks=add_metric_callbacks,

output_file_format=output_file_format,

rubber_stamp=model_util.has_rubber_stamp(eval_shared_models)))

inputs: beam.pvalue.PCollection) -> beam.pvalue.PCollection:

"""Converts serialized inputs (e.g. examples) to Extracts if not already."""

def to_extracts(x: Union[bytes, str, types.Extracts]) -> types.Extracts:

result = {}

if isinstance(x, dict):

result.update(x)

else:

result[constants.INPUT_KEY] = x

return result

batched_inputs: beam.pvalue.PCollection) -> beam.pvalue.PCollection:

"""Converts Arrow RecordBatch inputs to Extracts."""

def to_extracts(

x: Union[bytes, types.Extracts, pa.RecordBatch]) -> types.Extracts:

result = {}

if isinstance(x, dict):

result.update(x)

else:

result[constants.ARROW_RECORD_BATCH_KEY] = x

return result

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

evaluators: List[evaluator.Evaluator]) -> evaluator.Evaluation:

"""Performs Extractions and Evaluations in provided order."""

# evaluation[k] = list of values for k

evaluation = {}

def update(evaluation: Dict[str, Any], new_evaluation: Dict[str, Any]):

for k, v in new_evaluation.items():

if k not in evaluation:

evaluation[k] = []

evaluation[k].append(v)

return evaluation

_ = extracts | 'TrackInputBytes' >> _TrackBytesProcessed() # pylint: disable=no-value-for-parameter

# Run evaluators that run before extraction (i.e. that only require

# the incoming input extract added by ReadInputs)

for v in evaluators:

if not v.run_after:

update(evaluation, extracts | v.stage_name >> v.ptransform)

for x in extractors:

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

for v in evaluators:

if v.run_after == x.stage_name:

update(evaluation, extracts | v.stage_name >> v.ptransform)

for v in evaluators:

if v.run_after == extractor.LAST_EXTRACTOR_STAGE_NAME:

update(evaluation, extracts | v.stage_name >> v.ptransform)

# Merge multi-valued keys if necessary.

result = {}

for k, v in evaluation.items():

if len(v) == 1:

result[k] = v[0]

continue

# Note that we assume that if a key is multivalued, its values are

# dictionaries with disjoint keys. The combined value will simply be the

# disjoint union of all the dictionaries.

result[k] = (

v

| 'FlattenEvaluationOutput(%s)' % k >> beam.Flatten()

| 'CombineEvaluationOutput(%s)' % k >> beam.CombinePerKey(

_CombineEvaluationDictionariesFn()))

"""CombineFn to combine dictionaries generated by different evaluators."""

def create_accumulator(self) -> Dict[str, Any]:

return {}

def _merge(self, accumulator: Dict[str, Any], output_dict: Dict[str,

Any]) -> None:

intersection = set(accumulator) & set(output_dict)

if intersection:

raise ValueError(

'Dictionaries generated by different evaluators should have '

'different keys, but keys %s appeared in the output of multiple '

'evaluators' % intersection)

accumulator.update(output_dict)

def add_input(self, accumulator: Dict[str, Any],

output_dict: Dict[str, Any]) -> Dict[str, Any]:

if not isinstance(output_dict, dict):

raise TypeError(

'for outputs written to by multiple evaluators, the outputs must all '

'be dictionaries, but got output of type %s, value %s' %

(type(output_dict), str(output_dict)))

self._merge(accumulator, output_dict)

return accumulator

def merge_accumulators(

self, accumulators: Iterable[Dict[str, Any]]) -> Dict[str, Any]:

accumulators = iter(accumulators)

result = next(accumulators)

for acc in accumulators:

self._merge(result, acc)

return result

def extract_output(self, accumulator: Dict[str, Any]) -> Dict[str, Any]:

evaluation_or_validation: Union[evaluator.Evaluation, validator.Validation],

writers: List[writer.Writer]) -> Dict[str, beam.PCollection]:

"""Writes Evaluation or Validation results using given writers.

Args:

evaluation_or_validation: Evaluation or Validation output.

writers: Writes to use for writing out output.

Raises:

ValueError: If Evaluation or Validation is empty.

Returns:

A dict of writer results keyed by the writer stage name.

"""

if not evaluation_or_validation: