Repository: onnx/sklearn-onnx
Pull request #1137 (2024-11-07 16:12:08)

feat: add support to sklearn TargetEncoder


This PR implements a converter and a shape calculator for the TargetEncoder class introduced in Scikit-learn 1.5. The code follows much of the implementation of the converter for Ordinal Encoder.

A partial suit of tests is already implemented, but there is at least a couple of additional tests that I would like to add (missing values and using the smooth parameter from sklearn, even though I think it shouldn't matter).

xadupre: Thanks for the contribution. One line should be removed. Everything else looks good.
boccaff: Thanks for the comments @xadupre. I've removed the line, and solved a couple of the CodeQL suggestions (removed an unused import and an unused variable). The rest of the CodeQL suggested changes would diverge from the other implementations. For the `.assertTrue` I can just follow the suggestion, but it would diverge from other tests. Is it ok? For the `except: pass`, maybe we can add a warning (example above on the respective CodeQL comment).


khoover: Hi @boccaff, @xadupre, I'm wondering whether there's a timeline on getting this merged in. My team's been interested in using ONNX, but almost all of our models use target encoding.
xadupre: The PR looks ok. I'll merge it when the CI test pass.
xadupre: You should write ``convert_sklearn or to_onnx (..., target_opset={"": 18, "ai.onnx.ml": 3})``. That will fix the opset.
boccaff: I've noticed that the converter for the `OrdinalEncoder` sets `op_version=2` when [adding nodes](https://github.com/onnx/sklearn-onnx/blob/main/skl2onnx/operator_converters/ordinal_encoder.py#L133). Would that be preferable? If so, should I [raise](https://github.com/onnx/sklearn-onnx/blob/a245fb0962b6e0d91e1d167ef18912db6883fd22/skl2onnx/operator_converters/label_encoder.py#L26) an error as it is done on `LabelEncoder`?

PS: should we also raise on OrdinalEncoder? I could follow up on a different issue.
xadupre: 2 or 3 should not matter. If you set 3, onnx won't find any version for opset 3 so it chooses the latest available, so 2.
boccaff: I'll try to set up a ubuntu docker image and the failing config (sklearn 1.3.2 and py 3.11) to debug this failing test.
I had local tests running ok, but on more recent env (sklearn 1.6.1 and python 3.13.1).

py 3.12 and sklearn==1.4.2 appears [ok](https://github.com/onnx/sklearn-onnx/actions/runs/13424697949/job/37527885778?pr=1137#step:12:133), 1.5.2 [also](https://github.com/onnx/sklearn-onnx/actions/runs/13424697949/job/37527884671?pr=1137#step:12:133) and [1.6.2](https://github.com/onnx/sklearn-onnx/actions/runs/13424697949/job/37527883726?pr=1137#step:12:133).

xadupre: According to this https://scikit-learn.org/1.3/modules/generated/sklearn.preprocessing.TargetEncoder.html, the attribute does not exist in that version. You can either choose to retreive the information from another attribute or say the converter only works after sklearn>=1.4. Since it is a new converter, I don't see any issue with that. We may address older versions if other users post an issue on github. So feel free to raise an exception if the attribute is missing giving the information the user needs to understand why it does not work.
boccaff: Opted for checking only `target_type_` that now should be only 'continuos' or 'binary'.

Patch:
diff --git a/skl2onnx/_supported_operators.py b/skl2onnx/_supported_operators.py
index 7343bb1..77fa184 100644
--- a/skl2onnx/_supported_operators.py
+++ b/skl2onnx/_supported_operators.py
@@ -239,6 +239,12 @@ try:
 except ImportError:
     # not available in 0.19
     KBinsDiscretizer = None
+try:
+    from sklearn.preprocessing import TargetEncoder
+except ImportError:
+    # Not available in scikit-learn < 1.3
+    TargetEncoder = None
+
 from sklearn.preprocessing import (
     LabelBinarizer,
     LabelEncoder,
@@ -511,6 +517,7 @@ def build_sklearn_operator_name_map():
             RidgeClassifierCV: "SklearnLinearClassifier",
             SGDRegressor: "SklearnLinearRegressor",
             StandardScaler: "SklearnScaler",
+            TargetEncoder: "SklearnTargetEncoder",
             TheilSenRegressor: "SklearnLinearRegressor",
         }
     )
diff --git a/skl2onnx/operator_converters/__init__.py b/skl2onnx/operator_converters/__init__.py
index 0a9dadf..f5a7885 100644
--- a/skl2onnx/operator_converters/__init__.py
+++ b/skl2onnx/operator_converters/__init__.py
@@ -61,6 +61,7 @@ from . import sgd_classifier
 from . import sgd_oneclass_svm
 from . import stacking
 from . import support_vector_machines
+from . import target_encoder
 from . import text_vectoriser
 from . import tfidf_transformer
 from . import tfidf_vectoriser
@@ -128,6 +129,7 @@ __all__ = [
     sgd_oneclass_svm,
     stacking,
     support_vector_machines,
+    target_encoder,
     text_vectoriser,
     tfidf_transformer,
     tfidf_vectoriser,
diff --git a/skl2onnx/operator_converters/target_encoder.py b/skl2onnx/operator_converters/target_encoder.py
new file mode 100644
index 0000000..9d71eb4
--- /dev/null
+++ b/skl2onnx/operator_converters/target_encoder.py
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+import numpy as np
+
+from ..common._apply_operation import apply_cast, apply_concat, apply_reshape
+from ..common._container import ModelComponentContainer
+from ..common.data_types import (
+    FloatTensorType,
+    Int64TensorType,
+)
+from ..common._registration import register_converter
+from ..common._topology import Scope, Operator
+from ..proto import onnx_proto
+
+
+def convert_sklearn_target_encoder(
+    scope: Scope, operator: Operator, container: ModelComponentContainer
+):
+    op = operator.raw_operator
+    result = []
+    input_idx = 0
+    dimension_idx = 0
+
+    if op.target_type_ not in ("binary", "continuous"):
+        raise NotImplementedError(
+            "Current implementation of the converter only support TargetEncoder for"
+            " binary classification or 1d regression (sklearn target types binary"
+            " or continuous). See scikit-learn type_of_target documentation for details."
+        )
+    for categories, encodings in zip(op.categories_, op.encodings_):
+        if len(categories) == 0:
+            continue
+
+        current_input = operator.inputs[input_idx]
+        if current_input.get_second_dimension() == 1:
+            feature_column = current_input
+            input_idx += 1
+        else:
+            index_name = scope.get_unique_variable_name("index")
+            container.add_initializer(
+                index_name, onnx_proto.TensorProto.INT64, [], [dimension_idx]
+            )
+
+            feature_column = scope.declare_local_variable(
+                "feature_column",
+                current_input.type.__class__([current_input.get_first_dimension(), 1]),
+            )
+
+            container.add_node(
+                "ArrayFeatureExtractor",
+                [current_input.onnx_name, index_name],
+                feature_column.onnx_name,
+                op_domain="ai.onnx.ml",
+                name=scope.get_unique_operator_name("ArrayFeatureExtractor"),
+            )
+
+            dimension_idx += 1
+            if dimension_idx == current_input.get_second_dimension():
+                dimension_idx = 0
+                input_idx += 1
+
+        attrs = {"name": scope.get_unique_operator_name("LabelEncoder")}
+        if isinstance(feature_column.type, FloatTensorType):
+            attrs["keys_floats"] = np.array(
+                [float(s) for s in categories], dtype=np.float32
+            )
+        elif isinstance(feature_column.type, Int64TensorType):
+            attrs["keys_int64s"] = np.array(
+                [int(s) for s in categories], dtype=np.int64
+            )
+        else:
+            attrs["keys_strings"] = np.array(
+                [str(s).encode("utf-8") for s in categories]
+            )
+        attrs["values_floats"] = encodings
+        attrs["default_float"] = op.target_mean_
+
+        result.append(scope.get_unique_variable_name("ordinal_output"))
+        label_encoder_output = scope.get_unique_variable_name("label_encoder")
+
+        container.add_node(
+            "LabelEncoder",
+            feature_column.onnx_name,
+            label_encoder_output,
+            op_domain="ai.onnx.ml",
+            op_version=2,
+            **attrs,
+        )
+        apply_reshape(
+            scope,
+            label_encoder_output,
+            result[-1],
+            container,
+            desired_shape=(-1, 1),
+        )
+
+    concat_result_name = scope.get_unique_variable_name("concat_result")
+    apply_concat(scope, result, concat_result_name, container, axis=1)
+    apply_cast(
+        scope,
+        concat_result_name,
+        operator.output_full_names,
+        container,
+        to=onnx_proto.TensorProto.FLOAT,
+    )
+
+
+register_converter("SklearnTargetEncoder", convert_sklearn_target_encoder)
diff --git a/skl2onnx/shape_calculators/__init__.py b/skl2onnx/shape_calculators/__init__.py
index 71f9042..3cea54b 100644
--- a/skl2onnx/shape_calculators/__init__.py
+++ b/skl2onnx/shape_calculators/__init__.py
@@ -48,6 +48,7 @@ from . import sequence
 from . import sgd_oneclass_svm
 from . import svd
 from . import support_vector_machines
+from . import target_encoder
 from . import text_vectorizer
 from . import tuned_threshold_classifier
 from . import tfidf_transformer
@@ -101,6 +102,7 @@ __all__ = [
     sgd_oneclass_svm,
     svd,
     support_vector_machines,
+    target_encoder,
     text_vectorizer,
     tfidf_transformer,
     tuned_threshold_classifier,
diff --git a/skl2onnx/shape_calculators/target_encoder.py b/skl2onnx/shape_calculators/target_encoder.py
new file mode 100644
index 0000000..89ec73c
--- /dev/null
+++ b/skl2onnx/shape_calculators/target_encoder.py
@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+from ..common._registration import register_shape_calculator
+from ..common.data_types import FloatTensorType
+from ..common.data_types import Int64TensorType, StringTensorType
+from ..common.utils import check_input_and_output_numbers
+from ..common.utils import check_input_and_output_types
+
+
+def calculate_sklearn_target_encoder_output_shapes(operator):
+    """
+    This function just copy the input shape to the output because target
+    encoder only alters input features' values, not their shape.
+    """
+    check_input_and_output_numbers(operator, output_count_range=1)
+    check_input_and_output_types(
+        operator, good_input_types=[FloatTensorType, Int64TensorType, StringTensorType]
+    )
+
+    N = operator.inputs[0].get_first_dimension()
+    shape = [N, len(operator.raw_operator.categories_)]
+
+    operator.outputs[0].type = FloatTensorType(shape=shape)
+
+
+register_shape_calculator(
+    "SklearnTargetEncoder", calculate_sklearn_target_encoder_output_shapes
+)


Test patch:
diff --git a/tests/test_sklearn_target_encoder_converter.py b/tests/test_sklearn_target_encoder_converter.py
new file mode 100644
index 0000000..4846def
--- /dev/null
+++ b/tests/test_sklearn_target_encoder_converter.py
@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests scikit-learn's TargetEncoder converter."""
+
+from skl2onnx import convert_sklearn, to_onnx
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.pipeline import make_pipeline
+from sklearn.compose import ColumnTransformer
+from sklearn import __version__ as sklearn_version
+import onnxruntime
+import pandas as pd
+from numpy.testing import assert_almost_equal
+import unittest
+import packaging.version as pv
+import numpy as np
+from onnxruntime import __version__ as ort_version
+from skl2onnx.common.data_types import (
+    Int64TensorType,
+    StringTensorType,
+)
+from test_utils import dump_data_and_model, TARGET_OPSET
+
+try:
+    from sklearn.preprocessing import TargetEncoder
+except ImportError:
+    # Not available for scikit-learn < 1.3
+    pass
+
+
+ort_version = ".".join(ort_version.split(".")[:2])
+
+
+def target_encoder_support():
+    # pv.Version does not work with development versions
+    vers = ".".join(sklearn_version.split(".")[:2])
+    if pv.Version(vers) < pv.Version("1.3.0"):
+        return False
+    if pv.Version(onnxruntime.__version__) < pv.Version("0.3.0"):
+        return False
+    return pv.Version(vers) >= pv.Version("1.3.0")
+
+
+def is_scikit_13():
+    vers = ".".join(sklearn_version.split(".")[:2])
+    return pv.Version(vers) < pv.Version("1.4.0")
+
+
+def set_output_support():
+    vers = ".".join(sklearn_version.split(".")[:2])
+    return pv.Version(vers) >= pv.Version("1.2")
+
+
+class TestSklearnTargetEncoderConverter(unittest.TestCase):
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    def test_model_target_encoder(self):
+        model = TargetEncoder()
+        X = np.array(["str3", "str2", "str0", "str1", "str3"]).reshape(-1, 1)
+        y = np.array([0.0, 1.0, 1.0, 0.0, 1.0])
+        model.fit(X, y)
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn target encoder",
+            [("input", StringTensorType([None, X.shape[1]]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertTrue(model_onnx is not None)
+        self.assertTrue(model_onnx.graph.node is not None)
+        if model_onnx.ir_version >= 7 and TARGET_OPSET < 12:
+            raise AssertionError("Incompatbilities")
+        dump_data_and_model(X, model, model_onnx, basename="SklearnTargetEncoder")
+
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    def test_model_target_encoder_int(self):
+        model = TargetEncoder()
+        X = np.array([0, 0, 1, 0, 0, 1, 1], dtype=np.int64).reshape(-1, 1)
+        y = np.array([1, 1, 1, 0, 0, 0, 0], dtype=np.int64)
+        X_test = np.array([0, 1, 2, 1, 0], dtype=np.int64).reshape(-1, 1)
+
+        model.fit(X, y)
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn label encoder",
+            [("input", Int64TensorType([None, X.shape[1]]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertTrue(model_onnx is not None)
+        self.assertTrue(model_onnx.graph.node is not None)
+        if model_onnx.ir_version >= 7 and TARGET_OPSET < 12:
+            raise AssertionError("Incompatbilities")
+        dump_data_and_model(
+            X_test, model, model_onnx, basename="SklearnTargetEncoderInt"
+        )
+
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    def test_target_encoder_twocats(self):
+        data = [["cat2"], ["cat1"]]
+        label = [0, 1]
+        model = TargetEncoder(categories="auto")
+        model.fit(data, label)
+        inputs = [("input1", StringTensorType([None, 1]))]
+        model_onnx = convert_sklearn(
+            model, "Target encoder two string cats", inputs, target_opset=TARGET_OPSET
+        )
+        self.assertTrue(model_onnx is not None)
+        dump_data_and_model(
+            data, model, model_onnx, basename="SklearnTargetEncoderTwoStringCat"
+        )
+
+    @unittest.skipIf(
+        not set_output_support(),
+        reason="'ColumnTransformer' object has no attribute 'set_output'",
+    )
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    def test_target_encoder_pipeline_f32(self):
+        from onnxruntime import InferenceSession
+
+        data = pd.DataFrame({"cat": ["cat2", "cat1"] * 10, "num": [0, 1, 1, 0] * 5})
+        data["num"] = data["num"].astype(np.float32)
+        y = np.array([0, 1, 0, 1] * 5, dtype=np.float32)
+        # target encoder uses cross-fitting and have cv=5 as default, which
+        # caused some folds to have constant y.
+        preprocessor = ColumnTransformer(
+            transformers=[
+                ("cat", TargetEncoder(cv=2), ["cat"]),
+                ("num", "passthrough", ["num"]),
+            ],
+            sparse_threshold=1,
+            verbose_feature_names_out=False,
+        ).set_output(transform="pandas")
+        model = make_pipeline(
+            preprocessor, RandomForestRegressor(n_estimators=3, max_depth=2)
+        )
+        model.fit(data, y)
+        expected = model.predict(data)
+        model_onnx = to_onnx(model, data, target_opset=TARGET_OPSET)
+        sess = InferenceSession(
+            model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        got = sess.run(
+            None,
+            {
+                "cat": data["cat"].values.reshape((-1, 1)),
+                "num": data["num"].values.reshape((-1, 1)),
+            },
+        )
+        assert_almost_equal(expected, got[0].ravel())
+
+    @unittest.skipIf(
+        not set_output_support(),
+        reason="'ColumnTransformer' object has no attribute 'set_output'",
+    )
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    def test_target_encoder_pipeline_string_int64(self):
+        from onnxruntime import InferenceSession
+
+        data = pd.DataFrame(
+            {
+                "C1": ["cat2", "cat1", "cat3"] * 10,
+                "C2": [1, 0, 1] * 10,
+                "num": [0, 1, 1] * 10,
+            }
+        )
+        data["num"] = data["num"].astype(np.float32)
+        data["C2"] = data["C2"].astype(np.int64)
+        y = np.array([0, 1, 0, 1, 0, 1] * 5, dtype=np.int64)
+        preprocessor = ColumnTransformer(
+            transformers=[
+                ("cat", TargetEncoder(cv=2), ["C1", "C2"]),
+                ("num", "passthrough", ["num"]),
+            ],
+            sparse_threshold=1,
+            verbose_feature_names_out=False,
+        ).set_output(transform="pandas")
+        model = make_pipeline(
+            preprocessor, RandomForestRegressor(n_estimators=3, max_depth=2)
+        )
+        model.fit(data, y)
+        expected = model.predict(data)
+        model_onnx = to_onnx(model, data, target_opset=TARGET_OPSET)
+        sess = InferenceSession(
+            model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        got = sess.run(
+            None,
+            {
+                "C1": data["C1"].values.reshape((-1, 1)),
+                "C2": data["C2"].values.reshape((-1, 1)),
+                "num": data["num"].values.reshape((-1, 1)),
+            },
+        )
+        assert_almost_equal(expected, got[0].ravel())
+
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    @unittest.skipIf(TARGET_OPSET < 9, reason="not available")
+    def test_target_encoder_mixed_string_int_pandas(self):
+        col1 = "col1"
+        col2 = "col2"
+        col3 = "col3"
+        data_pd = pd.DataFrame(
+            {
+                col1: np.array(["c0.4", "c1.4", "c0.2", "c0.2", "c0.2", "c0.2"]),
+                col2: np.array(["c0.2", "c1.2", "c2.2", "c2.2", "c2.2", "c2.2"]),
+                col3: np.array([3, 0, 1, 1, 1, 1]),
+            }
+        )
+        data_label = np.array([0, 0, 1, 1, 1, 0])
+        test_pd = pd.DataFrame(
+            {
+                col1: np.array(["c0.2"]),
+                col2: np.array(["c2.2"]),
+                col3: np.array([1]),
+            }
+        )
+        model = TargetEncoder()
+        model.fit(data_pd, data_label)
+        inputs = [
+            ("input1", StringTensorType([None, 2])),
+            ("input2", Int64TensorType([None, 1])),
+        ]
+        model_onnx = convert_sklearn(
+            model, "ordinal encoder", inputs, target_opset=TARGET_OPSET
+        )
+        self.assertIsNotNone(model_onnx)
+        dump_data_and_model(
+            test_pd,
+            model,
+            model_onnx,
+            basename="SklearnTargetEncoderMixedStringIntPandas",
+        )
+
+    @unittest.skipIf(
+        not target_encoder_support(),
+        reason="TargetEncoder was not available before 1.3",
+    )
+    @unittest.skipIf(TARGET_OPSET < 9, reason="not available")
+    @unittest.skipIf(
+        is_scikit_13(),
+        reason="scikit-learn 1.3 does not allow multiclass TargetEncoder",
+    )
+    def test_target_encoder_multiclass_assertion(self):
+        model = TargetEncoder()
+        X = np.array([0, 0, 1, 0, 0, 1, 1], dtype=np.int64).reshape(-1, 1)
+        y = np.array([0, 1, 2, 0, 1, 2, 0], dtype=np.int64)
+
+        model.fit(X, y)
+        with self.assertRaises(NotImplementedError):
+            # sklearn 1.3 does not allow for multiclass TargetEncoder.
+            # for >= 1.4 we must ensure that the output is binary or continuous
+            convert_sklearn(
+                model,
+                "scikit-learn target encoder",
+                [("input", Int64TensorType([None, X.shape[1]]))],
+                target_opset=TARGET_OPSET,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)