bikes.core.models

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