Code / scripts/ml_training/essay_scoring/features/pipeline.py
scripts/ml_training/essay_scoring/features/pipeline.py
1
"""Ledger-first feature extraction pipeline orchestration.
2
3
Purpose:
4
Provide a single high-level entrypoint (`FeaturePipeline.extract`) for producing model-ready
5
feature matrices from essay records. The pipeline coordinates:
6
- Writing issue ledger assembly over the source essay text.
7
- Deterministic parser-stabilized text derivation from spell corrections in the ledger.
8
- Optional Hemma offload (`/v1/extract`) for embeddings + handcrafted tiers.
9
- Local extraction for embeddings + tiered handcrafted features (Tier 1–2).
10
11
Relationships:
12
- Builds Tier 1 writing issue features from ledgers, then tiered features via
13
`tier1_error_readability` and
14
`tier2_syntactic_cohesion`.
15
- Optionally delegates remote work to
16
`scripts.ml_training.essay_scoring.offload.extract_client`.
17
- Outputs `FeatureMatrix` that is later consumed by training/cross-validation runners.
18
"""
19
20
from __future__ import annotations
21
22
import logging
23
import time
24
from dataclasses import dataclass
25
from typing import TYPE_CHECKING
26
27
import numpy as np
28
from huleedu_nlp_shared.feature_contracts import (
29
AcceptedFamilySemanticIdentityV1,
30
family_feature_names,
31
)
32
from huleedu_nlp_shared.writing_issues import (
33
NORMALIZED_SOURCE_UNREPAIRED,
34
PARSER_STABILIZED_REPAIRED,
35
ParserTextTreatment,
36
derive_embedding_feature_text_stream,
37
derive_spacy_feature_text_stream,
38
feature_inputs_for_parser_text_treatment,
39
)
40
41
from ..config import (
42
OffloadBackend,
43
OffloadConfig,
44
SpellNormalizationConfig,
45
)
46
from ..embedding_config import EmbeddingConfig
47
from ..features.combiner import FeatureMatrix
48
from ..features.feature_input import FeaturePipelineInput
49
from ..features.feature_outcome_projection import (
50
all_computed_outcomes,
51
compact_feature_provenance,
52
)
53
from ..features.feature_result_builder import (
54
build_feature_extraction_result,
55
)
56
from ..features.parser_text_authority_projection import (
57
ResearchParserTextArtifact,
58
)
59
from ..features.pipeline_helpers import (
60
applied_counts_for_selected_families,
61
feature_selection_for_pipeline,
62
required_consumer_ref,
63
)
64
from ..features.protocols import (
65
EmbeddingExtractorProtocol,
66
RemoteExtractClientProtocol,
67
)
68
from ..features.remote_extract_request import (
69
build_remote_extract_request,
70
)
71
from ..features.results import FeatureExtractionResult
72
from ..features.schema import feature_names_for
73
from ..features.writing_issue_ledger import (
74
ResearchLedgerArtifact,
75
ResearchLedgerAssemblerProtocol,
76
ResearchWritingIssueLedgerAssembler,
77
)
78
from ..offload.embedding_profile_contracts import (
79
CandidateEmbeddingProfileId,
80
)
81
from ..offload.extract_feature_selection import (
82
LEXICAL_CANDIDATE_FEATURE_SET_ID,
83
CandidateFeatureSetId,
84
ExtractAcceptedScorerInputV3,
85
)
86
from ..offload.metrics import OffloadMetricsCollector
87
from ..offload.prompt_relevance_mean_fit_transport import (
88
PromptRelevanceMeanFitTransportV1,
89
validate_prompt_relevance_mean_fit_transport,
90
)
91
from ..runtime.logging_utils import ProgressWriter
92
93
logger = logging.getLogger(__name__)
94
95
if TYPE_CHECKING:
96
from ..offload.extract_models import ExtractMeta
97
98
99
@dataclass
100
class FeaturePipeline:
101
"""Coordinate extraction of embeddings and tiered features."""
102
103
embedding_config: EmbeddingConfig
104
offload: OffloadConfig | None = None
105
spell_normalization: SpellNormalizationConfig | None = None
106
spacy_model: str = "en_core_web_sm"
107
embedder: EmbeddingExtractorProtocol | None = None
108
remote_extractor: RemoteExtractClientProtocol | None = None
109
ledger_assembler: ResearchLedgerAssemblerProtocol | None = None
110
offload_metrics: OffloadMetricsCollector | None = None
111
last_offload_meta: ExtractMeta | None = None
112
113
def extract(
114
self,
115
inputs: list[FeaturePipelineInput],
116
accepted: ExtractAcceptedScorerInputV3,
117
*,
118
family_semantic_identities: tuple[AcceptedFamilySemanticIdentityV1, ...],
119
progress: ProgressWriter | None = None,
120
candidate_feature_set_ids: tuple[CandidateFeatureSetId, ...] = (),
121
candidate_embedding_profile_ids: tuple[CandidateEmbeddingProfileId, ...] = (),
122
prepared_ledger_artifacts: list[ResearchLedgerArtifact] | None = None,
123
prepared_parser_text_artifacts: list[ResearchParserTextArtifact] | None = None,
124
prompt_relevance_mean_fit: PromptRelevanceMeanFitTransportV1 | None = None,
125
parser_text_treatment: ParserTextTreatment = PARSER_STABILIZED_REPAIRED,
126
) -> FeatureExtractionResult:
127
"""Extract features and provenance for a collection of records."""
128
129
self._validate_offload_configuration()
130
scorer_input_id = accepted.scorer_input_id
131
accepted_registry = accepted.registry_snapshot
132
selected_candidate_ids = candidate_feature_set_ids
133
feature_selection = feature_selection_for_pipeline(
134
accepted,
135
selected_candidate_ids,
136
)
137
selected_accepted_family_ids = accepted.requested_family_ids
138
selected_accepted_families = tuple(
139
accepted_registry.family(family_id) for family_id in selected_accepted_family_ids
140
)
141
selected_input_role_ids = {
142
role.role_id for family in selected_accepted_families for role in family.input_roles
143
}
144
if selected_candidate_ids and (
145
self.offload is None or self.offload.backend != OffloadBackend.HEMMA
146
):
147
raise ValueError("Candidate sidecars require backend=hemma.")
148
needs_linguistic_scalars = any(
149
family.value_kind == "scalar" for family in selected_accepted_families
150
)
151
if needs_linguistic_scalars and (
152
self.spell_normalization is None or not self.spell_normalization.enabled
153
):
154
raise RuntimeError(
155
"Accepted scalar extraction requires an explicit enabled spell normalization."
156
)
157
if needs_linguistic_scalars and (
158
self.offload is None or self.offload.backend != OffloadBackend.HEMMA
159
):
160
raise RuntimeError(
161
"Accepted handcrafted and combined scorer surfaces require Hemma/offload "
162
"extraction after TASK-0803 because accepted GUM/eRST primary-tree "
163
"features are parser-derived."
164
)
165
validate_prompt_relevance_mean_fit_transport(
166
feature_selection=feature_selection,
167
mean_fit=prompt_relevance_mean_fit,
168
)
169
prompts = [feature_input.prompt_label for feature_input in inputs]
170
ordered_prompt_authorities = [feature_input.prompt_authority for feature_input in inputs]
171
transport_prompt_authorities = (
172
ordered_prompt_authorities if "prompt_text" in selected_input_role_ids else []
173
)
174
ledger_artifacts = (
175
self.prepare_ledger_artifacts(
176
inputs=inputs,
177
require_language_tool="final_writing_issue_ledger" in selected_input_role_ids,
178
)
179
if prepared_ledger_artifacts is None
180
else prepared_ledger_artifacts
181
)
182
if len(ledger_artifacts) != len(inputs) or any(
183
artifact.source_text != feature_input.source_text
184
or artifact.offline_bundle.record_id != feature_input.record_id
185
for artifact, feature_input in zip(ledger_artifacts, inputs, strict=True)
186
):
187
raise ValueError("Prepared ledger artifacts do not match extraction records.")
188
if prepared_parser_text_artifacts is not None:
189
if parser_text_treatment == NORMALIZED_SOURCE_UNREPAIRED:
190
raise ValueError("prepared parser text requires a repaired parser treatment")
191
if len(prepared_parser_text_artifacts) != len(inputs) or any(
192
parser_artifact.source_text_ref != ledger_artifact.source_text_ref
193
for parser_artifact, ledger_artifact in zip(
194
prepared_parser_text_artifacts, ledger_artifacts, strict=True
195
)
196
):
197
raise ValueError("Prepared parser text artifacts do not match ledger records.")
198
syntax_parser_texts = [
199
derive_spacy_feature_text_stream(
200
(
201
artifact.source_text
202
if parser_text_treatment == NORMALIZED_SOURCE_UNREPAIRED
203
else artifact.parser_stabilized_text
204
)
205
if prepared_parser_text_artifacts is None
206
else prepared_parser_text_artifacts[index].text
207
)
208
for index, artifact in enumerate(ledger_artifacts)
209
]
210
primary_embedding_texts = [
211
derive_embedding_feature_text_stream(artifact.source_text)
212
for artifact in ledger_artifacts
213
]
214
treated_feature_inputs = [
215
(
216
feature_inputs_for_parser_text_treatment(
217
source_text=artifact.source_text,
218
feature_inputs=artifact.offline_bundle.policy_manifest.feature_inputs,
219
treatment=parser_text_treatment,
220
)
221
if prepared_parser_text_artifacts is None
222
else prepared_parser_text_artifacts[index].feature_inputs
223
)
224
for index, artifact in enumerate(ledger_artifacts)
225
]
226
syntax_parser_input_refs = [
227
feature_inputs.syntax_parser_input for feature_inputs in treated_feature_inputs
228
]
229
coherence_boundary_input_refs = [
230
required_consumer_ref(
231
feature_inputs.coherence_boundary_input,
232
"coherence_boundary_input",
233
)
234
for feature_inputs in treated_feature_inputs
235
]
236
rst_parser_input_refs = [
237
required_consumer_ref(
238
feature_inputs.rst_parser_input,
239
"rst_parser_input",
240
)
241
for feature_inputs in treated_feature_inputs
242
]
243
embedding_input_refs = [
244
artifact.offline_bundle.policy_manifest.feature_inputs.embedding_input
245
for artifact in ledger_artifacts
246
]
247
lexical_input_refs = [
248
feature_inputs.lexical_input for feature_inputs in treated_feature_inputs
249
]
250
needs_lexical_stream = (
251
"lexical_source_text" in selected_input_role_ids
252
or LEXICAL_CANDIDATE_FEATURE_SET_ID in selected_candidate_ids
253
)
254
lexical_candidate_texts = list(syntax_parser_texts) if needs_lexical_stream else None
255
parser_stabilized_policy_identities = [
256
(
257
artifact.offline_bundle.policy_manifest.parser_stabilized_policy_identity
258
if prepared_parser_text_artifacts is None
259
else prepared_parser_text_artifacts[index].derivation.policy_identity
260
)
261
for index, artifact in enumerate(ledger_artifacts)
262
]
263
writing_issue_ledgers = [artifact.ledger for artifact in ledger_artifacts]
264
applied_spelling_counts = applied_counts_for_selected_families(
265
prepared_parser_text_artifacts=prepared_parser_text_artifacts,
266
selected_family_ids=selected_accepted_family_ids,
267
)
268
269
if self.offload and self.offload.backend == OffloadBackend.HEMMA:
270
offload_service_url = self.offload.offload_service_url
271
if offload_service_url is None:
272
raise RuntimeError("backend=hemma requires offload_service_url.")
273
if self.offload_metrics is None:
274
self.offload_metrics = OffloadMetricsCollector()
275
276
from ..offload.extract_client import ( # noqa: PLC0415
277
RemoteExtractClient,
278
)
279
280
remote_extractor = self.remote_extractor
281
if remote_extractor is None:
282
remote_extractor = RemoteExtractClient(
283
base_url=offload_service_url,
284
offload_config=self.offload,
285
metrics=self.offload_metrics,
286
progress=progress,
287
)
288
self.remote_extractor = remote_extractor
289
290
logger.info(
291
"Offload extract start (backend=hemma, scorer_input_id=%s, records=%d)",
292
scorer_input_id,
293
len(syntax_parser_texts),
294
)
295
start = time.monotonic()
296
request = build_remote_extract_request(
297
syntax_parser_texts=syntax_parser_texts,
298
embedding_texts=primary_embedding_texts,
299
syntax_parser_input_refs=syntax_parser_input_refs,
300
coherence_boundary_input_refs=coherence_boundary_input_refs,
301
rst_parser_input_refs=rst_parser_input_refs,
302
embedding_input_refs=embedding_input_refs,
303
parser_stabilized_policy_identities=parser_stabilized_policy_identities,
304
prompts=prompts,
305
writing_issue_ledgers=(
306
writing_issue_ledgers
307
if "final_writing_issue_ledger" in selected_input_role_ids
308
else None
309
),
310
applied_spelling_counts=applied_spelling_counts,
311
semantic_prompt_authorities=transport_prompt_authorities,
312
feature_selection=feature_selection,
313
candidate_language=("en" if needs_lexical_stream else None),
314
lexical_candidate_texts=lexical_candidate_texts,
315
lexical_candidate_input_refs=(lexical_input_refs if needs_lexical_stream else None),
316
lexical_candidate_policy_identities=(
317
parser_stabilized_policy_identities if needs_lexical_stream else None
318
),
319
candidate_embedding_profile_ids=candidate_embedding_profile_ids,
320
prompt_relevance_mean_fit=prompt_relevance_mean_fit,
321
parser_text_treatment=parser_text_treatment,
322
)
323
result = remote_extractor.extract(
324
request,
325
family_semantic_identities=family_semantic_identities,
326
)
327
elapsed = time.monotonic() - start
328
logger.info(
329
"Offload extract complete (backend=hemma, records=%d) in %.2fs "
330
"(embeddings=%s, linguistic_scalars=%s)",
331
len(syntax_parser_texts),
332
elapsed,
333
None if result.embeddings is None else result.embeddings.shape,
334
(None if result.linguistic_scalars is None else result.linguistic_scalars.shape),
335
)
336
self.last_offload_meta = result.meta
337
provenance = compact_feature_provenance(
338
outcomes=result.outcomes,
339
ledger_artifacts=ledger_artifacts,
340
prompt_authorities=ordered_prompt_authorities,
341
)
342
343
selected_names = [
344
name
345
for family_id in selected_accepted_family_ids
346
for name in family_feature_names(accepted_registry.family(family_id))
347
]
348
if result.embeddings is not None and result.linguistic_scalars is None:
349
if result.embeddings is None:
350
raise RuntimeError(
351
"Offload returned no embeddings for embeddings-only feature set."
352
)
353
return build_feature_extraction_result(
354
feature_matrix=FeatureMatrix(
355
matrix=result.embeddings,
356
feature_names=selected_names,
357
),
358
ledger_artifacts=provenance.ledger_artifacts,
359
offload_meta=result.meta,
360
prompt_authorities=provenance.prompt_authorities,
361
outcomes=result.outcomes,
362
candidate_artifacts=result.candidate_artifacts,
363
candidate_embedding_artifacts=result.candidate_embedding_artifacts,
364
transport_attempts=result.transport_attempts,
365
)
366
367
if result.linguistic_scalars is not None and result.embeddings is None:
368
if result.linguistic_scalars is None:
369
raise RuntimeError("Offload returned no whitebox linguistic scalars.")
370
return build_feature_extraction_result(
371
feature_matrix=FeatureMatrix(
372
matrix=result.linguistic_scalars,
373
feature_names=selected_names,
374
),
375
ledger_artifacts=provenance.ledger_artifacts,
376
offload_meta=result.meta,
377
prompt_authorities=provenance.prompt_authorities,
378
outcomes=result.outcomes,
379
candidate_artifacts=result.candidate_artifacts,
380
candidate_embedding_artifacts=result.candidate_embedding_artifacts,
381
transport_attempts=result.transport_attempts,
382
)
383
384
if result.embeddings is None or result.linguistic_scalars is None:
385
raise RuntimeError(
386
"Offload returned incomplete feature payload for combined feature set."
387
)
388
matrix = np.concatenate([result.embeddings, result.linguistic_scalars], axis=1)
389
return build_feature_extraction_result(
390
feature_matrix=FeatureMatrix(
391
matrix=matrix,
392
feature_names=selected_names,
393
),
394
ledger_artifacts=provenance.ledger_artifacts,
395
offload_meta=result.meta,
396
prompt_authorities=provenance.prompt_authorities,
397
outcomes=result.outcomes,
398
candidate_artifacts=result.candidate_artifacts,
399
candidate_embedding_artifacts=result.candidate_embedding_artifacts,
400
transport_attempts=result.transport_attempts,
401
)
402
403
embedding_extractor = self._primary_embedding_extractor()
404
start = time.monotonic()
405
logger.info("Embedding extraction start (n=%d)", len(primary_embedding_texts))
406
embedding_vectors = embedding_extractor.embed(primary_embedding_texts)
407
elapsed = time.monotonic() - start
408
logger.info(
409
"Embedding extraction complete (shape=%s) in %.2fs",
410
embedding_vectors.shape,
411
elapsed,
412
)
413
return build_feature_extraction_result(
414
feature_matrix=FeatureMatrix(
415
matrix=embedding_vectors,
416
feature_names=feature_names_for(
417
scorer_input_id,
418
embedding_vectors.shape[1],
419
),
420
),
421
ledger_artifacts=ledger_artifacts,
422
prompt_authorities=ordered_prompt_authorities,
423
outcomes=all_computed_outcomes(len(inputs)),
424
)
425
426
def _primary_embedding_extractor(self) -> EmbeddingExtractorProtocol:
427
if self.embedder is None:
428
from ..features.embeddings import ( # noqa: PLC0415
429
TransformerMeanLast4Embedder,
430
)
431
432
self.embedder = TransformerMeanLast4Embedder(self.embedding_config)
433
return self.embedder
434
435
def _validate_offload_configuration(self) -> None:
436
if self.offload is None or self.offload.backend != OffloadBackend.HEMMA:
437
return
438
if not self.offload.offload_service_url:
439
raise RuntimeError("backend=hemma requires offload_service_url.")
440
if self.offload.language_tool_service_url:
441
raise RuntimeError(
442
"backend=hemma routes LanguageTool evidence through offload_service_url /v1/check; "
443
"do not pass language_tool_service_url for decision-grade Hemma runs."
444
)
445
446
def _assemble_ledgers(
447
self,
448
*,
449
inputs: list[FeaturePipelineInput],
450
require_language_tool: bool,
451
) -> list[ResearchLedgerArtifact]:
452
spell_config = self.spell_normalization or SpellNormalizationConfig()
453
language_tool_url = self._language_tool_evidence_url()
454
if require_language_tool and language_tool_url is None and self.ledger_assembler is None:
455
raise RuntimeError("LanguageTool service URL is required for research ledger evidence.")
456
if self.ledger_assembler is None:
457
self.ledger_assembler = ResearchWritingIssueLedgerAssembler(
458
spell_normalization=spell_config,
459
language="en-US",
460
language_tool_url=language_tool_url,
461
request_timeout_s=self.offload.request_timeout_s if self.offload else 60.0,
462
language_tool_cache_dir=self.offload.language_tool_cache_dir
463
if self.offload
464
else None,
465
)
466
return self.ledger_assembler.assemble_batch(
467
inputs,
468
require_language_tool=require_language_tool,
469
)
470
471
def prepare_ledger_artifacts(
472
self,
473
*,
474
inputs: list[FeaturePipelineInput],
475
require_language_tool: bool,
476
) -> list[ResearchLedgerArtifact]:
477
"""Build reusable exact ledger/input evidence before remote family lookup."""
478
479
return self._assemble_ledgers(
480
inputs=inputs,
481
require_language_tool=require_language_tool,
482
)
483
484
def _language_tool_evidence_url(self) -> str | None:
485
if self.offload is None:
486
return None
487
if self.offload.backend == OffloadBackend.HEMMA:
488
return self.offload.offload_service_url
489
return self.offload.language_tool_service_url