Code / scripts/ml_training/essay_scoring/features/protocols.py
scripts/ml_training/essay_scoring/features/protocols.py
1
"""Protocols for essay-scoring feature extraction components.
2
3
Purpose:
4
Define small `typing.Protocol` interfaces used by the research feature pipeline
5
to decouple orchestration code from concrete implementations.
6
7
Relationships:
8
- Consumed by `scripts.ml_training.essay_scoring.features.pipeline.FeaturePipeline`.
9
- Implemented by local and offloaded embedding extractors and tier extractors.
10
"""
11
12
from __future__ import annotations
13
14
from collections.abc import Sequence
15
from typing import TYPE_CHECKING, Protocol
16
17
import numpy as np
18
from huleedu_nlp_shared.feature_contracts import AcceptedFamilySemanticIdentityV1
19
from huleedu_nlp_shared.writing_issues import WritingIssueLedgerV2
20
21
from ..datasets.prompt_authority import SemanticPromptAuthority
22
from ..features.schema import (
23
ConstructComplexNominalFeatures,
24
Tier1Features,
25
Tier2Features,
26
)
27
28
if TYPE_CHECKING:
29
from spacy.tokens import Doc
30
31
from ..features.tier2_syntactic_cohesion import (
32
Tier2BatchResult,
33
)
34
from ..offload.extract_client import RemoteExtractResult
35
from ..offload.extract_request_models import ExtractRequest
36
37
38
class EmbeddingExtractorProtocol(Protocol):
39
"""Protocol for embedding extraction."""
40
41
def embed(self, texts: list[str]) -> np.ndarray:
42
"""Return embeddings for a list of texts."""
43
44
45
class RemoteExtractClientProtocol(Protocol):
46
"""Narrow remote boundary that accepts one fully validated request."""
47
48
def extract(
49
self,
50
request: ExtractRequest,
51
*,
52
family_semantic_identities: tuple[AcceptedFamilySemanticIdentityV1, ...],
53
) -> RemoteExtractResult:
54
"""Return accepted and requested candidate features for one request."""
55
56
57
class Tier1ExtractorProtocol(Protocol):
58
"""Protocol for Tier 1 feature extraction."""
59
60
def extract(self, text: str, *, writing_issue_ledger: WritingIssueLedgerV2) -> Tier1Features:
61
"""Extract Tier 1 features from a single essay."""
62
63
def extract_batch(
64
self,
65
texts: list[str],
66
*,
67
writing_issue_ledgers: list[WritingIssueLedgerV2],
68
) -> list[Tier1Features]:
69
"""Extract Tier 1 features for a batch of essays.
70
71
Args:
72
texts: Ledger-derived feature texts used for spaCy-tokenized Tier 1 statistics.
73
writing_issue_ledgers: Final source-text ledgers used for issue count features.
74
"""
75
76
77
class Tier2ExtractorProtocol(Protocol):
78
"""Protocol for Tier 2 feature extraction."""
79
80
def extract(self, text: str, prompt_authority: SemanticPromptAuthority) -> Tier2Features:
81
"""Extract Tier 2 features from a single essay and governed prompt authority."""
82
83
def extract_batch(
84
self,
85
texts: list[str],
86
prompt_authorities: list[SemanticPromptAuthority],
87
) -> list[Tier2Features]:
88
"""Extract Tier 2 features for a batch of essays and governed prompt authorities."""
89
90
def extract_batch_with_artifacts(
91
self,
92
texts: list[str],
93
prompt_authorities: list[SemanticPromptAuthority],
94
) -> Tier2BatchResult:
95
"""Extract Tier 2 features and parse artifacts for downstream features."""
96
97
98
class ConstructComplexNominalExtractorProtocol(Protocol):
99
"""Protocol for accepted construct complex-nominal feature extraction."""
100
101
def extract_batch_from_docs(
102
self,
103
docs: Sequence[Doc],
104
) -> list[ConstructComplexNominalFeatures]:
105
"""Extract accepted construct features from parsed spaCy docs."""