DeepHealth/model_architectures.py

"""Model-architecture identifiers and checkpoint validation helpers."""

from __future__ import annotations

import re
from collections.abc import Mapping


TRANSFORMER_FFN_ARCHITECTURE = "transformer_ffn_v1"
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
DEFAULT_MODEL_ARCHITECTURE = TRANSFORMER_FFN_ARCHITECTURE
SUPPORTED_MODEL_ARCHITECTURES = (
    TRANSFORMER_FFN_ARCHITECTURE,
    TRAJ_MIXER_ARCHITECTURE,
)


_FFN_STATE_KEY = re.compile(
    r"(?:^|\.)blocks\.\d+\.mlp\.w[123]\.(?:weight|bias)$"
)
_TRAJ_MIXER_STATE_KEY = re.compile(
    r"(?:^|\.)blocks\.\d+\.mlp\.(?:"
    r"norm\.(?:weight|bias)|"
    r"intra_gate_proj|"
    r"intra_value_proj|"
    r"intra_output_proj|"
    r"intra_gate_logits|"
    r"gate_proj|"
    r"value_proj|"
    r"output_proj"
    r")$"
)


def _validate_model_architecture(model_architecture: object) -> str:
    if not isinstance(model_architecture, str):
        raise ValueError(
            "model_architecture must be one of "
            f"{SUPPORTED_MODEL_ARCHITECTURES}, got {model_architecture!r}"
        )
    if model_architecture not in SUPPORTED_MODEL_ARCHITECTURES:
        raise ValueError(
            f"Unsupported model_architecture={model_architecture!r}; "
            f"expected one of {SUPPORTED_MODEL_ARCHITECTURES}."
        )
    return model_architecture


def detect_model_architecture_from_state_dict(
    state_dict: Mapping[str, object],
) -> str:
    """Infer the architecture from block parameter names.

    Detection deliberately accepts any ``blocks.<index>`` prefix rather than
    assuming that block zero is present.
    """

    if not isinstance(state_dict, Mapping):
        raise TypeError(
            "state_dict must be a mapping, got "
            f"{type(state_dict).__name__}"
        )

    has_ffn = False
    has_traj_mixer = False
    for raw_key in state_dict:
        key = str(raw_key)
        has_ffn = has_ffn or _FFN_STATE_KEY.search(key) is not None
        has_traj_mixer = (
            has_traj_mixer
            or _TRAJ_MIXER_STATE_KEY.search(key) is not None
        )
        if has_ffn and has_traj_mixer:
            raise ValueError(
                "Checkpoint contains both Transformer FFN and TrajMixer "
                "block parameters; its model architecture is ambiguous."
            )

    if has_ffn:
        return TRANSFORMER_FFN_ARCHITECTURE
    if has_traj_mixer:
        return TRAJ_MIXER_ARCHITECTURE
    raise ValueError(
        "Could not detect model architecture from checkpoint parameters. "
        "Expected a blocks.<index>.mlp FFN or TrajMixer parameter."
    )


def resolve_model_architecture(
    config_or_marker: Mapping[str, object] | str | None = None,
    state_dict: Mapping[str, object] | None = None,
) -> str:
    """Resolve and cross-check a configured and checkpoint architecture.

    Every saved run must provide an explicit architecture marker. Checkpoint
    parameter names are used only to verify that the marker describes the
    weights being loaded.
    """

    if isinstance(config_or_marker, Mapping):
        configured = config_or_marker.get("model_architecture")
    elif isinstance(config_or_marker, str) or config_or_marker is None:
        configured = config_or_marker
    else:
        raise TypeError(
            "config_or_marker must be a config mapping, string, or None, got "
            f"{type(config_or_marker).__name__}"
        )

    resolved_config = (
        _validate_model_architecture(configured)
        if configured is not None
        else None
    )
    detected = (
        detect_model_architecture_from_state_dict(state_dict)
        if state_dict is not None
        else None
    )

    if resolved_config is None:
        raise ValueError(
            "model_architecture is required; expected one of "
            f"{SUPPORTED_MODEL_ARCHITECTURES}."
        )
    if detected is not None and resolved_config != detected:
        raise ValueError(
            "Configured model architecture conflicts with checkpoint: "
            f"config={resolved_config!r}, checkpoint={detected!r}."
        )
    return resolved_config