MODELCATALOG

Write your own estimator that behaves like the built-ins.

sklearn.base

base

Write your own estimator that behaves like the built-ins.

1 Get data Something to practise on
2 Split & tune An honest test set
3 Prepare Into numbers a model accepts
4 Compose One object, no leakage
5 Model The algorithm itself
6 Unsupervised No labels available
7 Measure Is it any good, and why
8 Utilities Glue and extension

What it is for

Inherit from these and your custom class works inside Pipeline, GridSearchCV and cross_validate with no special handling.

Contents

BaseEstimator Gives you get_params and set_params, which the search tools need
TransformerMixin Gives you fit_transform for free
ClassifierMixin, RegressorMixin Gives you a default score method
clone Copy an unfitted estimator with the same settings

Worked examples 2

Runnable, with the arguments that change the outcome called out. Output shapes are shown as comments where they help.

sklearn.base

from sklearn.base import BaseEstimator, TransformerMixin

class TextLengthFeatures(BaseEstimator, TransformerMixin):
    """Adds length signals - often useful alongside TF-IDF, since long
    complaints and short ones behave differently."""

    def fit(self, X, y=None):
        return self                     # nothing to learn

    def transform(self, X):
        import numpy as np
        return np.c_[
            [len(t) for t in X],
            [len(t.split()) for t in X],
            [t.count("!") for t in X],
        ]

# Drops straight into a Pipeline or FeatureUnion.

A custom transformer that behaves like a built-in

Inherit from these two and your class works inside Pipeline, ColumnTransformer and GridSearchCV with no special handling.

from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np

class TextStats(BaseEstimator, TransformerMixin):
    """Length signals to sit alongside TF-IDF. Long complaints and short
    ones behave differently, and a bag of words cannot see that."""

    def __init__(self, include_punct=True):
        # Store constructor args UNCHANGED on self. get_params relies on
        # this; validation belongs in fit, not here.
        self.include_punct = include_punct

    def fit(self, X, y=None):
        return self                       # nothing to learn

    def transform(self, X):
        feats = [[len(t), len(t.split()), np.mean([len(w) for w in t.split()])]
                 for t in X]
        if self.include_punct:
            feats = [f + [t.count("!"), t.count("?")]
                     for f, t in zip(feats, X)]
        return np.array(feats)

    def get_feature_names_out(self, input_features=None):
        names = ["n_chars", "n_words", "avg_word_len"]
        return np.array(names + (["n_excl", "n_quest"]
                                 if self.include_punct else []))

# Drops straight into the stack, and GridSearchCV can tune it:
#   {"stats__include_punct": [True, False]}