Code / scripts/ml_training/essay_scoring/embedding_config.py
scripts/ml_training/essay_scoring/embedding_config.py
1
"""Import-light configuration authority for local embedding extraction.
2
3
Purpose:
4
Define the model aliases and validated values required to construct the
5
local embedding implementation without importing research pipeline runtime
6
configuration.
7
8
Relationships:
9
- Used by the local embedding implementation and CLI configuration callers.
10
- Embedded by ``ExperimentConfig`` in the private research configuration
11
module without importing its paths, caches, or offload settings here.
12
"""
13
14
from __future__ import annotations
15
16
from enum import Enum
17
18
from pydantic import BaseModel, ConfigDict
19
20
21
class EmbeddingBackbone(str, Enum):
22
"""Short aliases for supported embedding backbones.
23
24
These aliases exist to reduce CLI/config churn and avoid typos when running
25
controlled comparisons across multiple Hugging Face checkpoints.
26
"""
27
28
DEBERTA_V3_BASE = "deberta_v3_base"
29
ROBERTA_BASE = "roberta_base"
30
MODERNBERT_BASE = "modernbert_base"
31
32
33
EMBEDDING_BACKBONE_MODEL_NAMES: dict[EmbeddingBackbone, str] = {
34
EmbeddingBackbone.DEBERTA_V3_BASE: "microsoft/deberta-v3-base",
35
EmbeddingBackbone.ROBERTA_BASE: "roberta-base",
36
EmbeddingBackbone.MODERNBERT_BASE: "answerdotai/ModernBERT-base",
37
}
38
39
40
def resolve_embedding_model_name(*, backbone: EmbeddingBackbone | None, model_name: str) -> str:
41
"""Resolve an effective Hugging Face model name from a backbone alias.
42
43
Args:
44
backbone: Optional short alias that maps to a known Hugging Face model id.
45
model_name: Explicit Hugging Face model id to use when no alias is set.
46
47
Returns:
48
Hugging Face model id to pass to Transformers ``from_pretrained(...)``.
49
"""
50
51
if backbone is None:
52
return model_name
53
return EMBEDDING_BACKBONE_MODEL_NAMES[backbone]
54
55
56
class EmbeddingConfig(BaseModel):
57
"""Configuration for embedding extraction."""
58
59
model_config = ConfigDict(extra="forbid")
60
61
model_name: str = "microsoft/deberta-v3-base"
62
max_length: int = 512
63
batch_size: int = 8
64
device: str | None = None