diff --git a/rubicon_ml/schema/registry.py b/rubicon_ml/schema/registry.py index 2dc69955..468da5d3 100644 --- a/rubicon_ml/schema/registry.py +++ b/rubicon_ml/schema/registry.py @@ -15,6 +15,9 @@ "h2o__H2ORandomForestEstimator": lambda: _load_schema( os.path.join("schema", "h2o__H2ORandomForestEstimator.yaml") ), + "h2o__H2OTargetEncoderEstimator": lambda: _load_schema( + os.path.join("schema", "h2o__H2OTargetEncoderEstimator.yaml") + ), "lightgbm__LGBMModel": lambda: _load_schema(os.path.join("schema", "lightgbm__LGBMModel.yaml")), "lightgbm__LGBMClassifier": lambda: _load_schema( os.path.join("schema", "lightgbm__LGBMClassifier.yaml") diff --git a/rubicon_ml/schema/schema/h2o__H2OTargetEncoderEstimator.yaml b/rubicon_ml/schema/schema/h2o__H2OTargetEncoderEstimator.yaml new file mode 100644 index 00000000..1703cd05 --- /dev/null +++ b/rubicon_ml/schema/schema/h2o__H2OTargetEncoderEstimator.yaml @@ -0,0 +1,32 @@ +name: h2o__H2OTargetEncoderEstimator +version: 1.0.0 + +compatibility: + lightgbm: + max_version: + min_version: 3.44.0.1 +docs_url: https://docs.h2o.ai/h2o/latest-stable/h2o-py/docs/modeling.html#h2otargetencoderestimator + +parameters: + - name: blending + value_attr: blending + - name: columns_to_encode + value_attr: columns_to_encode + - name: data_leakage_handling + value_attr: data_leakage_handling + - name: fold_column + value_attr: fold_column + - name: ignored_columns + value_attr: ignored_columns + - name: inflection_point + value_attr: inflection_point + - name: keep_original_categorical_columns + value_attr: keep_original_categorical_columns + - name: noise + value_attr: noise + - name: response_column + value_attr: response_column + - name: seed + value_attr: seed + - name: smoothing + value_attr: smoothing diff --git a/tests/integration/test_schema.py b/tests/integration/test_schema.py index c6051b6f..1e95ee11 100644 --- a/tests/integration/test_schema.py +++ b/tests/integration/test_schema.py @@ -5,6 +5,7 @@ from h2o.estimators.gbm import H2OGradientBoostingEstimator from h2o.estimators.glm import H2OGeneralizedLinearEstimator from h2o.estimators.random_forest import H2ORandomForestEstimator +from h2o.estimators.targetencoder import H2OTargetEncoderEstimator from lightgbm import LGBMClassifier, LGBMRegressor from sklearn.ensemble import RandomForestClassifier from xgboost import XGBClassifier, XGBRegressor @@ -14,6 +15,7 @@ H2OGeneralizedLinearEstimator, H2OGradientBoostingEstimator, H2ORandomForestEstimator, + H2OTargetEncoderEstimator, ] PANDAS_SCHEMA_CLS = [ LGBMClassifier,