Files
DeepHealth/attention_types.py

27 lines
881 B
Python

"""Attention implementation identifiers and validation helpers."""
STANDARD_ATTENTION = "standard"
DIFF_ATTENTION = "diff"
DEFAULT_ATTENTION_TYPE = STANDARD_ATTENTION
SUPPORTED_ATTENTION_TYPES = (
STANDARD_ATTENTION,
DIFF_ATTENTION,
)
def resolve_attention_type(attention_type: str | None = None) -> str:
"""Resolve an attention selector, defaulting old configs to standard."""
resolved = DEFAULT_ATTENTION_TYPE if attention_type is None else attention_type
if not isinstance(resolved, str):
raise ValueError(
"attention_type must be one of "
f"{SUPPORTED_ATTENTION_TYPES}, got {resolved!r}"
)
if resolved not in SUPPORTED_ATTENTION_TYPES:
raise ValueError(
f"Unsupported attention_type={resolved!r}; "
f"expected one of {SUPPORTED_ATTENTION_TYPES}."
)
return resolved