Code / scripts/ml_training/essay_scoring/training/prediction_mapping.py
scripts/ml_training/essay_scoring/training/prediction_mapping.py
1
"""Prediction mapping and post-hoc calibration utilities for essay scoring.
2
3
Purpose:
4
Convert continuous regression outputs (`y_pred_raw`) into discrete half-band predictions
5
(`y_pred`) using either:
6
- the baseline mapping (round-to-nearest-half-band + clip), or
7
- post-hoc calibrated mappings that learn monotone cutpoints and optionally constrain tails.
8
9
These mappings support scorer-calibration experiments that reduce grade-band compression and
10
tail bias without changing the underlying feature extraction pipeline.
11
12
Relationships:
13
- Used by `scripts.ml_training.essay_scoring.cross_validation` to optionally map predictions
14
before computing CV metrics and residual diagnostics artifacts/reports.
15
- Residual diagnostics persist both `y_pred_raw` and mapped `y_pred`.
16
"""
17
18
from __future__ import annotations
19
20
from dataclasses import dataclass
21
from enum import Enum
22
23
import numpy as np
24
25
26
class PredictionMapping(str, Enum):
27
"""Mapping mode from `y_pred_raw` to discrete half-band predictions."""
28
29
ROUND_HALF_BAND = "round_half_band"
30
QWK_CUTPOINTS_LFO = "qwk_cutpoints_lfo"
31
SPLIT_TAIL_LOW_0_50_HIGH_1_00 = "split_tail_low_0_50_high_1_00"
32
33
34
@dataclass(frozen=True)
35
class CutpointCalibration:
36
"""Monotone cutpoints that map continuous predictions to half-band classes."""
37
38
min_band: float
39
max_band: float
40
band_values: np.ndarray
41
cutpoints: np.ndarray
42
43
def apply(self, y_pred_raw: np.ndarray) -> np.ndarray:
44
return apply_cutpoints(
45
y_pred_raw=y_pred_raw,
46
cutpoints=self.cutpoints,
47
band_values=self.band_values,
48
min_band=self.min_band,
49
max_band=self.max_band,
50
)
51
52
53
def half_band_values(*, min_band: float, max_band: float) -> np.ndarray:
54
"""Return the supported half-band values in [min_band, max_band]."""
55
56
step = 0.5
57
n = int(round((max_band - min_band) / step)) + 1
58
values = min_band + step * np.arange(n, dtype=float)
59
# Defensive: avoid tiny float drift (e.g. 3.5000000004)
60
return np.round(values * 2.0) / 2.0
61
62
63
def map_round_half_band(*, y_pred_raw: np.ndarray, min_band: float, max_band: float) -> np.ndarray:
64
"""Baseline mapping: round-to-nearest-half-band then clip."""
65
66
rounded = np.round(y_pred_raw.astype(float) * 2.0) / 2.0
67
return np.clip(rounded, min_band, max_band)
68
69
70
def apply_cutpoints(
71
*,
72
y_pred_raw: np.ndarray,
73
cutpoints: np.ndarray,
74
band_values: np.ndarray,
75
min_band: float,
76
max_band: float,
77
) -> np.ndarray:
78
"""Apply cutpoints to map predictions into discrete half-band values."""
79
80
if int(band_values.size) < 2:
81
raise ValueError("band_values must contain at least 2 classes")
82
if int(cutpoints.size) != int(band_values.size) - 1:
83
raise ValueError("cutpoints must have length len(band_values)-1")
84
85
clipped = np.clip(y_pred_raw.astype(float), min_band, max_band)
86
class_ids = np.digitize(clipped, cutpoints.astype(float), right=False)
87
return band_values[class_ids]
88
89
90
def mapping_requires_qwk_cutpoints(prediction_mapping: PredictionMapping) -> bool:
91
"""Return whether the mapping learns fold-safe QWK cutpoints before application."""
92
93
return prediction_mapping in {
94
PredictionMapping.QWK_CUTPOINTS_LFO,
95
PredictionMapping.SPLIT_TAIL_LOW_0_50_HIGH_1_00,
96
}
97
98
99
def effective_cutpoints_for_mapping(
100
*,
101
prediction_mapping: PredictionMapping,
102
learned_cutpoints: np.ndarray,
103
band_values: np.ndarray,
104
) -> np.ndarray:
105
"""Convert learned LFO cutpoints into the effective cutpoints for a mapping mode."""
106
107
if prediction_mapping == PredictionMapping.QWK_CUTPOINTS_LFO:
108
return learned_cutpoints.astype(float)
109
if prediction_mapping == PredictionMapping.SPLIT_TAIL_LOW_0_50_HIGH_1_00:
110
return split_tail_low_0_50_high_1_00_cutpoints(
111
learned_cutpoints=learned_cutpoints,
112
band_values=band_values,
113
)
114
raise ValueError(f"Mapping does not use learned cutpoints: {prediction_mapping.value}")
115
116
117
def split_tail_low_0_50_high_1_00_cutpoints(
118
*,
119
learned_cutpoints: np.ndarray,
120
band_values: np.ndarray,
121
) -> np.ndarray:
122
"""Blend LFO tail cutpoints while keeping middle cutpoints at half-band rounding.
123
124
The mapping keeps ordinary rounding thresholds for the middle score bands, blends the
125
low-tail thresholds halfway toward the learned LFO thresholds, and applies the learned
126
thresholds fully for the high tail. The constants match TASK-0553 Slice 2's lead diagnostic
127
candidate, `split_tail_low_0_50_high_1_00`.
128
"""
129
130
base_cutpoints = half_band_rounding_cutpoints(band_values=band_values)
131
if int(learned_cutpoints.size) != int(base_cutpoints.size):
132
raise ValueError("learned_cutpoints must have length len(band_values)-1")
133
134
effective = base_cutpoints.astype(float).copy()
135
low_tail_indices = np.array([0, 1, 2], dtype=int)
136
high_tail_indices = np.array(
137
[base_cutpoints.size - 3, base_cutpoints.size - 2, base_cutpoints.size - 1],
138
dtype=int,
139
)
140
effective[low_tail_indices] = base_cutpoints[low_tail_indices] + 0.5 * (
141
learned_cutpoints.astype(float)[low_tail_indices] - base_cutpoints[low_tail_indices]
142
)
143
effective[high_tail_indices] = learned_cutpoints.astype(float)[high_tail_indices]
144
return _enforce_strictly_increasing(effective)
145
146
147
def half_band_rounding_cutpoints(*, band_values: np.ndarray) -> np.ndarray:
148
"""Return ordinary nearest-half-band cutpoints for the provided band values."""
149
150
if int(band_values.size) < 2:
151
raise ValueError("band_values must contain at least 2 classes")
152
lower = np.asarray(band_values[:-1], dtype=float)
153
upper = np.asarray(band_values[1:], dtype=float)
154
return np.asarray((lower + upper) / 2.0, dtype=float)
155
156
157
def fit_qwk_cutpoints_coordinate_ascent(
158
*,
159
y_true: np.ndarray,
160
y_pred_raw: np.ndarray,
161
min_band: float,
162
max_band: float,
163
max_iter: int = 10,
164
candidate_quantiles: int = 64,
165
) -> CutpointCalibration:
166
"""Fit monotone cutpoints to maximize QWK via coordinate ascent.
167
168
Notes:
169
- This is a discrete, piecewise-constant objective. Coordinate ascent is deterministic and
170
fast enough for CV use at ELLIPSE scale.
171
- Predictions are clipped to [min_band, max_band] before thresholding.
172
- Initialization matches the marginal grade distribution using quantiles of predictions.
173
"""
174
175
band_vals = half_band_values(min_band=min_band, max_band=max_band)
176
n_classes = int(band_vals.size)
177
true_ids = _to_class_ids(y_true.astype(float), min_band=min_band)
178
179
preds = np.clip(y_pred_raw.astype(float), min_band, max_band)
180
if preds.size == 0:
181
raise ValueError("Empty y_pred_raw for cutpoint fitting")
182
183
# Quantile initialization: match predicted class proportions to true label proportions.
184
counts = np.bincount(true_ids, minlength=n_classes).astype(float)
185
cum = np.cumsum(counts)[:-1]
186
qs = (cum / float(np.sum(counts))).astype(float)
187
cutpoints = np.quantile(preds, qs).astype(float)
188
cutpoints = _enforce_strictly_increasing(cutpoints)
189
190
best_qwk = _qwk_from_class_ids(true_ids, np.digitize(preds, cutpoints, right=False), n_classes)
191
best = cutpoints.copy()
192
193
if candidate_quantiles < 8:
194
raise ValueError("candidate_quantiles must be >= 8 for stable cutpoint search")
195
196
for _ in range(max_iter):
197
improved = False
198
for idx in range(best.size):
199
low = best[idx - 1] if idx > 0 else float(min_band)
200
high = best[idx + 1] if idx + 1 < best.size else float(max_band)
201
if not low < high:
202
continue
203
204
mask = (preds > low) & (preds < high)
205
region = preds[mask]
206
if int(region.size) < 4:
207
continue
208
209
qs_region = np.linspace(0.05, 0.95, candidate_quantiles, dtype=float)
210
candidates = np.unique(np.quantile(region, qs_region).astype(float))
211
for candidate in candidates:
212
if not (low < float(candidate) < high):
213
continue
214
trial = best.copy()
215
trial[idx] = float(candidate)
216
trial = _enforce_strictly_increasing(trial)
217
if np.any(trial <= min_band) or np.any(trial >= max_band):
218
# Keep cutpoints inside the clipped prediction range.
219
continue
220
221
qwk = _qwk_from_class_ids(
222
true_ids, np.digitize(preds, trial, right=False), n_classes
223
)
224
if qwk > best_qwk + 1e-7:
225
best_qwk = qwk
226
best = trial
227
improved = True
228
229
if not improved:
230
break
231
232
return CutpointCalibration(
233
min_band=float(min_band),
234
max_band=float(max_band),
235
band_values=band_vals.astype(float),
236
cutpoints=best.astype(float),
237
)
238
239
240
def _to_class_ids(values: np.ndarray, *, min_band: float) -> np.ndarray:
241
"""Convert half-band floats to contiguous class IDs starting at min_band."""
242
243
min_half_id = int(round(float(min_band) * 2.0))
244
half_ids = np.round(values.astype(float) * 2.0).astype(int)
245
return half_ids - min_half_id
246
247
248
def _qwk_from_class_ids(y_true_ids: np.ndarray, y_pred_ids: np.ndarray, n_classes: int) -> float:
249
"""Fast QWK implementation using confusion matrix math (quadratic weights)."""
250
251
if int(n_classes) < 2:
252
return 0.0
253
254
y_true = y_true_ids.astype(int)
255
y_pred = y_pred_ids.astype(int)
256
y_true = np.clip(y_true, 0, n_classes - 1)
257
y_pred = np.clip(y_pred, 0, n_classes - 1)
258
259
n = int(y_true.size)
260
if n == 0:
261
return 0.0
262
263
conf = np.bincount(y_true * n_classes + y_pred, minlength=n_classes * n_classes).reshape(
264
(n_classes, n_classes)
265
)
266
true_hist = conf.sum(axis=1).astype(float)
267
pred_hist = conf.sum(axis=0).astype(float)
268
expected = np.outer(true_hist, pred_hist) / float(n)
269
270
w = np.zeros((n_classes, n_classes), dtype=float)
271
denom = float((n_classes - 1) ** 2)
272
for i in range(n_classes):
273
for j in range(n_classes):
274
w[i, j] = ((i - j) ** 2) / denom
275
276
observed_score = float(np.sum(w * conf))
277
expected_score = float(np.sum(w * expected))
278
if expected_score <= 0:
279
return 0.0
280
return 1.0 - (observed_score / expected_score)
281
282
283
def _enforce_strictly_increasing(values: np.ndarray, *, eps: float = 1e-6) -> np.ndarray:
284
"""Ensure a strictly increasing cutpoint vector (monotone thresholds)."""
285
286
out = values.astype(float).copy()
287
for i in range(1, out.size):
288
if out[i] <= out[i - 1]:
289
out[i] = out[i - 1] + eps
290
return out