def setup()

in tfx_addons/xgboost_evaluator/xgboost_predict_extractor.py [0:0]

• 24 lines of code
• 10 McCabe index (conditional complexity)

  def setup(self):
    """Function that accesses and saves the chosen feature keys and label key."""

    # models are loaded and stored into self._loaded_models below
    super().setup()
    self._feature_keys = None
    self._label_key = None

    # check to see whether each of these loaded models has a corresponding model spec
    if self._eval_config:
      label_config = self.extract_model_specs()
      for name in self._loaded_models:
        if name not in label_config:
          raise ValueError(f"Missing model spec for loaded model {name}.")

    for name, loaded_model in self._loaded_models.items():
      feature_keys = loaded_model.feature_names
      if self._feature_keys and self._label_key:
        assert self._feature_keys == feature_keys, (
            f'Features mismatch in loaded models. Expected {self._feature_keys}'
            f', got {feature_keys} instead.')
        assert self._label_key == label_config[name], (
            f'Label mismatch in loaded models. Expected "{self._label_key}"'
            f', got "{label_config[name]}" instead.')
      elif feature_keys and label_config[name]:
        self._feature_keys = feature_keys
        self._label_key = label_config[name]
      else:
        raise ValueError(
            f'Missing feature or label keys in loaded model {name}.')