bikes.utils.searchers

  1"""Find the best hyperparameters for a model."""
  2
  3# %% IMPORTS
  4
  5import abc
  6import typing as T
  7
  8import pandas as pd
  9import pydantic as pdt
 10from sklearn import model_selection
 11
 12from bikes.core import metrics, models, schemas
 13from bikes.utils import splitters
 14
 15# %% TYPES
 16
 17# Grid of model params
 18Grid = dict[models.ParamKey, list[models.ParamValue]]
 19
 20# Results of a model search
 21Results = tuple[
 22    T.Annotated[pd.DataFrame, "details"],
 23    T.Annotated[float, "best score"],
 24    T.Annotated[models.Params, "best params"],
 25]
 26
 27# Cross-validation options for searchers
 28CrossValidation = int | splitters.TrainTestSplits | splitters.Splitter
 29
 30# %% SEARCHERS
 31
 32
 33class Searcher(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 34    """Base class for a searcher.
 35
 36    Use searcher to fine-tune models.
 37    i.e., to find the best model params.
 38
 39    Parameters:
 40        param_grid (Grid): mapping of param key -> values.
 41    """
 42
 43    KIND: str
 44
 45    param_grid: Grid
 46
 47    @abc.abstractmethod
 48    def search(
 49        self,
 50        model: models.Model,
 51        metric: metrics.Metric,
 52        inputs: schemas.Inputs,
 53        targets: schemas.Targets,
 54        cv: CrossValidation,
 55    ) -> Results:
 56        """Search the best model for the given inputs and targets.
 57
 58        Args:
 59            model (models.Model): AI/ML model to fine-tune.
 60            metric (metrics.Metric): main metric to optimize.
 61            inputs (schemas.Inputs): model inputs for tuning.
 62            targets (schemas.Targets): model targets for tuning.
 63            cv (CrossValidation): choice for cross-fold validation.
 64
 65        Returns:
 66            Results: all the results of the searcher execution process.
 67        """
 68
 69
 70class GridCVSearcher(Searcher):
 71    """Grid searcher with cross-fold validation.
 72
 73    Convention: metric returns higher values for better models.
 74
 75    Parameters:
 76        n_jobs (int, optional): number of jobs to run in parallel.
 77        refit (bool): refit the model after the tuning.
 78        verbose (int): set the searcher verbosity level.
 79        error_score (str | float): strategy or value on error.
 80        return_train_score (bool): include train scores if True.
 81    """
 82
 83    KIND: T.Literal["GridCVSearcher"] = "GridCVSearcher"
 84
 85    n_jobs: int | None = None
 86    refit: bool = True
 87    verbose: int = 3
 88    error_score: str | float = "raise"
 89    return_train_score: bool = False
 90
 91    @T.override
 92    def search(
 93        self,
 94        model: models.Model,
 95        metric: metrics.Metric,
 96        inputs: schemas.Inputs,
 97        targets: schemas.Targets,
 98        cv: CrossValidation,
 99    ) -> Results:
100        searcher = model_selection.GridSearchCV(
101            estimator=model,
102            scoring=metric.scorer,
103            cv=cv,
104            param_grid=self.param_grid,
105            n_jobs=self.n_jobs,
106            refit=self.refit,
107            verbose=self.verbose,
108            error_score=self.error_score,
109            return_train_score=self.return_train_score,
110        )
111        searcher.fit(inputs, targets)
112        results = pd.DataFrame(searcher.cv_results_)
113        return results, searcher.best_score_, searcher.best_params_
114
115
116SearcherKind = GridCVSearcher