Code / scripts/ml_training/essay_scoring/training/feature_selection.py
scripts/ml_training/essay_scoring/training/feature_selection.py
1
"""Predictor feature selection helpers for essay scoring experiments.
2
3
Purpose:
4
Provide a small, reproducible way to select a subset of *columns* from an existing feature
5
matrix (usually loaded from a feature store), without re-extracting features. This is used
6
for controlled experiments like "pruned handcrafted subset" while keeping embeddings intact.
7
8
Relationships:
9
- Used by `scripts.ml_training.essay_scoring.cross_validation` to pass `keep_feature_indices`
10
into fold training (`scripts.ml_training.essay_scoring.evaluation.oof.fold_execution`).
11
- Intended to keep feature-selection logic out of the runners so they stay below the repo's
12
file size limits.
13
"""
14
15
from __future__ import annotations
16
17
from collections.abc import Iterable, Sequence
18
19
from huleedu_nlp_shared.feature_contracts import is_candidate_sidecar_feature_name
20
21
from scripts.ml_training.json_types import JsonValue
22
23
from ..features.feature_schema_identity import (
24
CurrentScorerInputId,
25
serialized_feature_set_label,
26
)
27
28
_EMBEDDING_PREFIX = "embedding_"
29
30
31
class HandcraftedFeatureSelectionError(Exception):
32
"""Typed refusal for an invalid handcrafted keep/drop selection.
33
34
This is deliberately not a `ValueError`: pydantic converts `ValueError` and
35
`AssertionError` raised inside validators into a `ValidationError` while
36
propagating every other exception unchanged. Staying outside that
37
conversion lets the cv CLI seams present the refusal as one
38
`typer.BadParameter` diagnostic naming the offending selection and the
39
violated rule (the TASK-HULE-REP-0031 seam shape).
40
"""
41
42
def __init__(self, *, rule: str, detail: str) -> None:
43
self.rule = rule
44
self.detail = detail
45
super().__init__(f"{detail} [violated rule: {rule}]")
46
47
48
def resolve_keep_feature_indices(
49
feature_names: Sequence[str],
50
*,
51
scorer_input_id: CurrentScorerInputId,
52
handcrafted_keep: Sequence[str] | None,
53
handcrafted_drop: Sequence[str] | None = None,
54
) -> list[int] | None:
55
"""Resolve feature indices to keep for training.
56
57
For the combined scorer input, all embedding dimensions are always included. When an
58
allowlist of handcrafted features is provided, only those handcrafted features are kept
59
in addition to all embeddings. Alternatively, a denylist can drop specific handcrafted
60
features.
61
62
Args:
63
feature_names: Ordered list of feature names for the matrix columns.
64
scorer_input_id: Exact scorer input used to produce the matrix.
65
handcrafted_keep: Optional allowlist of handcrafted feature names to keep.
66
handcrafted_drop: Optional denylist of handcrafted feature names to drop.
67
68
Returns:
69
`None` when no filtering is requested (keep all columns). Otherwise a list of indices to
70
keep, preserving the original column order.
71
"""
72
73
if handcrafted_keep and handcrafted_drop:
74
raise HandcraftedFeatureSelectionError(
75
rule="exclusive_keep_or_drop",
76
detail=(
77
"Use only one of handcrafted_keep or handcrafted_drop (not both): "
78
f"handcrafted_keep={sorted(_normalize_keep_set(handcrafted_keep))}, "
79
f"handcrafted_drop={sorted(_normalize_keep_set(handcrafted_drop))}"
80
),
81
)
82
83
if not handcrafted_keep and not handcrafted_drop:
84
return None
85
86
if scorer_input_id == "embeddings.primary.v1":
87
requested_selection = sorted(
88
_normalize_keep_set([*(handcrafted_keep or []), *(handcrafted_drop or [])])
89
)
90
raise HandcraftedFeatureSelectionError(
91
rule="embeddings_only_selection_unsupported",
92
detail=(
93
"handcrafted_keep/handcrafted_drop is not supported for embeddings-only "
94
f"feature sets: requested={requested_selection}"
95
),
96
)
97
98
keep_set = _normalize_keep_set(handcrafted_keep or [])
99
drop_set = _normalize_keep_set(handcrafted_drop or [])
100
101
embedding_drops = sorted(name for name in drop_set if name.startswith(_EMBEDDING_PREFIX))
102
if embedding_drops:
103
raise HandcraftedFeatureSelectionError(
104
rule="embedding_dimensions_not_droppable",
105
detail=(
106
"Dropping embedding dimensions is not supported: "
107
f"handcrafted_drop={embedding_drops}"
108
),
109
)
110
111
requested = keep_set or drop_set
112
candidate_requests = sorted(
113
name for name in requested if is_candidate_sidecar_feature_name(name)
114
)
115
if candidate_requests:
116
raise ValueError(
117
"Candidate sidecar features require an explicit research-only candidate recipe: "
118
f"{candidate_requests}"
119
)
120
missing = sorted(name for name in requested if name not in feature_names)
121
if missing:
122
raise HandcraftedFeatureSelectionError(
123
rule="unknown_handcrafted_feature_names",
124
detail=f"Unknown handcrafted feature names: {missing}",
125
)
126
127
keep_indices: list[int] = []
128
for idx, name in enumerate(feature_names):
129
if name.startswith(_EMBEDDING_PREFIX):
130
keep_indices.append(idx)
131
elif keep_set and name in keep_set:
132
keep_indices.append(idx)
133
elif drop_set and name not in drop_set:
134
keep_indices.append(idx)
135
136
if not keep_indices:
137
raise ValueError("Feature selection produced an empty keep set.")
138
139
return keep_indices
140
141
142
def feature_selection_summary(
143
feature_names: Sequence[str],
144
*,
145
scorer_input_id: CurrentScorerInputId,
146
keep_feature_indices: Sequence[int] | None,
147
handcrafted_keep: Sequence[str] | None,
148
handcrafted_drop: Sequence[str] | None = None,
149
) -> dict[str, JsonValue]:
150
"""Build a small JSON-serializable summary of predictor feature selection."""
151
152
embedding_dim = sum(1 for name in feature_names if name.startswith(_EMBEDDING_PREFIX))
153
154
if keep_feature_indices is None:
155
return {
156
"mode": "all",
157
"feature_set": serialized_feature_set_label(scorer_input_id),
158
"embedding_dim": int(embedding_dim),
159
"feature_count_total": int(len(feature_names)),
160
"kept_feature_count": int(len(feature_names)),
161
"kept_handcrafted_feature_count": int(len(feature_names) - embedding_dim),
162
}
163
164
kept_names = [feature_names[int(idx)] for idx in keep_feature_indices]
165
kept_handcrafted = [name for name in kept_names if not name.startswith(_EMBEDDING_PREFIX)]
166
167
if handcrafted_drop:
168
return {
169
"mode": "handcrafted_denylist",
170
"feature_set": serialized_feature_set_label(scorer_input_id),
171
"embedding_dim": int(embedding_dim),
172
"feature_count_total": int(len(feature_names)),
173
"kept_feature_count": int(len(kept_names)),
174
"kept_handcrafted_feature_count": int(len(kept_handcrafted)),
175
"handcrafted_drop": _json_string_list(_normalize_keep_set(handcrafted_drop)),
176
"kept_handcrafted_features": _json_string_list(kept_handcrafted),
177
}
178
179
return {
180
"mode": "handcrafted_allowlist",
181
"feature_set": serialized_feature_set_label(scorer_input_id),
182
"embedding_dim": int(embedding_dim),
183
"feature_count_total": int(len(feature_names)),
184
"kept_feature_count": int(len(kept_names)),
185
"kept_handcrafted_feature_count": int(len(kept_handcrafted)),
186
"handcrafted_keep": _json_string_list(_normalize_keep_set(handcrafted_keep or [])),
187
"kept_handcrafted_features": _json_string_list(kept_handcrafted),
188
}
189
190
191
def _normalize_keep_set(values: Iterable[str]) -> set[str]:
192
keep_set: set[str] = set()
193
for raw in values:
194
cleaned = raw.strip()
195
if cleaned:
196
keep_set.add(cleaned)
197
return keep_set
198
199
200
def _json_string_list(values: Iterable[str]) -> list[JsonValue]:
201
return [value for value in values]