""" 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] = ("", "", "", ""), append_toks: Sequence[str] = ("", "", ""), 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"") 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[""] self.padding_idx = self.get_idx("") self.cls_idx = self.get_idx("") self.mask_idx = self.get_idx("") self.eos_idx = self.get_idx("") self.all_special_tokens = ["", "", "", "", ""] 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 = ("", "", "", "") append_toks = ("",) 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("" + seq + "") 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