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