Skip to content

Commit 6a2745c

Browse files
committed
Add TargetEncoder for scikit-learn
1 parent 76ca370 commit 6a2745c

File tree

2 files changed

+29
-2
lines changed

2 files changed

+29
-2
lines changed

stubs/sklearn/preprocessing/__init__.pyi

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ from ._data import (
1919
scale as scale,
2020
)
2121
from ._discretization import KBinsDiscretizer as KBinsDiscretizer
22-
from ._encoders import OneHotEncoder as OneHotEncoder, OrdinalEncoder as OrdinalEncoder
22+
from ._encoders import OneHotEncoder as OneHotEncoder, OrdinalEncoder as OrdinalEncoder, TargetEncoder as TargetEncoder
2323
from ._function_transformer import FunctionTransformer as FunctionTransformer
2424
from ._label import (
2525
LabelBinarizer as LabelBinarizer,
@@ -47,6 +47,7 @@ __all__ = [
4747
"RobustScaler",
4848
"SplineTransformer",
4949
"StandardScaler",
50+
"TargetEncoder",
5051
"add_dummy_feature",
5152
"PolynomialFeatures",
5253
"binarize",

stubs/sklearn/preprocessing/_encoders.pyi

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ from ..base import BaseEstimator, OneToOneFeatureMixin, TransformerMixin
1313
# Joris Van den Bossche <jorisvandenbossche@gmail.com>
1414
# License: BSD 3 clause
1515

16-
__all__ = ["OneHotEncoder", "OrdinalEncoder"]
16+
__all__ = ["OneHotEncoder", "OrdinalEncoder", "TargetEncoder"]
1717

1818
class _BaseEncoder(TransformerMixin, BaseEstimator): ...
1919

@@ -67,3 +67,29 @@ class OrdinalEncoder(OneToOneFeatureMixin, _BaseEncoder):
6767
def fit(self, X: MatrixLike, y: Series | None = None) -> Self: ...
6868
def transform(self, X: MatrixLike) -> ndarray: ...
6969
def inverse_transform(self, X: MatrixLike) -> ndarray: ...
70+
71+
class TargetEncoder(OneToOneFeatureMixin, _BaseEncoder):
72+
feature_names_in_: ndarray = ...
73+
n_features_in_: int = ...
74+
categories_: list[ndarray] = ...
75+
encodings_: list[ndarray] = ...
76+
target_type_: str = ...
77+
target_mean_: float = ...
78+
classes_: ndarray | None = ...
79+
80+
_parameter_constraints: ClassVar[dict] = ...
81+
82+
def __init__(
83+
self,
84+
*,
85+
categories: Sequence[ArrayLike] | Literal["auto"] = "auto",
86+
target_type: Literal["auto", "continuous", "binary", "multiclass"] = "auto",
87+
smooth: Literal["auto"] | float = "auto",
88+
cv: int = 5,
89+
shuffle: bool = True,
90+
random_state: Int | None = None,
91+
) -> None: ...
92+
def fit(self, X: MatrixLike, y: ArrayLike) -> Self: ...
93+
def transform(self, X: MatrixLike) -> ndarray: ...
94+
def fit_transform(self, X: MatrixLike, y: ArrayLike) -> ndarray: ...
95+
def get_feature_names_out(self, input_features: None | ArrayLike = None) -> ndarray: ...

0 commit comments

Comments
 (0)