modeling_mint_stability.py

"""
Self-contained model file for HuggingFace Hub (trust_remote_code=True).

Inlines all dependencies from the MINT package (Ullanat et al., 2026) so that
users can load the model without installing mint:

    from transformers import AutoModel
    model = AutoModel.from_pretrained("...", trust_remote_code=True)

Original MINT source: https://github.com/VarunUllanat/mint (MIT License)
ESM2 backbone: Meta Platforms, Inc. (MIT License)

Generated by: python -m mint_stability.build_hf
"""

import inspect
import itertools
import math
import os
import uuid
from typing import Dict, List, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.nn import Parameter
from transformers import PreTrainedModel

try:
    from .configuration_mint_stability import MintStabilityConfig
except ImportError:  # standalone / direct import (no package context)
    import importlib as _il
    _cfg = _il.import_module("configuration_mint_stability")
    MintStabilityConfig = _cfg.MintStabilityConfig


# --- torch.load compat (inlined from _compat.py) ---
_TORCH_LOAD_PARAMS = set(inspect.signature(torch.load).parameters)

def torch_load(f, **kwargs):
    if "weights_only" not in _TORCH_LOAD_PARAMS:
        kwargs.pop("weights_only", None)
    return torch.load(f, **kwargs)


# ============================================================================
# Backbone (inlined from mint_stability/backbone.py)
# ============================================================================

# ============================================================================
# Constants (from mint/constants.py)
# ============================================================================

proteinseq_toks = {
    "toks": [
        "L", "A", "G", "V", "S", "E", "R", "T", "I", "D",
        "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C",
        "X", "B", "U", "Z", "O", ".", "-",
    ]
}

# ============================================================================
# Alphabet (from mint/data.py)
# ============================================================================


class Alphabet(object):
    def __init__(
        self,
        standard_toks: Sequence[str],
        prepend_toks: Sequence[str] = ("<null_0>", "<pad>", "<eos>", "<unk>"),
        append_toks: Sequence[str] = ("<cls>", "<mask>", "<sep>"),
        prepend_bos: bool = True,
        append_eos: bool = False,
        use_msa: bool = False,
    ):
        self.standard_toks = list(standard_toks)
        self.prepend_toks = list(prepend_toks)
        self.append_toks = list(append_toks)
        self.prepend_bos = prepend_bos
        self.append_eos = append_eos
        self.use_msa = use_msa

        self.all_toks = list(self.prepend_toks)
        self.all_toks.extend(self.standard_toks)
        for i in range((8 - (len(self.all_toks) % 8)) % 8):
            self.all_toks.append(f"<null_{i + 1}>")
        self.all_toks.extend(self.append_toks)

        self.tok_to_idx = {tok: i for i, tok in enumerate(self.all_toks)}

        self.unk_idx = self.tok_to_idx["<unk>"]
        self.padding_idx = self.get_idx("<pad>")
        self.cls_idx = self.get_idx("<cls>")
        self.mask_idx = self.get_idx("<mask>")
        self.eos_idx = self.get_idx("<eos>")
        self.all_special_tokens = ["<eos>", "<unk>", "<pad>", "<cls>", "<mask>"]
        self.unique_no_split_tokens = self.all_toks

    def __len__(self):
        return len(self.all_toks)

    def get_idx(self, tok):
        return self.tok_to_idx.get(tok, self.unk_idx)

    def get_tok(self, ind):
        return self.all_toks[ind]

    def to_dict(self):
        return self.tok_to_idx.copy()

    @classmethod
    def from_architecture(cls, name: str) -> "Alphabet":
        if name in ("ESM-1b", "roberta_large"):
            standard_toks = proteinseq_toks["toks"]
            prepend_toks = ("<cls>", "<pad>", "<eos>", "<unk>")
            append_toks = ("<mask>",)
            prepend_bos = True
            append_eos = True
            use_msa = False
        else:
            raise ValueError(f"Unknown architecture: {name}")
        return cls(standard_toks, prepend_toks, append_toks, prepend_bos, append_eos, use_msa)

    def _tokenize(self, text) -> str:
        return text.split()

    def tokenize(self, text, **kwargs) -> List[str]:
        def split_on_token(tok, text):
            result = []
            split_text = text.split(tok)
            for i, sub_text in enumerate(split_text):
                if i < len(split_text) - 1:
                    sub_text = sub_text.rstrip()
                if i > 0:
                    sub_text = sub_text.lstrip()
                if i == 0 and not sub_text:
                    result.append(tok)
                elif i == len(split_text) - 1:
                    if sub_text:
                        result.append(sub_text)
                else:
                    if sub_text:
                        result.append(sub_text)
                    result.append(tok)
            return result

        def split_on_tokens(tok_list, text):
            if not text.strip():
                return []
            tokenized_text = []
            text_list = [text]
            for tok in tok_list:
                tokenized_text = []
                for sub_text in text_list:
                    if sub_text not in self.unique_no_split_tokens:
                        tokenized_text.extend(split_on_token(tok, sub_text))
                    else:
                        tokenized_text.append(sub_text)
                text_list = tokenized_text
            return list(
                itertools.chain.from_iterable(
                    self._tokenize(token)
                    if token not in self.unique_no_split_tokens
                    else [token]
                    for token in tokenized_text
                )
            )

        no_split_token = self.unique_no_split_tokens
        tokenized_text = split_on_tokens(no_split_token, text)
        return tokenized_text

    def encode(self, text):
        return [self.get_idx(tok) for tok in self.tokenize(text)]


# ============================================================================
# Rotary embeddings (from mint/rotary_embedding.py)
# ============================================================================


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(x, cos, sin):
    cos = cos[:, : x.shape[-2], :]
    sin = sin[:, : x.shape[-2], :]
    return (x * cos) + (rotate_half(x) * sin)


class RotaryEmbedding(torch.nn.Module):
    def __init__(self, dim: int, *_, **__):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self._seq_len_cached = None
        self._cos_cached = None
        self._sin_cached = None

    def _update_cos_sin_tables(self, x, seq_dimension=1):
        seq_len = x.shape[seq_dimension]
        if (seq_len != self._seq_len_cached
                or self._cos_cached.device != x.device
                or self._cos_cached.dtype != x.dtype):
            self._seq_len_cached = seq_len
            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(
                self.inv_freq
            )
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self._cos_cached = emb.cos()[None, :, :].to(x.dtype)
            self._sin_cached = emb.sin()[None, :, :].to(x.dtype)
        return self._cos_cached, self._sin_cached

    def forward(
        self, q: torch.Tensor, k: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(
            k, seq_dimension=-2
        )
        return (
            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
        )


# ============================================================================
# Multi-head attention (from mint/multihead_attention.py)
# ============================================================================


def utils_softmax(x, dim: int, onnx_trace: bool = False):
    if onnx_trace:
        return F.softmax(x.float(), dim=dim)
    else:
        return F.softmax(x, dim=dim, dtype=torch.float32)


class FairseqIncrementalState(object):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.init_incremental_state()

    def init_incremental_state(self):
        self._incremental_state_id = str(uuid.uuid4())

    def _get_full_incremental_state_key(self, key: str) -> str:
        return "{}.{}".format(self._incremental_state_id, key)

    def get_incremental_state(
        self,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
        key: str,
    ) -> Optional[Dict[str, Optional[Tensor]]]:
        full_key = self._get_full_incremental_state_key(key)
        if incremental_state is None or full_key not in incremental_state:
            return None
        return incremental_state[full_key]

    def set_incremental_state(
        self,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
        key: str,
        value: Dict[str, Optional[Tensor]],
    ) -> Optional[Dict[str, Dict[str, Optional[Tensor]]]]:
        if incremental_state is not None:
            full_key = self._get_full_incremental_state_key(key)
            incremental_state[full_key] = value
        return incremental_state


def with_incremental_state(cls):
    cls.__bases__ = (FairseqIncrementalState,) + tuple(
        b for b in cls.__bases__ if b != FairseqIncrementalState
    )
    return cls


@with_incremental_state
class MultiheadAttention(nn.Module):
    """Multi-headed attention (Vaswani et al., 2017)."""

    def __init__(
        self,
        embed_dim,
        num_heads,
        kdim=None,
        vdim=None,
        dropout=0.0,
        bias=True,
        add_bias_kv: bool = False,
        add_zero_attn: bool = False,
        self_attention: bool = False,
        encoder_decoder_attention: bool = False,
        use_rotary_embeddings: bool = False,
        no_proj=False,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.kdim = kdim if kdim is not None else embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim
        self.scaling = self.head_dim ** -0.5
        self.self_attention = self_attention
        self.encoder_decoder_attention = encoder_decoder_attention

        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        if no_proj:
            self.out_proj = None
        else:
            self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        if add_bias_kv:
            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
        else:
            self.bias_k = self.bias_v = None

        self.add_zero_attn = add_zero_attn
        self.reset_parameters()
        self.onnx_trace = False
        self.rot_emb = None
        if use_rotary_embeddings:
            self.rot_emb = RotaryEmbedding(dim=self.head_dim)

    def reset_parameters(self):
        if self.qkv_same_dim:
            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
        else:
            nn.init.xavier_uniform_(self.k_proj.weight)
            nn.init.xavier_uniform_(self.v_proj.weight)
            nn.init.xavier_uniform_(self.q_proj.weight)
        if self.out_proj is not None:
            nn.init.xavier_uniform_(self.out_proj.weight)
            if self.out_proj.bias is not None:
                nn.init.constant_(self.out_proj.bias, 0.0)
        if self.bias_k is not None:
            nn.init.xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            nn.init.xavier_normal_(self.bias_v)

    def forward(
        self,
        query,
        key: Optional[Tensor],
        value: Optional[Tensor],
        key_padding_mask: Optional[Tensor] = None,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
        need_weights: bool = True,
        static_kv: bool = False,
        attn_mask: Optional[Tensor] = None,
        before_softmax: bool = False,
        need_head_weights: bool = False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        if need_head_weights:
            need_weights = True

        tgt_len, bsz, embed_dim = query.size()
        assert embed_dim == self.embed_dim

        if incremental_state is not None:
            saved_state = self._get_input_buffer(incremental_state)
            if saved_state is not None and "prev_key" in saved_state:
                if static_kv:
                    key = value = None
        else:
            saved_state = None

        if self.self_attention:
            q = self.q_proj(query)
            k = self.k_proj(query)
            v = self.v_proj(query)
        elif self.encoder_decoder_attention:
            q = self.q_proj(query)
            if key is None:
                assert value is None
                k = v = None
            else:
                k = self.k_proj(key)
                v = self.v_proj(key)
        else:
            assert key is not None and value is not None
            q = self.q_proj(query)
            k = self.k_proj(key)
            v = self.v_proj(value)
        q *= self.scaling

        if self.bias_k is not None:
            assert self.bias_v is not None
            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
            if attn_mask is not None:
                attn_mask = torch.cat(
                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
                )
            if key_padding_mask is not None:
                key_padding_mask = torch.cat(
                    [key_padding_mask, key_padding_mask.new_zeros(key_padding_mask.size(0), 1)],
                    dim=1,
                )

        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if k is not None:
            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if v is not None:
            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)

        if saved_state is not None:
            if "prev_key" in saved_state:
                _prev_key = saved_state["prev_key"]
                assert _prev_key is not None
                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
                if static_kv:
                    k = prev_key
                else:
                    assert k is not None
                    k = torch.cat([prev_key, k], dim=1)
            if "prev_value" in saved_state:
                _prev_value = saved_state["prev_value"]
                assert _prev_value is not None
                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
                if static_kv:
                    v = prev_value
                else:
                    assert v is not None
                    v = torch.cat([prev_value, v], dim=1)
            prev_key_padding_mask: Optional[Tensor] = None
            if "prev_key_padding_mask" in saved_state:
                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
            assert k is not None and v is not None
            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
                key_padding_mask=key_padding_mask,
                prev_key_padding_mask=prev_key_padding_mask,
                batch_size=bsz,
                src_len=k.size(1),
                static_kv=static_kv,
            )
            saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
            saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
            saved_state["prev_key_padding_mask"] = key_padding_mask
            assert incremental_state is not None
            incremental_state = self._set_input_buffer(incremental_state, saved_state)

        assert k is not None
        src_len = k.size(1)

        if key_padding_mask is not None and key_padding_mask.dim() == 0:
            key_padding_mask = None

        if key_padding_mask is not None:
            assert key_padding_mask.size(0) == bsz
            assert key_padding_mask.size(1) == src_len

        if self.add_zero_attn:
            assert v is not None
            src_len += 1
            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
            if attn_mask is not None:
                attn_mask = torch.cat(
                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
                )
            if key_padding_mask is not None:
                key_padding_mask = torch.cat(
                    [key_padding_mask, torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask)],
                    dim=1,
                )

        if self.rot_emb:
            q, k = self.rot_emb(q, k)

        attn_weights = torch.bmm(q, k.transpose(1, 2))
        attn_weights = MultiheadAttention.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]

        if attn_mask is not None:
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.masked_fill(attn_mask.unsqueeze(1), float("-inf"))
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if key_padding_mask is not None:
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
            )
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if before_softmax:
            return (
                attn_weights.view(bsz, self.num_heads, tgt_len, src_len),
                v.view(bsz, self.num_heads, src_len, self.head_dim),
            )

        attn_weights_float = utils_softmax(attn_weights, dim=-1, onnx_trace=self.onnx_trace)
        attn_weights = attn_weights_float.type_as(attn_weights)
        attn_probs = F.dropout(
            attn_weights_float.type_as(attn_weights),
            p=self.dropout,
            training=self.training,
        )
        assert v is not None
        attn = torch.bmm(attn_probs, v)
        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
        attn = self.out_proj(attn)
        attn_weights_result: Optional[Tensor] = None
        if need_weights:
            attn_weights_result = (
                attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len)
                .type_as(attn)
                .transpose(1, 0)
            )
            if not need_head_weights:
                attn_weights_result = attn_weights_result.mean(dim=0)
        return attn, attn_weights_result

    @staticmethod
    def _append_prev_key_padding_mask(
        key_padding_mask: Optional[Tensor],
        prev_key_padding_mask: Optional[Tensor],
        batch_size: int,
        src_len: int,
        static_kv: bool,
    ) -> Optional[Tensor]:
        if prev_key_padding_mask is not None and static_kv:
            new_key_padding_mask = prev_key_padding_mask
        elif prev_key_padding_mask is not None and key_padding_mask is not None:
            new_key_padding_mask = torch.cat(
                [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
            )
        elif prev_key_padding_mask is not None:
            filler = torch.zeros(
                (batch_size, src_len - prev_key_padding_mask.size(1)),
                device=prev_key_padding_mask.device,
            )
            new_key_padding_mask = torch.cat(
                [prev_key_padding_mask.float(), filler.float()], dim=1
            )
        elif key_padding_mask is not None:
            filler = torch.zeros(
                (batch_size, src_len - key_padding_mask.size(1)),
                device=key_padding_mask.device,
            )
            new_key_padding_mask = torch.cat(
                [filler.float(), key_padding_mask.float()], dim=1
            )
        else:
            new_key_padding_mask = prev_key_padding_mask
        return new_key_padding_mask

    def _get_input_buffer(
        self,
        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
    ) -> Dict[str, Optional[Tensor]]:
        result = self.get_incremental_state(incremental_state, "attn_state")
        if result is not None:
            return result
        else:
            empty_result: Dict[str, Optional[Tensor]] = {}
            return empty_result

    def _set_input_buffer(
        self,
        incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
        buffer: Dict[str, Optional[Tensor]],
    ):
        return self.set_incremental_state(incremental_state, "attn_state", buffer)

    def apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int):
        return attn_weights


# ============================================================================
# Transformer modules (from mint/modules.py)
# ============================================================================


def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))


try:
    from apex.normalization import FusedLayerNorm as _FusedLayerNorm

    class ESM1bLayerNorm(_FusedLayerNorm):
        @torch.jit.unused
        def forward(self, x):
            if not x.is_cuda:
                return super().forward(x)
            else:
                with torch.cuda.device(x.device):
                    return super().forward(x)

except ImportError:
    from torch.nn import LayerNorm as ESM1bLayerNorm


class TransformerLayer(nn.Module):
    """Transformer layer block with optional multimer cross-chain attention."""

    def __init__(
        self,
        embed_dim,
        ffn_embed_dim,
        attention_heads,
        add_bias_kv=True,
        use_esm1b_layer_norm=False,
        use_rotary_embeddings: bool = False,
        use_multimer=False,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.ffn_embed_dim = ffn_embed_dim
        self.attention_heads = attention_heads
        self.use_rotary_embeddings = use_rotary_embeddings
        self.use_multimer = use_multimer
        self._init_submodules(add_bias_kv, use_esm1b_layer_norm)

    def _init_submodules(self, add_bias_kv, use_esm1b_layer_norm):
        BertLayerNorm = ESM1bLayerNorm if use_esm1b_layer_norm else nn.LayerNorm

        self.self_attn = MultiheadAttention(
            self.embed_dim,
            self.attention_heads,
            add_bias_kv=add_bias_kv,
            add_zero_attn=False,
            use_rotary_embeddings=self.use_rotary_embeddings,
        )

        if self.use_multimer:
            self.multimer_attn = MultiheadAttention(
                self.embed_dim,
                self.attention_heads,
                add_bias_kv=add_bias_kv,
                add_zero_attn=False,
                use_rotary_embeddings=False,
                no_proj=True,
            )

        self.self_attn_layer_norm = BertLayerNorm(self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, self.ffn_embed_dim)
        self.fc2 = nn.Linear(self.ffn_embed_dim, self.embed_dim)
        self.final_layer_norm = BertLayerNorm(self.embed_dim)

    def forward(
        self, x, self_attn_mask=None, self_attn_padding_mask=None, need_head_weights=False
    ):
        residual = x
        x = self.self_attn_layer_norm(x)
        if self.use_multimer:
            self_attn, self_v = self.self_attn(
                query=x, key=x, value=x,
                key_padding_mask=self_attn_padding_mask,
                before_softmax=True,
            )
            multimer_attn, multimer_v = self.multimer_attn(
                query=x, key=x, value=x,
                key_padding_mask=self_attn_padding_mask,
                before_softmax=True,
            )
            attn_weights = torch.where(
                self_attn_mask.unsqueeze(1), multimer_attn, self_attn
            )
            attn_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(
                attn_weights
            )
            attn_probs_dropout = F.dropout(
                attn_probs, p=self.self_attn.dropout, training=self.training
            )
            self_attn_probs = attn_probs_dropout.masked_fill(
                self_attn_mask.unsqueeze(1), 0.0
            )
            multimer_attn_probs = attn_probs_dropout.masked_fill(
                ~self_attn_mask.unsqueeze(1), 0.0
            )
            attn_out = torch.matmul(self_attn_probs, self_v) + torch.matmul(
                multimer_attn_probs, multimer_v
            )
            attn_out = attn_out.transpose(1, 2).contiguous()
            attn_out = attn_out.view(*attn_out.shape[:2], -1)
            x = self.self_attn.out_proj(attn_out).transpose(0, 1).contiguous()
            if need_head_weights:
                attn = attn_probs.transpose(0, 1).contiguous()
            else:
                attn = attn_probs.mean(1)
        else:
            x, attn = self.self_attn(
                query=x, key=x, value=x,
                key_padding_mask=self_attn_padding_mask,
                need_weights=True,
                need_head_weights=need_head_weights,
                attn_mask=self_attn_mask,
            )
        x = residual + x

        residual = x
        x = self.final_layer_norm(x)
        x = gelu(self.fc1(x))
        x = self.fc2(x)
        x = residual + x
        return x, attn


class RobertaLMHead(nn.Module):
    """Head for masked language modeling."""

    def __init__(self, embed_dim, output_dim, weight):
        super().__init__()
        self.dense = nn.Linear(embed_dim, embed_dim)
        self.layer_norm = ESM1bLayerNorm(embed_dim)
        self.weight = weight
        self.bias = nn.Parameter(torch.zeros(output_dim))

    def forward(self, features):
        x = self.dense(features)
        x = gelu(x)
        x = self.layer_norm(x)
        x = F.linear(x, self.weight) + self.bias
        return x


# ============================================================================
# ESM2 backbone (from mint/model/esm2.py)
# ============================================================================


class ESM2(nn.Module):
    def __init__(
        self,
        num_layers: int = 33,
        embed_dim: int = 1280,
        attention_heads: int = 20,
        alphabet: Union[Alphabet, str] = "ESM-1b",
        token_dropout: bool = True,
        use_multimer: bool = False,
    ):
        super().__init__()
        self.num_layers = num_layers
        self.embed_dim = embed_dim
        self.attention_heads = attention_heads
        if not isinstance(alphabet, Alphabet):
            alphabet = Alphabet.from_architecture(alphabet)
        self.alphabet = alphabet
        self.alphabet_size = len(alphabet)
        self.padding_idx = alphabet.padding_idx
        self.mask_idx = alphabet.mask_idx
        self.cls_idx = alphabet.cls_idx
        self.eos_idx = alphabet.eos_idx
        self.prepend_bos = alphabet.prepend_bos
        self.append_eos = alphabet.append_eos
        self.token_dropout = token_dropout
        self.use_multimer = use_multimer
        self._init_submodules()

    def _init_submodules(self):
        self.embed_scale = 1
        self.embed_tokens = nn.Embedding(
            self.alphabet_size, self.embed_dim, padding_idx=self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                TransformerLayer(
                    self.embed_dim,
                    4 * self.embed_dim,
                    self.attention_heads,
                    add_bias_kv=False,
                    use_esm1b_layer_norm=True,
                    use_rotary_embeddings=True,
                    use_multimer=self.use_multimer,
                )
                for _ in range(self.num_layers)
            ]
        )
        self.emb_layer_norm_after = ESM1bLayerNorm(self.embed_dim)
        self.lm_head = RobertaLMHead(
            embed_dim=self.embed_dim,
            output_dim=self.alphabet_size,
            weight=self.embed_tokens.weight,
        )

    def forward(
        self,
        tokens,
        chain_ids=None,
        repr_layers=[],
        need_head_weights=False,
        return_contacts=False,
    ):
        if return_contacts:
            need_head_weights = True

        assert tokens.ndim == 2
        padding_mask = tokens.eq(self.padding_idx)

        if chain_ids is None:
            chain_ids = torch.zeros_like(tokens)

        if self.use_multimer:
            # Cross-chain mask: True where tokens are from different chains.
            # Used by TransformerLayer to select multimer vs self attention.
            self_attn_mask = ~torch.eq(
                chain_ids.unsqueeze(-1), chain_ids.unsqueeze(-2)
            )
        else:
            # Without multimer attention, no cross-chain masking is needed.
            # All tokens can attend to all other tokens (standard ESM2).
            self_attn_mask = None

        x = self.embed_scale * self.embed_tokens(tokens)

        if self.token_dropout:
            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
            mask_ratio_train = 0.15 * 0.8
            src_lengths = (~padding_mask).sum(-1)
            mask_ratio_observed = (
                (tokens == self.mask_idx).sum(-1).to(x.dtype) / src_lengths
            )
            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]

        if padding_mask is not None:
            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))

        repr_layers = set(repr_layers)
        hidden_representations = {}
        if 0 in repr_layers:
            hidden_representations[0] = x

        if need_head_weights:
            attn_weights = []

        x = x.transpose(0, 1)

        if not padding_mask.any():
            padding_mask = None

        for layer_idx, layer in enumerate(self.layers):
            x, attn = layer(
                x,
                self_attn_padding_mask=padding_mask,
                need_head_weights=need_head_weights,
                self_attn_mask=self_attn_mask,
            )
            if (layer_idx + 1) in repr_layers:
                hidden_representations[layer_idx + 1] = x.transpose(0, 1)
            if need_head_weights:
                attn_weights.append(attn.transpose(1, 0))

        x = self.emb_layer_norm_after(x)
        x = x.transpose(0, 1)

        if (layer_idx + 1) in repr_layers:
            hidden_representations[layer_idx + 1] = x
        x = self.lm_head(x)

        result = {"logits": x, "representations": hidden_representations}
        if need_head_weights:
            attentions = torch.stack(attn_weights, 1)
            if padding_mask is not None:
                attention_mask = 1 - padding_mask.type_as(attentions)
                attention_mask = (
                    attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)
                )
                attentions = attentions * attention_mask[:, None, None, :, :]
            result["attentions"] = attentions

        return result


# ============================================================================
# Base PreTrainedModel (inlined from mint_stability/modeling_base.py)
# ============================================================================

class StabilityPreTrainedModel(PreTrainedModel):
    """Shared base for MINT and SPEARMINT stability models.

    Provides:
      - Common weight initialization (_init_weights)
      - Robust from_pretrained() that works across transformers versions
    """

    base_model_prefix = ""
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        """Default weight init -- overridden by from_pretrained() loading."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        """Load model with explicit state_dict loading.

        The default PreTrainedModel.from_pretrained() loading pipeline varies
        significantly across transformers versions (meta-device init, key
        prefix stripping, tied-weight handling, _fast_init, etc.).  This
        override sidesteps all of that and just does a straightforward
        ``load_state_dict`` which works identically everywhere.
        """
        # --- resolve to a local directory --------------------------------
        if os.path.isdir(pretrained_model_name_or_path):
            model_dir = pretrained_model_name_or_path
        else:
            from huggingface_hub import snapshot_download

            model_dir = snapshot_download(
                pretrained_model_name_or_path,
                cache_dir=kwargs.get("cache_dir"),
                token=kwargs.get("token") or kwargs.get("use_auth_token"),
                revision=kwargs.get("revision"),
            )

        # --- config -------------------------------------------------------
        config = cls.config_class.from_pretrained(model_dir)

        # --- build model (random init, post_init runs _init_weights) ------
        model = cls(config)

        # --- load checkpoint weights --------------------------------------
        sf_path = os.path.join(model_dir, "model.safetensors")
        bin_path = os.path.join(model_dir, "pytorch_model.bin")

        if os.path.exists(sf_path):
            from safetensors.torch import load_file

            state_dict = load_file(sf_path)
        elif os.path.exists(bin_path):
            state_dict = torch_load(bin_path, map_location="cpu", weights_only=True)
        else:
            raise FileNotFoundError(
                f"No model weights found in {model_dir}"
            )

        # Use _keys_to_ignore_on_load_missing from the subclass
        allowed_missing = set(getattr(cls, "_keys_to_ignore_on_load_missing", []))

        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        real_missing = [k for k in missing if k not in allowed_missing]
        if unexpected:
            raise RuntimeError(
                f"Unexpected keys in checkpoint: {unexpected}"
            )
        if real_missing:
            raise RuntimeError(
                f"Missing keys in checkpoint: {real_missing}"
            )

        # --- optional: move to device / dtype ----------------------------
        device_map = kwargs.get("device_map")
        if device_map == "auto" or device_map == "cuda":
            model = model.cuda()
        elif device_map is not None:
            import warnings
            warnings.warn(
                f"device_map={device_map!r} is not supported; ignoring. "
                "Use 'auto' or 'cuda'.",
                stacklevel=2,
            )
        dtype = kwargs.get("torch_dtype")
        if dtype is not None:
            model = model.to(dtype=dtype)

        model.eval()
        return model


# ============================================================================
# HuggingFace PreTrainedModel wrapper
# ============================================================================

class MintStabilityPreTrainedModel(StabilityPreTrainedModel):
    """Base class for MINT stability models."""

    config_class = MintStabilityConfig


class MintStabilityForRegression(MintStabilityPreTrainedModel):
    """MINT backbone + projection head for pMHC stability prediction.

    Architecture: ESM2-650M with optional cross-chain multimer attention,
    mean-pooled over non-special tokens, followed by a projection head
    (Linear -> ReLU -> Dropout -> Linear) outputting a scalar prediction
    (half-life in hours).
    """

    config_class = MintStabilityConfig
    # lm_head.weight is tied to embed_tokens.weight in ESM2 but is never
    # used during stability inference -- suppress the missing-key warning.
    _keys_to_ignore_on_load_missing = ["model.lm_head.weight"]

    def __init__(self, config: MintStabilityConfig):
        super().__init__(config)
        self.config = config

        self.model = ESM2(
            num_layers=config.num_layers,
            embed_dim=config.embed_dim,
            attention_heads=config.attention_heads,
            token_dropout=config.token_dropout,
            use_multimer=config.use_multimer,
        )

        self.project = nn.Sequential(
            nn.Linear(config.embed_dim, config.hidden_dim),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.hidden_dim, config.output_size),
        )

        self.sigmoid_output = config.sigmoid_output
        self.post_init()

    def forward(
        self,
        chains: torch.Tensor,
        chain_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
    ) -> Dict[str, Optional[torch.Tensor]]:
        """
        Args:
            chains: Token IDs, shape (batch, seq_len).
            chain_ids: Chain membership IDs, shape (batch, seq_len).
                       0 = peptide, 1 = MHC.
            labels: Optional regression targets, shape (batch,).

        Returns:
            dict with "loss" (if labels provided) and "logits".
        """
        mask = (
            (~chains.eq(self.model.cls_idx))
            & (~chains.eq(self.model.eos_idx))
            & (~chains.eq(self.model.padding_idx))
        )
        chain_out = self.model(
            chains, chain_ids, repr_layers=[self.config.num_layers]
        )["representations"][self.config.num_layers]

        mask_expanded = mask.unsqueeze(-1).expand_as(chain_out)
        masked_chain_out = chain_out * mask_expanded
        sum_masked = masked_chain_out.sum(dim=1)
        mask_counts = mask.sum(dim=1, keepdim=True).float()
        mean_chain_out = sum_masked / mask_counts

        out = self.project(mean_chain_out)
        if self.sigmoid_output:
            out = torch.sigmoid(out)

        loss = None
        if labels is not None:
            loss = F.mse_loss(out.squeeze(-1), labels)

        return {"loss": loss, "logits": out}


# ============================================================================
# Tokenizer helper
# ============================================================================


class MintTokenizer:
    """Tokenize peptide + MHC sequences for MINT S1/S2 models.

    Usage:
        tokenizer = MintTokenizer()
        chains, chain_ids = tokenizer.prepare_input("GILGFVFTL", "MAVMAPRTL...")
        output = model(chains.unsqueeze(0), chain_ids.unsqueeze(0))
    """

    def __init__(self):
        self.alphabet = Alphabet.from_architecture("ESM-1b")

    def _encode_sequence(self, seq: str) -> torch.Tensor:
        seq = seq.replace("J", "L")
        encoded = self.alphabet.encode("<cls>" + seq + "<eos>")
        return torch.tensor(encoded, dtype=torch.int64)

    def prepare_input(
        self, peptide: str, mhc_sequence: str
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Tokenize a single peptide-MHC pair.

        Returns:
            chains: Token IDs, shape (seq_len,). Add batch dim with .unsqueeze(0).
            chain_ids: Chain membership, shape (seq_len,). 0=peptide, 1=MHC.
        """
        pep_tokens = self._encode_sequence(peptide)
        mhc_tokens = self._encode_sequence(mhc_sequence)

        chains = torch.cat([pep_tokens, mhc_tokens], dim=0)
        chain_ids = torch.cat(
            [
                torch.zeros(len(pep_tokens), dtype=torch.int32),
                torch.ones(len(mhc_tokens), dtype=torch.int32),
            ],
            dim=0,
        )
        return chains, chain_ids

    def prepare_batch(
        self, peptides: List[str], mhc_sequences: List[str]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Tokenize a batch of peptide-MHC pairs with padding.

        Returns:
            chains: (batch, max_seq_len)
            chain_ids: (batch, max_seq_len)
        """
        all_chains = []
        all_chain_ids = []
        for pep, mhc in zip(peptides, mhc_sequences):
            c, cid = self.prepare_input(pep, mhc)
            all_chains.append(c)
            all_chain_ids.append(cid)

        max_len = max(len(c) for c in all_chains)
        padded_chains = torch.full(
            (len(all_chains), max_len),
            self.alphabet.padding_idx,
            dtype=torch.int64,
        )
        padded_chain_ids = torch.zeros(
            (len(all_chains), max_len), dtype=torch.int32
        )

        for i, (c, cid) in enumerate(zip(all_chains, all_chain_ids)):
            padded_chains[i, : len(c)] = c
            padded_chain_ids[i, : len(cid)] = cid

        return padded_chains, padded_chain_ids