Fix Issue 2114: Accept estimators that don't inherit from sklearn.BaseEstimator in tabular_pipeline

CHANGES.rst

@@ -47,6 +47,14 @@ Changes
   :pr:`2222` by :user:`Ashwin V. Mohanan <ashwinvis>`, with guidance from
   :user:`Jérôme Dockès <jeromedockes>`.
 
+- Made the following changes to :func:`tabular_pipeline`:
+  - Estimators are no longer required to inherit from :class:`sklearn.BaseEstimator`.
+    Instead, scikit-learn compatibility check is based on presence of the methods:
+    `get_params`, `set_params`, `fit`, `predict`.
+  - Requirement for special treatment for tree ensemble/HGBT models is determined
+    based on class name substring matching, rather than exact type matching.
+
+  :pr:`2225` by :user:`Laurence Dyer <ljdyer>`. 
 
 Bugfixes
 --------

doc/modules/default_wrangling/tabular_pipeline.rst

@@ -15,7 +15,7 @@
 Building robust ML baselines with |tabular_pipeline|
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-The |tabular_pipeline| is a function that, given a scikit-learn estimator,
+The |tabular_pipeline| is a function that, given a scikit-learn compatible estimator,
 returns a full scikit-learn |Pipeline| that contains a |TableVectorizer|
 followed by the given estimator.
 If the estimator is a linear model (e.g., ``Ridge``, ``LogisticRegression``),
@@ -50,8 +50,8 @@ problems, but may not beat properly tuned ad-hoc pipelines.
    :widths: 25 25 25 25
 
    * - Parameter
-     - ``RandomForest`` models
-     - ``HistGradientBoosting`` models
+     - Tree ensemble models (e.g. ``RandomForest``)
+     - ``HistGradientBoosting`` models 
      - Linear models and others
    * - Low-cardinality encoder
      - :class:`~sklearn.preprocessing.OrdinalEncoder`

skrub/_tabular_pipeline.py

@@ -1,5 +1,4 @@
 from sklearn import ensemble
-from sklearn.base import BaseEstimator
 from sklearn.impute import SimpleImputer
 from sklearn.pipeline import make_pipeline
 from sklearn.preprocessing import OrdinalEncoder
@@ -11,22 +10,34 @@
 from ._table_vectorizer import TableVectorizer
 from ._to_categorical import ToCategorical
 
-_HGBT_CLASSES = (
-    ensemble.HistGradientBoostingClassifier,
-    ensemble.HistGradientBoostingRegressor,
-)
-_TREE_ENSEMBLE_CLASSES = (
-    ensemble.HistGradientBoostingClassifier,
-    ensemble.HistGradientBoostingRegressor,
-    ensemble.RandomForestClassifier,
-    ensemble.RandomForestRegressor,
+_HGBT_CLASS_NAME_SUBSTRINGS = ("HistGradientBoosting",)
+_TREE_ENSEMBLE_CLASS_NAME_SUBSTRINGS = (
+    "HistGradientBoosting",
+    "RandomForest",
+    "XGB",
+    "LGBM",
 )
 
 
+def is_scikit_learn_compatible_estimator(estimator) -> tuple[bool, str | None]:
+    """Determine whether a candidate object is a valid scikit learn-compatiable
+    estimator. Return True or False, plus an optional string stating the failure
+    reason."""
+
+    REQUIRED_METHOD_NAMES = ["get_params", "set_params", "fit", "predict"]
+    for method_name in REQUIRED_METHOD_NAMES:
+        if not hasattr(estimator, method_name):
+            return False, f"The estimator must have a {method_name} attribute."
+    for method_name in REQUIRED_METHOD_NAMES:
+        if not callable(getattr(estimator, method_name)):
+            return False, f"The estimator's {method_name} attribute must be callable."
+    return True, None
+
+
 def tabular_pipeline(estimator, *, n_jobs=None):
     """Get a simple machine-learning pipeline for tabular data.
 
-    Given either a scikit-learn estimator or one of the special-cased strings
+    Given either a scikit-learn compatible estimator or one of the special-cased strings
     ``'regressor'``, ``'regression'``, ``'classifier'``, ``'classification'``, this
     function creates a scikit-learn pipeline that extracts numeric features, imputes
     missing values and scales the data if necessary, then applies the estimator.
@@ -47,7 +58,9 @@ def tabular_pipeline(estimator, *, n_jobs=None):
 
     Parameters
     ----------
-    estimator : {"regressor", "regression", "classifier", "classification"} or sklearn.base.BaseEstimator
+    estimator : {"regressor", "regression", "classifier", "classification"} or scikit-learn
+        compatible estimator
+
         The estimator to use as the final step in the pipeline. Based on the type of
         estimator, the previous preprocessing steps and their respective parameters are
         chosen. The possible values are:
@@ -58,7 +71,8 @@ def tabular_pipeline(estimator, *, n_jobs=None):
         - ``'classifier'`` or ``'classification'``: a
           :obj:`~sklearn.ensemble.HistGradientBoostingClassifier` is used as the final
           step;
-        - a scikit-learn estimator: the provided estimator is used as the final step.
+        - a scikit-learn compatible estimator: the provided estimator is used as the final
+          step.
 
     n_jobs : int, default=None
         Number of jobs to run in parallel in the :obj:`TableVectorizer` step. ``None``
@@ -240,31 +254,45 @@ def tabular_pipeline(estimator, *, n_jobs=None):
             "If ``estimator`` is a string it should be 'regressor', 'regression',"
             " 'classifier' or 'classification'."
         )
-    if isinstance(estimator, type) and issubclass(estimator, BaseEstimator):
+
+    if isinstance(estimator, type):
         raise TypeError(
-            "tabular_pipeline expects a scikit-learn estimator as its first"
-            f" argument. Pass an instance of {estimator.__name__} rather than the class"
-            " itself."
+            "tabular_pipeline expects a scikit-learn compatible estimator instance as"
+            " its first argument, but you have passed a type. Pass an instance of the"
+            " estimator rather than the class itself."
         )
-    if not isinstance(estimator, BaseEstimator):
+
+    is_scikit_learn_compatible, incompatable_reason = (
+        is_scikit_learn_compatible_estimator(estimator)
+    )
+    if not is_scikit_learn_compatible:
         raise TypeError(
-            "tabular_pipeline expects a scikit-learn estimator, 'regressor',"
-            " or 'classifier' as its first argument."
+            "tabular_pipeline expects a scikit-learn compatible estimator as its first"
+            " argument. " + incompatable_reason
         )
 
     is_estimator_from_tabicl = estimator.__class__.__name__ in (
         "TabICLClassifier",
         "TabICLRegressor",
     )
+    is_hgbt_estimator = any(
+        x.lower() in estimator.__class__.__name__.lower()
+        for x in _HGBT_CLASS_NAME_SUBSTRINGS
+    )
+    is_tree_ensemble_estimator = any(
+        x.lower() in estimator.__class__.__name__.lower()
+        for x in _TREE_ENSEMBLE_CLASS_NAME_SUBSTRINGS
+    )
+
     if (
-        isinstance(estimator, _HGBT_CLASSES)
+        is_hgbt_estimator
         and getattr(estimator, "categorical_features", None) == "from_dtype"
     ):
         vectorizer.set_params(
             low_cardinality=ToCategorical(),
             high_cardinality=StringEncoder(),
         )
-    elif isinstance(estimator, _TREE_ENSEMBLE_CLASSES):
+    elif is_tree_ensemble_estimator:
         vectorizer.set_params(
             low_cardinality=OrdinalEncoder(
                 handle_unknown="use_encoded_value",
@@ -283,9 +311,15 @@ def tabular_pipeline(estimator, *, n_jobs=None):
         vectorizer.set_params(datetime=DatetimeEncoder(periodic_encoding="spline"))
     steps = [vectorizer]
     if not is_estimator_from_tabicl:
-        if not get_tags(estimator).input_tags.allow_nan:
+        # Check whether we need imputation
+        try:
+            allow_nan = get_tags(estimator).input_tags.allow_nan
+        except AttributeError:
+            allow_nan = False
+        if not allow_nan:
             steps.append(SimpleImputer(add_indicator=True))
-        if not isinstance(estimator, _TREE_ENSEMBLE_CLASSES):
+        # Check whether we need squashing scalar
+        if not is_tree_ensemble_estimator:
             steps.append(SquashingScaler(max_absolute_value=5))
 
     steps.append(estimator)

skrub/tests/test_tabular_pipeline.py

@@ -1,6 +1,7 @@
+import numpy as np
+import pandas as pd
 import pytest
 from sklearn import ensemble
-from sklearn.base import BaseEstimator
 from sklearn.impute import SimpleImputer
 from sklearn.linear_model import Ridge
 from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder
@@ -36,14 +37,80 @@ def test_bad_learner():
         match=".*should be 'regressor', 'regression', 'classifier' or 'classification'",
     ):
         tabular_pipeline("bad")
+    with pytest.raises(TypeError, match=".*Pass an instance"):
+        tabular_pipeline(ensemble.HistGradientBoostingRegressor)
     with pytest.raises(
-        TypeError, match=".*Pass an instance of HistGradientBoostingRegressor"
+        TypeError, match=".*expects a scikit-learn compatible estimator"
     ):
-        tabular_pipeline(ensemble.HistGradientBoostingRegressor)
-    with pytest.raises(TypeError, match=".*expects a scikit-learn estimator"):
         tabular_pipeline(object())
 
 
+def test_missing_required_attribute():
+    """Test that a TypeError is raised when the estimator does not have one of the
+    attributes required of a scikit learn-compatible estimator"""
+
+    class MissingSetParams:
+        def fit(self, X, y=None):
+            return self
+
+        def predict(self, X):
+            return np.zeros(X.shape[0])
+
+        def get_params(self):
+            return {}
+
+    with pytest.raises(
+        TypeError, match=".*expects a scikit-learn compatible estimator.*set_params"
+    ):
+        tabular_pipeline(MissingSetParams())
+
+
+def test_required_attribute_is_not_callable():
+    """Test that a TypeError is raised when the estimator has all of the required
+    attributes, but one of them is not callable"""
+
+    class PredictNotCallable:
+        def fit(self, X, y=None):
+            return self
+
+        predict = 1
+
+        def get_params(self):
+            return {}
+
+        def set_params(self, **params):
+            return self
+
+    with pytest.raises(
+        TypeError, match=".*expects a scikit-learn compatible estimator.*predict"
+    ):
+        tabular_pipeline(PredictNotCallable())
+
+
+class Regressor:
+    """Dummy regressor used for tests"""
+
+    def fit(self, X, y=None):
+        return self
+
+    def predict(self, X):
+        return np.zeros(X.shape[0])
+
+    def get_params(self):
+        return {}
+
+    def set_params(self, **params):
+        return self
+
+
+def test_sklearn_compatible_learner_returns_correct_pipeline():
+    """Test that no error is raised when the estimate have both `get_params`
+    and `set_params` attributes"""
+    pipeline = tabular_pipeline(Regressor())
+    X = pd.DataFrame({"feature": [1, 2, 3]})
+    pipeline.fit(X)
+
+
 def test_linear_learner():
     original_learner = Ridge()
     p = tabular_pipeline(original_learner)
@@ -66,6 +133,40 @@ def test_tree_learner():
     assert tv.datetime.periodic_encoding is None
 
 
+def test_tree_ensemble_treatment_for_any_random_forest():
+    """Test that special treatment for tree ensemble models is applied when
+    substring 'RandomForest' appears in estimator class name"""
+
+    class IAmARandomForestEstimator(Regressor):
+        pass
+
+    original_learner = IAmARandomForestEstimator()
+    p = tabular_pipeline(original_learner)
+    _, tv = p.steps[0]
+    _, learner = p.steps[-1]
+    assert learner is original_learner
+    assert isinstance(tv.high_cardinality, StringEncoder)
+    assert isinstance(tv.low_cardinality, OrdinalEncoder)
+    assert tv.datetime.periodic_encoding is None
+
+
+def test_tree_ensemble_treatment_for_xgboost():
+    """Test that special treatment for tree ensemble models is applied when
+    substring 'XGB' appears in estimator class name"""
+
+    class IAmXGB(Regressor):
+        pass
+
+    original_learner = IAmXGB()
+    p = tabular_pipeline(original_learner)
+    _, tv = p.steps[0]
+    _, learner = p.steps[-1]
+    assert learner is original_learner
+    assert isinstance(tv.high_cardinality, StringEncoder)
+    assert isinstance(tv.low_cardinality, OrdinalEncoder)
+    assert tv.datetime.periodic_encoding is None
+
+
 def test_from_dtype():
     p = tabular_pipeline(
         ensemble.HistGradientBoostingRegressor(categorical_features=())
@@ -77,13 +178,13 @@ def test_from_dtype():
     assert isinstance(p.named_steps["tablevectorizer"].low_cardinality, ToCategorical)
 
 
-class TabICLClassifier(BaseEstimator):
+class TabICLClassifier(Regressor):
     """Dummy class which pretends to be `tabicl.TabICLClassifier`"""
 
     pass
 
 
-class TabICLRegressor(BaseEstimator):
+class TabICLRegressor(Regressor):
     """Dummy class which pretends to be `tabicl.TabICLRegressor`"""
 
     pass