Code / scripts/ml_training/essay_scoring/training/xgboost/trainer.py
scripts/ml_training/essay_scoring/training/xgboost/trainer.py
1
"""Training utilities for the essay scoring research pipeline.
2
3
Purpose:
4
Train XGBoost models with a consistent early-stopping signal (QWK) and
5
return structured artifacts. The default mode is regression. CV experiments
6
can request an ordinal-as-multiclass mode to study grade-band compression
7
and tail behavior.
8
9
Relationships:
10
- Used by OOF fold execution and the retained final-evaluation experiment.
11
- Used by auxiliary-target and CJ/AES fusion training.
12
- Delegates training-mode-specific encoding and metrics to ``training_modes``.
13
"""
14
15
from __future__ import annotations
16
17
from dataclasses import dataclass
18
19
import numpy as np
20
import xgboost as xgb
21
22
from ...config import TrainingConfig
23
from ...features.combiner import FeatureMatrix
24
from ...training.xgboost.inner_prompt_holdout import (
25
InnerPromptHoldout,
26
)
27
from ...training.xgboost.training_modes import (
28
TrainingMode,
29
encode_ordinal_labels,
30
is_ordinal_multiclass,
31
qwk_custom_metric,
32
)
33
from ...training.xgboost_training_profile import (
34
effective_xgboost_params,
35
)
36
37
38
@dataclass(frozen=True)
39
class TrainingArtifacts:
40
"""Artifacts returned from model training."""
41
42
model: xgb.Booster
43
evals_result: dict[str, dict[str, list[float]]]
44
best_iteration: int
45
46
47
def train_model(
48
train_features: FeatureMatrix,
49
val_features: FeatureMatrix,
50
y_train: np.ndarray,
51
y_val: np.ndarray,
52
config: TrainingConfig,
53
*,
54
min_band: float,
55
max_band: float,
56
training_mode: TrainingMode = TrainingMode.REGRESSION,
57
sample_weight_train: np.ndarray | None = None,
58
) -> TrainingArtifacts:
59
"""Train an XGBoost model with QWK early stopping."""
60
61
if train_features.feature_names != val_features.feature_names:
62
raise ValueError("Feature name mismatch between train and validation sets.")
63
64
if is_ordinal_multiclass(training_mode):
65
y_train = encode_ordinal_labels(y_train, min_band=min_band, max_band=max_band)
66
y_val = encode_ordinal_labels(y_val, min_band=min_band, max_band=max_band)
67
68
dtrain = xgb.DMatrix(
69
train_features.matrix,
70
label=y_train,
71
feature_names=train_features.feature_names,
72
weight=sample_weight_train,
73
)
74
dval = xgb.DMatrix(val_features.matrix, label=y_val, feature_names=val_features.feature_names)
75
76
evals_result: dict[str, dict[str, list[float]]] = {}
77
model = xgb.train(
78
params=effective_xgboost_params(
79
config=config,
80
training_mode=training_mode,
81
min_band=min_band,
82
max_band=max_band,
83
seed=config.random_seed,
84
),
85
dtrain=dtrain,
86
num_boost_round=config.num_boost_round,
87
evals=[(dtrain, "train"), (dval, "val")],
88
early_stopping_rounds=config.early_stopping_rounds,
89
custom_metric=qwk_custom_metric(min_band=min_band, max_band=max_band, mode=training_mode),
90
maximize=True,
91
evals_result=evals_result,
92
)
93
94
best_iteration = int(model.best_iteration or 0)
95
return TrainingArtifacts(
96
model=model,
97
evals_result=evals_result,
98
best_iteration=best_iteration,
99
)
100
101
102
def train_model_with_inner_holdout(
103
train_features: FeatureMatrix,
104
y_train: np.ndarray,
105
holdout: InnerPromptHoldout,
106
config: TrainingConfig,
107
*,
108
min_band: float,
109
max_band: float,
110
training_mode: TrainingMode = TrainingMode.REGRESSION,
111
sample_weight_train: np.ndarray | None = None,
112
) -> TrainingArtifacts:
113
"""Train with nested early stopping and refit on all fold-train rows.
114
115
The fold model that produces out-of-fold predictions must see every fold-train row, while
116
the round count that governs its complexity comes from an inner holdout the outer
117
validation rows never touched. This function therefore fits twice: once on
118
`holdout.train_positions` with early stopping against `holdout.holdout_positions` to
119
discover the inner best iteration, and once more on the complete `train_features`/`y_train`
120
rows for exactly `inner_best_iteration + 1` rounds, with no early stopping, no held-out
121
evaluation set, and no custom metric. `best_iteration` in the returned artifacts is the
122
INNER best iteration, not a property of the returned booster; the returned booster carries
123
`inner_best_iteration + 1` trees, so a plain `predict` call on it uses exactly those trees
124
without needing an explicit `iteration_range`.
125
126
Args:
127
train_features: All of one outer fold's training rows.
128
y_train: Labels for every row of `train_features`, in band units.
129
holdout: The inner prompt-grouped partition of the fold-train rows. Its
130
`train_positions`/`holdout_positions` index into `train_features`.
131
config: Training configuration (round budget, early-stopping patience, booster params).
132
min_band: Minimum score band value (inclusive).
133
max_band: Maximum score band value (inclusive).
134
training_mode: Regression or ordinal-multiclass training mode.
135
sample_weight_train: Optional per-row weights for every row of `train_features`.
136
137
Returns:
138
Training artifacts whose `model` is the full-fold-train refit and whose
139
`best_iteration` is the inner holdout's best iteration.
140
"""
141
142
_validate_inner_holdout_partition(
143
train_features=train_features,
144
y_train=y_train,
145
holdout=holdout,
146
sample_weight_train=sample_weight_train,
147
)
148
149
if is_ordinal_multiclass(training_mode):
150
y_train = encode_ordinal_labels(y_train, min_band=min_band, max_band=max_band)
151
152
params = effective_xgboost_params(
153
config=config,
154
training_mode=training_mode,
155
min_band=min_band,
156
max_band=max_band,
157
seed=config.random_seed,
158
)
159
160
inner_train_positions = holdout.train_positions
161
inner_holdout_positions = holdout.holdout_positions
162
163
inner_train_weight = (
164
None if sample_weight_train is None else sample_weight_train[inner_train_positions]
165
)
166
d_inner_train = xgb.DMatrix(
167
train_features.matrix[inner_train_positions],
168
label=y_train[inner_train_positions],
169
feature_names=train_features.feature_names,
170
weight=inner_train_weight,
171
)
172
d_inner_holdout = xgb.DMatrix(
173
train_features.matrix[inner_holdout_positions],
174
label=y_train[inner_holdout_positions],
175
feature_names=train_features.feature_names,
176
)
177
178
evals_result: dict[str, dict[str, list[float]]] = {}
179
inner_model = xgb.train(
180
params=params,
181
dtrain=d_inner_train,
182
num_boost_round=config.num_boost_round,
183
evals=[(d_inner_train, "inner_train"), (d_inner_holdout, "inner_holdout")],
184
early_stopping_rounds=config.early_stopping_rounds,
185
custom_metric=qwk_custom_metric(min_band=min_band, max_band=max_band, mode=training_mode),
186
maximize=True,
187
evals_result=evals_result,
188
)
189
190
inner_best_iteration = int(inner_model.best_iteration or 0)
191
192
d_full_train = xgb.DMatrix(
193
train_features.matrix,
194
label=y_train,
195
feature_names=train_features.feature_names,
196
weight=sample_weight_train,
197
)
198
refit_model = xgb.train(
199
params=params,
200
dtrain=d_full_train,
201
num_boost_round=inner_best_iteration + 1,
202
verbose_eval=False,
203
)
204
205
return TrainingArtifacts(
206
model=refit_model,
207
evals_result=evals_result,
208
best_iteration=inner_best_iteration,
209
)
210
211
212
def _validate_inner_holdout_partition(
213
*,
214
train_features: FeatureMatrix,
215
y_train: np.ndarray,
216
holdout: InnerPromptHoldout,
217
sample_weight_train: np.ndarray | None,
218
) -> None:
219
"""Fail closed unless the holdout positions exactly partition the fold-train rows."""
220
221
if y_train.ndim != 1:
222
raise ValueError(f"y_train must be 1D, got y_train.ndim={y_train.ndim}.")
223
224
n_rows = train_features.matrix.shape[0]
225
if y_train.shape[0] != n_rows:
226
raise ValueError(
227
f"y_train row count must match train_features, "
228
f"y_train.shape[0]={y_train.shape[0]} train_features_rows={n_rows}."
229
)
230
231
if sample_weight_train is not None and sample_weight_train.shape[0] != n_rows:
232
raise ValueError(
233
f"sample_weight_train row count must match train_features, "
234
f"sample_weight_train.shape[0]={sample_weight_train.shape[0]} "
235
f"train_features_rows={n_rows}."
236
)
237
238
train_positions = holdout.train_positions
239
holdout_positions = holdout.holdout_positions
240
241
out_of_range_positions = sorted(
242
{
243
int(position)
244
for position in (*train_positions.tolist(), *holdout_positions.tolist())
245
if position < 0 or position >= n_rows
246
}
247
)
248
if out_of_range_positions:
249
raise ValueError(
250
f"inner holdout positions out of range for train_features_rows={n_rows}, "
251
f"out_of_range_positions={out_of_range_positions}."
252
)
253
254
train_position_set = set(train_positions.tolist())
255
holdout_position_set = set(holdout_positions.tolist())
256
257
intersecting_positions = sorted(train_position_set & holdout_position_set)
258
if intersecting_positions:
259
raise ValueError(
260
f"inner holdout train/holdout positions intersect, "
261
f"intersecting_positions={intersecting_positions}."
262
)
263
264
union_positions = train_position_set | holdout_position_set
265
expected_positions = set(range(n_rows))
266
if union_positions != expected_positions:
267
missing_positions = sorted(expected_positions - union_positions)
268
extra_positions = sorted(union_positions - expected_positions)
269
raise ValueError(
270
f"inner holdout positions must union to every fold-train row, "
271
f"missing_positions={missing_positions} extra_positions={extra_positions}."
272
)