class_docstr: A DoFn that loads the models and predicts.
class_docstr_tok: ["A","DoFn","that","loads","the","models","and","predicts","."]
docstr: Uses loaded models to make predictions on batches of data.
  
  Args:
    elem: An extract containing batched features.
  
  Yields:
    Copy of the original extracts with predictions added for each model. If
    there are multiple models, a list of dicts keyed on model names will be
    added, with each value corresponding to a prediction for a single sample.
docstr_tok: ["Uses","loaded","models","to","make","predictions","on","batches","of","data",".","Args",":","elem",":","An","extract","containing","batched","features",".","Yields",":","Copy","of","the","original","extracts","with","predictions","added","for","each","model",".","If","there","are","multiple","models",",","a","list","of","dicts","keyed","on","model","names","will","be","added",",","with","each","value","corresponding","to","a","prediction","for","a","single","sample","."]
code: def process(self, elem: tfma.Extracts) -> Iterable[tfma.Extracts]:
      """Uses loaded models to make predictions on batches of data.
  
      Args:
        elem: An extract containing batched features.
  
      Yields:
        Copy of the original extracts with predictions added for each model. If
        there are multiple models, a list of dicts keyed on model names will be
        added, with each value corresponding to a prediction for a single sample.
      """
      # Build feature and label vectors because sklearn cannot read tf.Examples.
      result = copy.copy(elem)
      features = []
      for key in self._feature_keys:
          for i, v in enumerate(result[tfma.FEATURES_KEY][key]):
              if i >= len(features):
                  features.append([])
              if len(v) == 1:
                  v = v.item()
              features[i].append(v)
      result[tfma.LABELS_KEY] = result[tfma.FEATURES_KEY][
          self._label_key
      ]
  
      # Generate predictions for each model.
      for model_name, loaded_model in self._loaded_models.items():
          preds = loaded_model.predict(features)
          if len(self._loaded_models) == 1:
              result[tfma.PREDICTIONS_KEY] = preds
          elif tfma.PREDICTIONS_KEY not in result:
              result[tfma.PREDICTIONS_KEY] = [
                  {model_name: pred} for pred in preds
              ]
          else:
              for i, pred in enumerate(preds):
                  result[tfma.PREDICTIONS_KEY][i][model_name] = pred
      yield result
  
file_path: tfx/tfx/examples/penguin/experimental/sklearn_predict_extractor.py
imports: import copy
  import os
  import pickle
  from typing import Dict, Iterable, List
  import apache_beam
  import tensorflow
  import tensorflow_model_analysis
  from tfx_bsl.tfxio import tensor_adapter
label_desc: Use image node_id 3 for calling the _TFMAPredictionDoFn obj's underlying member method code with example usage: obj.process(elem) without return types