Clip output from sklearn pipeline predict

Question

Let's see the following pipeline:

scaler = ScalerFactory.get_scaler(scaler_type)
model = MultiOutputRegressor(lgb.LGBMRegressor(metric='tweedie', **hyperparameters))
steps = [('scaler', scaler), ('model', model)]
pipeline = Pipeline(steps)
pipeline.fit(X, y, model__feature_name=list(X.columns))

I am trying to add another step to the pipeline, so when it predicts it rounds all the values that are between -1 and 1 to 0.

I am trying to create a new class:

from numpy.random import randint
from sklearn.base import BaseEstimator, TransformerMixin


class OutputClipper(BaseEstimator, TransformerMixin):
    def __init__(self) -> None:
        super().__init__()
        self.clipping = False

    def fit(self, X, y=None):
        return self
    
    def transform(self, X, y):
        y[(y>-1) & (y<1)] = 0
        return y

and the new pipeline becomes:

steps = [('scaler', scaler), ('model', model), ('clipping',OutputClipper) ]
pipeline = Pipeline(steps)

However, I feel like this doesn't quite work. I guess the transform happens when the pipeline is called with the .predict() method. I am not sure how to test it too.

Clip output from sklearn pipeline predict

0 Answers0