class_docstr: A DoFn that loads the models and predicts.
class_docstr_tok: ["A","DoFn","that","loads","the","models","and","predicts","."]
code: def setup(self):
      super().setup()
      self._feature_keys = None
      self._label_key = None
      for loaded_model in self._loaded_models.values():
          if self._feature_keys and self._label_key:
              assert self._feature_keys == loaded_model.feature_keys, (
                  f"Features mismatch in loaded models. Expected {self._feature_keys}"
                  f", got {loaded_model.feature_keys} instead."
              )
              assert self._label_key == loaded_model.label_key, (
                  f'Label mismatch in loaded models. Expected "{self._label_key}"'
                  f', got "{loaded_model.label_key}" instead.'
              )
          elif loaded_model.feature_keys and loaded_model.label_key:
              self._feature_keys = loaded_model.feature_keys
              self._label_key = loaded_model.label_key
          else:
              raise ValueError(
                  "Missing feature or label keys in loaded model."
              )
  
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 2 for calling the _TFMAPredictionDoFn obj's underlying member method code with example usage: obj.setup() without return types