tfma.extractors.PredictionsExtractor

Creates an extractor for performing predictions over a batch.

The extractor runs in two modes:

1) If one or more EvalSharedModels are provided

The extractor's PTransform loads and runs the serving saved_model(s) against every extract yielding a copy of the incoming extracts with an additional extract added for the predictions keyed by tfma.PREDICTIONS_KEY. The model inputs are searched for under tfma.FEATURES_KEY (keras only) or tfma.INPUT_KEY (if tfma.FEATURES_KEY is not set or the model is non-keras). If multiple models are used the predictions will be stored in a dict keyed by model name.

2) If no EvalSharedModels are provided

The extractor's PTransform uses the config's ModelSpec.prediction_key(s) to lookup the associated prediction values stored as features under the tfma.FEATURES_KEY in extracts. The resulting values are then added to the extracts under the key tfma.PREDICTIONS_KEY.

Note that the use of a prediction_key in the ModelSpecs serve two use cases: (a) as a key into the dict of predictions output (option 1) (b) as the key for a pre-computed prediction stored as a feature (option 2)

eval_config Eval config.
eval_shared_model Shared model (single-model evaluation) or list of shared models (multi-model evaluation) or None (predictions obtained from features).
tensor_adapter_config Tensor adapter config which specifies how to obtain tensors from the Arrow RecordBatch. If None, we feed the raw examples to the model.

Extractor for extracting predictions.