Code / scripts/ml_training/essay_scoring/features/embeddings.py
scripts/ml_training/essay_scoring/features/embeddings.py
1
"""Transformer embedding extraction for essay scoring accepted vectors.
2
3
Purpose:
4
Provide a small, model-agnostic embedding extractor that can load any
5
Hugging Face encoder model via `transformers.AutoModel` and return a
6
fixed-size float32 embedding per input text using the accepted scorer
7
mean-last4 hidden-state pooling contract.
8
9
Relationships:
10
- Instantiated by `scripts.ml_training.essay_scoring.features.pipeline.FeaturePipeline`
11
for local embedding extraction when Hemma offload is not configured.
12
- Instantiated by the Hemma offload server via
13
`scripts.ml_training.essay_scoring.offload.embedder_factory`.
14
"""
15
16
from __future__ import annotations
17
18
from dataclasses import dataclass
19
20
import numpy as np
21
import torch
22
from transformers import AutoModel, AutoTokenizer
23
24
from ..embedding_config import EmbeddingConfig
25
from ..offload.embedding_pooling_policies import (
26
attention_mask_mean_last4,
27
)
28
29
30
@dataclass
31
class TransformerMeanLast4Embedder:
32
"""Embedding extractor using attention-mask mean pooling over last four layers."""
33
34
config: EmbeddingConfig
35
36
def __post_init__(self) -> None:
37
self._device = self._resolve_device(self.config.device)
38
self._effective_batch_size = int(self.config.batch_size)
39
self._tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
40
self._model = AutoModel.from_pretrained(self.config.model_name)
41
self._model.to(self._device)
42
self._model.eval()
43
44
def embed(self, texts: list[str]) -> np.ndarray:
45
"""Embed texts using attention-mask mean pooling over the last four layers.
46
47
Args:
48
texts: Input texts to embed. Empty input returns an empty float32 matrix.
49
50
Returns:
51
Float32 matrix shaped [len(texts), hidden_size].
52
"""
53
54
if not texts:
55
return np.empty((0, 0), dtype=np.float32)
56
57
batches = [
58
texts[idx : idx + self._effective_batch_size]
59
for idx in range(0, len(texts), self._effective_batch_size)
60
]
61
embeddings: list[np.ndarray] = []
62
63
try:
64
for batch in batches:
65
embeddings.append(self._embed_batch_with_cuda_adaptation(batch))
66
67
return np.vstack(embeddings).astype(np.float32, copy=False)
68
finally:
69
self._release_cuda_cache()
70
71
@property
72
def resolved_device(self) -> str:
73
"""Return the effective torch device used by this embedder."""
74
75
return str(self._device)
76
77
@property
78
def effective_batch_size(self) -> int:
79
"""Return the current embedding batch size after any CUDA OOM adaptation."""
80
81
return int(self._effective_batch_size)
82
83
def _embed_batch_with_cuda_adaptation(self, batch: list[str]) -> np.ndarray:
84
try:
85
return self._embed_batch(batch)
86
except RuntimeError as exc:
87
if not self._is_cuda_oom(exc) or len(batch) <= 1:
88
raise
89
self._release_cuda_cache()
90
split_at = max(1, len(batch) // 2)
91
self._effective_batch_size = max(1, min(self._effective_batch_size, split_at))
92
left = self._embed_batch_with_cuda_adaptation(batch[:split_at])
93
right = self._embed_batch_with_cuda_adaptation(batch[split_at:])
94
return np.vstack((left, right)).astype(np.float32, copy=False)
95
96
def _embed_batch(self, batch: list[str]) -> np.ndarray:
97
encoded = self._tokenizer(
98
batch,
99
padding=True,
100
truncation=True,
101
max_length=self.config.max_length,
102
return_tensors="pt",
103
)
104
encoded = {key: value.to(self._device) for key, value in encoded.items()}
105
with torch.no_grad():
106
outputs = self._model(**encoded, output_hidden_states=True)
107
embeddings = attention_mask_mean_last4(outputs, encoded["attention_mask"])
108
return np.asarray(embeddings.detach().cpu().numpy(), dtype=np.float32)
109
110
@staticmethod
111
def _is_cuda_oom(exc: RuntimeError) -> bool:
112
message = str(exc).lower()
113
return "cuda" in message and "out of memory" in message
114
115
def _release_cuda_cache(self) -> None:
116
if self._device.type == "cuda" and torch.cuda.is_available():
117
torch.cuda.empty_cache()
118
119
@staticmethod
120
def _resolve_device(device_override: str | None) -> torch.device:
121
"""Resolve the torch device for embeddings."""
122
123
if device_override:
124
device = torch.device(device_override)
125
if device.type == "cuda" and not torch.cuda.is_available():
126
raise RuntimeError("embedding_cuda_device_configured_but_unavailable")
127
return device
128
if torch.cuda.is_available():
129
return torch.device("cuda")
130
if torch.backends.mps.is_available():
131
return torch.device("mps")
132
return torch.device("cpu")