ModelOutputType is now a Flag instead of regular enum. Combinations of the base flags are provided for all of the previous output types for convenience. All checks in the code now use the basic flags and not the complex types.

Signed-off-by: abigailt <abigailt@il.ibm.com>
This commit is contained in:
abigailt 2024-07-03 13:29:37 +03:00
parent 2895b40f05
commit 367cae679b
10 changed files with 126 additions and 100 deletions

View file

@ -16,7 +16,8 @@ from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from apt.utils.datasets import ArrayDataset, DATA_PANDAS_NUMPY_TYPE
from apt.utils.models import Model, SklearnRegressor, ModelOutputType, SklearnClassifier
from apt.utils.models import Model, SklearnRegressor, ModelOutputType, SklearnClassifier, \
CLASSIFIER_SINGLE_OUTPUT_CLASS_PROBABILITIES
@dataclass
@ -96,7 +97,7 @@ class GeneralizeToRepresentative(BaseEstimator, MetaEstimatorMixin, TransformerM
# model output type is not critical as it only affects computation of nb_classes, which is in any case
# the same currently for single and multi output probabilities.
self.estimator = SklearnClassifier(estimator,
ModelOutputType.CLASSIFIER_SINGLE_OUTPUT_CLASS_PROBABILITIES)
CLASSIFIER_SINGLE_OUTPUT_CLASS_PROBABILITIES)
self.target_accuracy = target_accuracy
self.cells = cells
self.categorical_features = []