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
 11import skops.io as sio
 12from sklearn import compose, ensemble, pipeline, preprocessing
 13from sklearn.base import BaseEstimator, RegressorMixin
 14
 15from bikes.core import schemas
 16
 17# %% TYPES
 18
 19# Model params
 20ParamKey = str
 21ParamValue = T.Any
 22Params = dict[ParamKey, ParamValue]
 23
 24# %% MODELS
 25
 26
 27class _RegressorTagger(RegressorMixin, BaseEstimator):
 28    """Source of default scikit-learn regressor tags (sklearn >= 1.6 tags API)."""
 29
 30
 31class Model(abc.ABC, pdt.BaseModel, strict=True, frozen=False, extra="forbid"):
 32    """Base class for a project model.
 33
 34    Use a model to adapt AI/ML frameworks.
 35    e.g., to swap easily one model with another.
 36    """
 37
 38    KIND: str
 39
 40    def __sklearn_tags__(self) -> T.Any:
 41        """Return scikit-learn estimator tags.
 42
 43        Meta-estimators (e.g. GridSearchCV) query this since scikit-learn 1.6. The
 44        project Model is a Pydantic adapter rather than a scikit-learn BaseEstimator,
 45        so it borrows the default regressor tags.
 46
 47        Returns:
 48            T.Any: scikit-learn estimator tags for a regressor.
 49        """
 50        return _RegressorTagger().__sklearn_tags__()
 51
 52    def get_params(self, deep: bool = True) -> Params:  # noqa: ARG002  # sklearn get_params interface
 53        """Get the model params.
 54
 55        Args:
 56            deep (bool, optional): ignored.
 57
 58        Returns:
 59            Params: internal model parameters.
 60        """
 61        params: Params = {}
 62        for key, value in self.model_dump().items():
 63            if not key.startswith("_") and not key.isupper():
 64                params[key] = value
 65        return params
 66
 67    def set_params(self, **params: ParamValue) -> T.Self:
 68        """Set the model params in place.
 69
 70        Returns:
 71            T.Self: instance of the model.
 72        """
 73        for key, value in params.items():
 74            setattr(self, key, value)
 75        return self
 76
 77    @abc.abstractmethod
 78    def fit(self, inputs: schemas.Inputs, targets: schemas.Targets) -> T.Self:
 79        """Fit the model on the given inputs and targets.
 80
 81        Args:
 82            inputs (schemas.Inputs): model training inputs.
 83            targets (schemas.Targets): model training targets.
 84
 85        Returns:
 86            T.Self: instance of the model.
 87        """
 88
 89    @abc.abstractmethod
 90    def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
 91        """Generate outputs with the model for the given inputs.
 92
 93        Args:
 94            inputs (schemas.Inputs): model prediction inputs.
 95
 96        Returns:
 97            schemas.Outputs: model prediction outputs.
 98        """
 99
100    def explain_model(self) -> schemas.FeatureImportances:
101        """Explain the internal model structure.
102
103        Returns:
104            schemas.FeatureImportances: feature importances.
105        """
106        raise NotImplementedError
107
108    def explain_samples(self, inputs: schemas.Inputs) -> schemas.SHAPValues:
109        """Explain model outputs on input samples.
110
111        Returns:
112            schemas.SHAPValues: SHAP values.
113        """
114        raise NotImplementedError
115
116    def get_internal_model(self) -> T.Any:
117        """Return the internal model in the object.
118
119        Raises:
120            NotImplementedError: method not implemented.
121
122        Returns:
123            T.Any: any internal model (either empty or fitted).
124        """
125        raise NotImplementedError
126
127    def save_internal_model(self, path: str) -> None:
128        """Save the fitted internal model to a file in a safe (non-pickle) format.
129
130        Args:
131            path (str): file path to write.
132
133        Raises:
134            NotImplementedError: method not implemented.
135        """
136        raise NotImplementedError
137
138    def load_internal_model(self, path: str) -> T.Self:
139        """Load the internal model from a file written by `save_internal_model`.
140
141        Args:
142            path (str): file path to read.
143
144        Raises:
145            NotImplementedError: method not implemented.
146
147        Returns:
148            T.Self: instance of the model.
149        """
150        raise NotImplementedError
151
152
153class BaselineSklearnModel(Model):
154    """Simple baseline model based on scikit-learn.
155
156    Parameters:
157        max_depth (int): maximum depth of the random forest.
158        n_estimators (int): number of estimators in the random forest.
159        random_state (int, optional): random state of the machine learning pipeline.
160    """
161
162    KIND: T.Literal["BaselineSklearnModel"] = "BaselineSklearnModel"
163
164    # skops only reloads the types it knows are safe. Tree storage holds raw node indices
165    # that scikit-learn reads without bounds checks, so it is trusted explicitly here,
166    # for files this project writes, instead of trusting whatever a file declares.
167    TRUSTED_TYPES: T.ClassVar[list[str]] = ["sklearn.tree._tree.Tree"]
168
169    # params
170    max_depth: int = 20
171    n_estimators: int = 200
172    random_state: int | None = 42
173    # private
174    _pipeline: pipeline.Pipeline | None = None
175    _numericals: list[str] = [
176        "yr",
177        "mnth",
178        "hr",
179        "holiday",
180        "weekday",
181        "workingday",
182        "temp",
183        "atemp",
184        "hum",
185        "windspeed",
186        # casual and registered are left out: they add up to the target (cnt = casual + registered)
187    ]
188    _categoricals: list[str] = [
189        "season",
190        "weathersit",
191    ]
192
193    @T.override
194    def fit(self, inputs: schemas.Inputs, targets: schemas.Targets) -> BaselineSklearnModel:
195        if not inputs.index.equals(targets.index):
196            raise ValueError("Input and target sample IDs must match in the same order.")
197        # subcomponents
198        categoricals_transformer = preprocessing.OneHotEncoder(sparse_output=False, handle_unknown="ignore")
199        # components
200        transformer = compose.ColumnTransformer(
201            [
202                ("categoricals", categoricals_transformer, self._categoricals),
203                ("numericals", "passthrough", self._numericals),
204            ],
205            remainder="drop",
206        )
207        regressor = ensemble.RandomForestRegressor(
208            max_depth=self.max_depth,
209            n_estimators=self.n_estimators,
210            random_state=self.random_state,
211        )
212        # pipeline
213        self._pipeline = pipeline.Pipeline(
214            steps=[
215                ("transformer", transformer),
216                ("regressor", regressor),
217            ]
218        )
219        self._pipeline.fit(X=inputs, y=targets[schemas.TargetsSchema.cnt])
220        return self
221
222    @T.override
223    def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
224        model = self.get_internal_model()
225        # round to the nearest count: the uint32 output schema would otherwise truncate
226        prediction = model.predict(inputs).round()
227        outputs_ = pd.DataFrame(data={schemas.OutputsSchema.prediction: prediction}, index=inputs.index)
228        return schemas.OutputsSchema.check(data=outputs_)
229
230    @T.override
231    def explain_model(self) -> schemas.FeatureImportances:
232        model = self.get_internal_model()
233        regressor = model.named_steps["regressor"]
234        transformer = model.named_steps["transformer"]
235        feature = transformer.get_feature_names_out()
236        feature_importances_ = pd.DataFrame(
237            data={
238                "feature": feature,
239                "importance": regressor.feature_importances_,
240            }
241        )
242        return schemas.FeatureImportancesSchema.check(data=feature_importances_)
243
244    @T.override
245    def explain_samples(self, inputs: schemas.Inputs) -> schemas.SHAPValues:
246        model = self.get_internal_model()
247        regressor = model.named_steps["regressor"]
248        transformer = model.named_steps["transformer"]
249        transformed = transformer.transform(X=inputs)
250        explainer = shap.TreeExplainer(model=regressor)
251        shap_values_ = pd.DataFrame(
252            data=explainer.shap_values(X=transformed),
253            columns=transformer.get_feature_names_out(),
254            index=inputs.index,
255        )
256        return schemas.SHAPValuesSchema.check(data=shap_values_)
257
258    @T.override
259    def get_internal_model(self) -> pipeline.Pipeline:
260        model = self._pipeline
261        if model is None:
262            raise ValueError("Model is not fitted yet!")
263        return model
264
265    @T.override
266    def save_internal_model(self, path: str) -> None:
267        sio.dump(self.get_internal_model(), path)
268
269    @T.override
270    def load_internal_model(self, path: str) -> T.Self:
271        self._pipeline = sio.load(path, trusted=self.TRUSTED_TYPES)
272        return self
273
274
275ModelKind = BaselineSklearnModel