bikes.core.models

  1"""Define trainable machine learning models."""
  2
  3# %% IMPORTS
  4
  5import abc
  6import typing as T
  7
  8import pandas as pd
  9import pydantic as pdt
 10import shap
 11from sklearn import compose, ensemble, pipeline, preprocessing
 12from sklearn.base import BaseEstimator, RegressorMixin
 13
 14from bikes.core import schemas
 15
 16# %% TYPES
 17
 18# Model params
 19ParamKey = str
 20ParamValue = T.Any
 21Params = dict[ParamKey, ParamValue]
 22
 23# %% MODELS
 24
 25
 26class _RegressorTagger(RegressorMixin, BaseEstimator):
 27    """Source of default scikit-learn regressor tags (sklearn >= 1.6 tags API)."""
 28
 29
 30class Model(abc.ABC, pdt.BaseModel, strict=True, frozen=False, extra="forbid"):
 31    """Base class for a project model.
 32
 33    Use a model to adapt AI/ML frameworks.
 34    e.g., to swap easily one model with another.
 35    """
 36
 37    KIND: str
 38
 39    def __sklearn_tags__(self) -> T.Any:
 40        """Return scikit-learn estimator tags.
 41
 42        Meta-estimators (e.g. GridSearchCV) query this since scikit-learn 1.6. The
 43        project Model is a Pydantic adapter rather than a scikit-learn BaseEstimator,
 44        so it borrows the default regressor tags.
 45
 46        Returns:
 47            T.Any: scikit-learn estimator tags for a regressor.
 48        """
 49        return _RegressorTagger().__sklearn_tags__()
 50
 51    def get_params(self, deep: bool = True) -> Params:  # noqa: ARG002  # sklearn get_params interface
 52        """Get the model params.
 53
 54        Args:
 55            deep (bool, optional): ignored.
 56
 57        Returns:
 58            Params: internal model parameters.
 59        """
 60        params: Params = {}
 61        for key, value in self.model_dump().items():
 62            if not key.startswith("_") and not key.isupper():
 63                params[key] = value
 64        return params
 65
 66    def set_params(self, **params: ParamValue) -> T.Self:
 67        """Set the model params in place.
 68
 69        Returns:
 70            T.Self: instance of the model.
 71        """
 72        for key, value in params.items():
 73            setattr(self, key, value)
 74        return self
 75
 76    @abc.abstractmethod
 77    def fit(self, inputs: schemas.Inputs, targets: schemas.Targets) -> T.Self:
 78        """Fit the model on the given inputs and targets.
 79
 80        Args:
 81            inputs (schemas.Inputs): model training inputs.
 82            targets (schemas.Targets): model training targets.
 83
 84        Returns:
 85            T.Self: instance of the model.
 86        """
 87
 88    @abc.abstractmethod
 89    def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
 90        """Generate outputs with the model for the given inputs.
 91
 92        Args:
 93            inputs (schemas.Inputs): model prediction inputs.
 94
 95        Returns:
 96            schemas.Outputs: model prediction outputs.
 97        """
 98
 99    def explain_model(self) -> schemas.FeatureImportances:
100        """Explain the internal model structure.
101
102        Returns:
103            schemas.FeatureImportances: feature importances.
104        """
105        raise NotImplementedError
106
107    def explain_samples(self, inputs: schemas.Inputs) -> schemas.SHAPValues:
108        """Explain model outputs on input samples.
109
110        Returns:
111            schemas.SHAPValues: SHAP values.
112        """
113        raise NotImplementedError
114
115    def get_internal_model(self) -> T.Any:
116        """Return the internal model in the object.
117
118        Raises:
119            NotImplementedError: method not implemented.
120
121        Returns:
122            T.Any: any internal model (either empty or fitted).
123        """
124        raise NotImplementedError
125
126
127class BaselineSklearnModel(Model):
128    """Simple baseline model based on scikit-learn.
129
130    Parameters:
131        max_depth (int): maximum depth of the random forest.
132        n_estimators (int): number of estimators in the random forest.
133        random_state (int, optional): random state of the machine learning pipeline.
134    """
135
136    KIND: T.Literal["BaselineSklearnModel"] = "BaselineSklearnModel"
137
138    # params
139    max_depth: int = 20
140    n_estimators: int = 200
141    random_state: int | None = 42
142    # private
143    _pipeline: pipeline.Pipeline | None = None
144    _numericals: list[str] = [
145        "yr",
146        "mnth",
147        "hr",
148        "holiday",
149        "weekday",
150        "workingday",
151        "temp",
152        "atemp",
153        "hum",
154        "windspeed",
155        "casual",
156        "registered",  # too correlated with target
157    ]
158    _categoricals: list[str] = [
159        "season",
160        "weathersit",
161    ]
162
163    @T.override
164    def fit(self, inputs: schemas.Inputs, targets: schemas.Targets) -> BaselineSklearnModel:
165        # subcomponents
166        categoricals_transformer = preprocessing.OneHotEncoder(sparse_output=False, handle_unknown="ignore")
167        # components
168        transformer = compose.ColumnTransformer(
169            [
170                ("categoricals", categoricals_transformer, self._categoricals),
171                ("numericals", "passthrough", self._numericals),
172            ],
173            remainder="drop",
174        )
175        regressor = ensemble.RandomForestRegressor(
176            max_depth=self.max_depth,
177            n_estimators=self.n_estimators,
178            random_state=self.random_state,
179        )
180        # pipeline
181        self._pipeline = pipeline.Pipeline(
182            steps=[
183                ("transformer", transformer),
184                ("regressor", regressor),
185            ]
186        )
187        self._pipeline.fit(X=inputs, y=targets[schemas.TargetsSchema.cnt])
188        return self
189
190    @T.override
191    def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
192        model = self.get_internal_model()
193        prediction = model.predict(inputs)
194        outputs_ = pd.DataFrame(data={schemas.OutputsSchema.prediction: prediction}, index=inputs.index)
195        return schemas.OutputsSchema.check(data=outputs_)
196
197    @T.override
198    def explain_model(self) -> schemas.FeatureImportances:
199        model = self.get_internal_model()
200        regressor = model.named_steps["regressor"]
201        transformer = model.named_steps["transformer"]
202        feature = transformer.get_feature_names_out()
203        feature_importances_ = pd.DataFrame(
204            data={
205                "feature": feature,
206                "importance": regressor.feature_importances_,
207            }
208        )
209        return schemas.FeatureImportancesSchema.check(data=feature_importances_)
210
211    @T.override
212    def explain_samples(self, inputs: schemas.Inputs) -> schemas.SHAPValues:
213        model = self.get_internal_model()
214        regressor = model.named_steps["regressor"]
215        transformer = model.named_steps["transformer"]
216        transformed = transformer.transform(X=inputs)
217        explainer = shap.TreeExplainer(model=regressor)
218        shap_values_ = pd.DataFrame(
219            data=explainer.shap_values(X=transformed),
220            columns=transformer.get_feature_names_out(),
221        )
222        return schemas.SHAPValuesSchema.check(data=shap_values_)
223
224    @T.override
225    def get_internal_model(self) -> pipeline.Pipeline:
226        model = self._pipeline
227        if model is None:
228            raise ValueError("Model is not fitted yet!")
229        return model
230
231
232ModelKind = BaselineSklearnModel