Join the SIG TFX-Addons community and help make TFX even better!


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. The model's signature will be invoked with those tensors (matched by names). If None, an attempt will be made to create an adapter based on the model's input signature otherwise the model will be invoked with raw examples (assuming a signature of a single 1-D string tensor).

Extractor for extracting predictions.