Skip to content

TrainerSupervised

The TrainerSupervised class is part of the jarvais.trainer module.

jarvais.trainer.TrainerSupervised

TrainerSupervised is a class for automating the process of feature reduction, model training, and evaluation for various machine learning tasks.

Parameters:

Name Type Description Default
output_dir str | Path

The output directory for saving the trained model and data.

 required
target_variable str | list[str]

The column name of the target variable, or a list of two column names for survival analysis.

 required
task str

The type of task to perform, e.g. 'binary', 'multiclass', 'regression', or 'survival'.

 required
stratify_on str | None

The column name of a variable to stratify the train-test split over. If None, no stratification will be performed.

 None
test_size float

The proportion of data to use for testing. Default is 0.2.

 0.2
k_folds int

The number of folds to use for cross-validation. Default is 5.

 5
reduction_method str | None

The method to use for feature reduction. If None, no feature reduction will be performed.

 None
keep_k int

The number of features to keep after reduction. Default is 2.

 2
random_state int

The random state for reproducibility. Default is 42.

 42
explain bool

Whether to generate explanations for the model. Default is False.

 False
Source code in src/jarvais/trainer/trainer.py 
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
class TrainerSupervised:
    """
    TrainerSupervised is a class for automating the process of feature reduction, 
    model training, and evaluation for various machine learning tasks.

    Parameters:
        output_dir (str | Path): The output directory for saving the trained model and data.
        target_variable (str | list[str]): The column name of the target variable, or a list of two column names for survival analysis.
        task (str): The type of task to perform, e.g. 'binary', 'multiclass', 'regression', or 'survival'.
        stratify_on (str | None): The column name of a variable to stratify the train-test split over. If None, no stratification will be performed.
        test_size (float): The proportion of data to use for testing. Default is 0.2.
        k_folds (int): The number of folds to use for cross-validation. Default is 5.
        reduction_method (str | None): The method to use for feature reduction. If None, no feature reduction will be performed.
        keep_k (int): The number of features to keep after reduction. Default is 2.
        random_state (int): The random state for reproducibility. Default is 42.
        explain (bool): Whether to generate explanations for the model. Default is False.
    """
    def __init__(
        self,
        output_dir: str | Path,
        target_variable: str | list[str],
        task: str,
        stratify_on: str | None = None,
        test_size: float = 0.2,
        k_folds: int = 5,
        reduction_method: str | None = None,
        keep_k: int = 2,
        random_state: int = 42,
        explain: bool = False
    ) -> None:

        self.reduction_module = FeatureReductionModule.build(
            method=reduction_method,
            task=task,
            keep_k=keep_k
        )

        if task == "survival":
            if set(target_variable) != {'time', 'event'}: 
                raise ValueError("Target variable must be a list of ['time', 'event'] for survival analysis.")

            self.trainer_module = SurvivalTrainerModule.build(
                output_dir=output_dir 
            )
        else:
            self.trainer_module = AutogluonTabularWrapper.build(
                output_dir=output_dir,
                target_variable=target_variable,
                task=task,
                k_folds=k_folds
            )

        self.settings = TrainerSettings(
            output_dir=Path(output_dir),
            target_variable=target_variable,
            task=task,
            stratify_on=stratify_on,
            test_size=test_size,
            random_state=random_state,
            explain=explain,
            reduction_module=self.reduction_module,
            trainer_module=self.trainer_module,
        )

    @classmethod
    def from_settings(
            cls, 
            settings_dict: dict,
        ) -> "TrainerSupervised":
        """
        Initialize a TrainerSupervised instance with a given settings dictionary.

        Args:
            dict: A dictionary containing the settings for the TrainerSupervised instance.

        Returns:
            TrainerSupervised: An instance of TrainerSupervised with the given settings.
        """
        settings = TrainerSettings.model_validate(settings_dict)

        trainer = cls(
            output_dir=settings.output_dir,
            target_variable=settings.target_variable,
            task=settings.task,
        )

        trainer.reduction_module = settings.reduction_module
        trainer.trainer_module = settings.trainer_module

        trainer.settings = settings

        return trainer

    def run(
            self,
            data: pd.DataFrame 
        ) -> None:
        self.data = data

        # Preprocess
        X = self.data.drop(self.settings.target_variable, axis=1)
        y = self.data[self.settings.target_variable]

        X, y = self.reduction_module(X, y)     

        if self.settings.task in {'binary', 'multiclass'}:
            stratify_col = (
                y.astype(str) + '_' + self.data[self.settings.stratify_on].astype(str)
                if self.settings.stratify_on is not None
                else y
            )
        else:
            stratify_col = None

        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            X, 
            y, 
            test_size=self.settings.test_size, 
            stratify=stratify_col, 
            random_state=self.settings.random_state
        )

        # Train
        self.predictor, self.X_val, self.y_val = self.trainer_module.fit(
            X_train=self.X_train, 
            y_train=self.y_train, 
            X_test=self.X_test, 
            y_test=self.y_test
        )

        self.X_train = self.X_train.drop(self.X_val.index)
        self.y_train = self.y_train.drop(self.y_val.index)

        data_dir = self.settings.output_dir / 'data'
        data_dir.mkdir(parents=True, exist_ok=True)
        self.X_train.to_csv((data_dir / 'X_train.csv'), index=False)
        self.X_test.to_csv((data_dir / 'X_test.csv'), index=False)
        self.X_val.to_csv((data_dir / 'X_val.csv'), index=False)
        self.y_train.to_csv((data_dir / 'y_train.csv'), index=False)
        self.y_test.to_csv((data_dir / 'y_test.csv'), index=False)
        self.y_val.to_csv((data_dir / 'y_val.csv'), index=False)

        if self.settings.explain:
            explainer = Explainer.from_trainer(self)
            explainer.run()

        # Save Settings
        schema_path = self.settings.output_dir / 'trainer_settings.schema.json'
        with schema_path.open("w") as f:
            json.dump(self.settings.model_json_schema(), f, indent=2)

        settings_path = self.settings.output_dir / 'trainer_settings.json'
        with settings_path.open('w') as f:
            json.dump({
                "$schema": str(schema_path.relative_to(self.settings.output_dir)),
                **self.settings.model_dump(mode="json") 
            }, f, indent=2)

    def model_names(self) -> list[str]:
        """
        Returns all trainer model names.

        This method retrieves the names of all models associated with the 
        current predictor. It requires that the predictor has been trained.

        Returns:
            list: A list of model names available in the predictor.
        """

        return self.predictor.model_names()

    def infer(self, data: pd.DataFrame, model: str | None = None) -> np.ndarray:
        """
        Make predictions on new data using the trained predictor.

        Args:
            data (pd.DataFrame): The new data to make predictions on.
            model (str | None): The model to use for prediction. If None, the best model will be used.

        Returns:
            np.ndarray: The predicted values.
        """

        if self.settings.task == 'survival':
            inference = self.predictor.predict(data, model)
        elif self.predictor.can_predict_proba:
            inference = self.predictor.predict_proba(data, model, as_pandas=False)[:, 1]
        else:
            inference = self.predictor.predict(data, model, as_pandas=False)

        return inference

    @classmethod
    def load_trainer(
            cls, 
            project_dir: str | Path
        ) -> "TrainerSupervised":

        from autogluon.tabular import TabularPredictor

        from jarvais.trainer.modules.survival_trainer import SurvivalPredictor

        project_dir = Path(project_dir)
        with (project_dir / 'trainer_settings.json').open('r') as f:
            trainer_config = json.load(f)

        trainer = cls.from_settings(trainer_config)

        if trainer.settings.task == 'survival':
            model_dir = (project_dir / 'survival_models')
            trainer.predictor = SurvivalPredictor.load(model_dir)
        else:
            model_dir = (project_dir / 'autogluon_models' / 'autogluon_models_best_fold')
            trainer.predictor = TabularPredictor.load(model_dir, verbosity=1)

        trainer.X_test = pd.read_csv(project_dir / 'data' / 'X_test.csv')
        trainer.X_val = pd.read_csv(project_dir / 'data' / 'X_val.csv')
        trainer.X_train = pd.read_csv(project_dir / 'data' / 'X_train.csv')
        trainer.y_test = pd.read_csv(project_dir / 'data' / 'y_test.csv').squeeze()
        trainer.y_val = pd.read_csv(project_dir / 'data' / 'y_val.csv').squeeze()
        trainer.y_train = pd.read_csv(project_dir / 'data' / 'y_train.csv').squeeze()

        return trainer

    def __rich_repr__(self) -> rich.repr.Result:
        yield self.settings

    def __repr__(self) -> str:
        return f"TrainerSupervised(settings={self.settings.model_dump_json(indent=2)})"

from_settings(settings_dict) classmethod

Initialize a TrainerSupervised instance with a given settings dictionary.

Parameters:

Name Type Description Default
dict

A dictionary containing the settings for the TrainerSupervised instance.

 required

Returns:

Name Type Description
TrainerSupervised TrainerSupervised

An instance of TrainerSupervised with the given settings.
Source code in src/jarvais/trainer/trainer.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
@classmethod
def from_settings(
        cls, 
        settings_dict: dict,
    ) -> "TrainerSupervised":
    """
    Initialize a TrainerSupervised instance with a given settings dictionary.

    Args:
        dict: A dictionary containing the settings for the TrainerSupervised instance.

    Returns:
        TrainerSupervised: An instance of TrainerSupervised with the given settings.
    """
    settings = TrainerSettings.model_validate(settings_dict)

    trainer = cls(
        output_dir=settings.output_dir,
        target_variable=settings.target_variable,
        task=settings.task,
    )

    trainer.reduction_module = settings.reduction_module
    trainer.trainer_module = settings.trainer_module

    trainer.settings = settings

    return trainer

model_names()

Returns all trainer model names.

This method retrieves the names of all models associated with the current predictor. It requires that the predictor has been trained.

Returns:

Name Type Description
list list[str]

A list of model names available in the predictor.
Source code in src/jarvais/trainer/trainer.py 
176
177
178
179
180
181
182
183
184
185
186
187
def model_names(self) -> list[str]:
    """
    Returns all trainer model names.

    This method retrieves the names of all models associated with the 
    current predictor. It requires that the predictor has been trained.

    Returns:
        list: A list of model names available in the predictor.
    """

    return self.predictor.model_names()

infer(data, model=None)

Make predictions on new data using the trained predictor.

Parameters:

Name Type Description Default
data DataFrame

The new data to make predictions on.

 required
model str | None

The model to use for prediction. If None, the best model will be used.

 None

Returns:

Type Description
ndarray

np.ndarray: The predicted values.
Source code in src/jarvais/trainer/trainer.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
def infer(self, data: pd.DataFrame, model: str | None = None) -> np.ndarray:
    """
    Make predictions on new data using the trained predictor.

    Args:
        data (pd.DataFrame): The new data to make predictions on.
        model (str | None): The model to use for prediction. If None, the best model will be used.

    Returns:
        np.ndarray: The predicted values.
    """

    if self.settings.task == 'survival':
        inference = self.predictor.predict(data, model)
    elif self.predictor.can_predict_proba:
        inference = self.predictor.predict_proba(data, model, as_pandas=False)[:, 1]
    else:
        inference = self.predictor.predict(data, model, as_pandas=False)

    return inference