tfma.extractors.InputExtractor

View source on GitHub

Creates an extractor for extracting features, labels, and example weights.

The extractor's PTransform parses tf.train.Example protos stored under the tfma.INPUT_KEY in the incoming extracts and adds the resulting features, labels, and example weights to the extracts under the keys tfma.FEATURES_KEY, tfma.LABELS_KEY, and tfma.EXAMPLE_WEIGHTS_KEY. If the eval_config contains a prediction_key and a corresponding key is found in the parse example, then predictions will also be extracted and stored under the tfma.PREDICTIONS_KEY. Any extracts that already exist will be merged with the values parsed by this extractor with this extractor's values taking precedence when duplicate keys are detected.

Note that the use of a prediction_key in an eval_config serves two use cases: (1) as a key into the dict of predictions output by predict extractor (2) as the key for a pre-computed prediction stored as a feature. The InputExtractor can be used to handle case (2). These cases are meant to be exclusive (i.e. if approach (2) is used then a predict extractor would not be configured and if (1) is used then a key matching the predictons would not be stored in the features). However, if a feature key happens to match the same name as the prediction output key then both paths may be executed. In this case, the value stored here will be replaced by the predict extractor (though it will still be popped from the features).

eval_config Eval config.

Extractor for extracting features, labels, and example weights inputs.