File size: 6,536 Bytes
6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 98ed1b7 6019d52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """Deprecated MitoInteract v1 architecture retained for artifact forensics.
The v1 checkpoint failed its evaluation gates and uses unsafe pickle
serialization. New work lives under ``recovery/``. Loading and inference are
disabled by default so this module cannot silently present an undocumented
legacy score as pKd.
"""
import torch
import torch.nn as nn
from transformers import EsmModel, AutoModel, AutoTokenizer
class MitoInteract(nn.Module):
def __init__(
self,
esm_model_name="facebook/esm2_t33_650M_UR50D",
mol_model_name="seyonec/ChemBERTa-zinc-base-v1",
protein_dim=1280,
mol_dim=768,
proj_dim=256,
n_heads=8,
dropout=0.1,
freeze_encoders=True,
):
super().__init__()
self.freeze_encoders = freeze_encoders
self.esm = EsmModel.from_pretrained(esm_model_name)
self.protein_dim = protein_dim
self.mol_encoder = AutoModel.from_pretrained(mol_model_name)
self.mol_dim = mol_dim
if freeze_encoders:
for p in self.esm.parameters():
p.requires_grad = False
for p in self.mol_encoder.parameters():
p.requires_grad = False
self.prot_proj = nn.Sequential(
nn.Linear(protein_dim, proj_dim),
nn.LayerNorm(proj_dim),
nn.ReLU(),
nn.Dropout(dropout),
)
self.mol_proj = nn.Sequential(
nn.Linear(mol_dim, proj_dim),
nn.LayerNorm(proj_dim),
nn.ReLU(),
nn.Dropout(dropout),
)
self.cross_attn_mol2prot = nn.MultiheadAttention(
proj_dim, n_heads, dropout=dropout, batch_first=True
)
self.cross_attn_prot2mol = nn.MultiheadAttention(
proj_dim, n_heads, dropout=dropout, batch_first=True
)
self.ln_mol2prot = nn.LayerNorm(proj_dim)
self.ln_prot2mol = nn.LayerNorm(proj_dim)
fused_dim = proj_dim * 2
self.mlp = nn.Sequential(
nn.Linear(fused_dim, 512),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(512, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(256, 128),
nn.BatchNorm1d(128),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(128, 1),
)
def encode_protein(self, input_ids, attention_mask):
ctx = torch.no_grad() if self.freeze_encoders else torch.enable_grad()
with ctx:
out = self.esm(input_ids=input_ids, attention_mask=attention_mask)
mask = attention_mask.unsqueeze(-1).float()
pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
return pooled, out.last_hidden_state
def encode_molecule(self, input_ids, attention_mask):
ctx = torch.no_grad() if self.freeze_encoders else torch.enable_grad()
with ctx:
out = self.mol_encoder(input_ids=input_ids, attention_mask=attention_mask)
return out.pooler_output, out.last_hidden_state
def forward(
self, prot_input_ids, prot_attention_mask, mol_input_ids, mol_attention_mask
):
prot_pooled, prot_seq = self.encode_protein(prot_input_ids, prot_attention_mask)
mol_pooled, mol_seq = self.encode_molecule(mol_input_ids, mol_attention_mask)
prot_seq_proj = self.prot_proj(prot_seq)
mol_seq_proj = self.mol_proj(mol_seq)
prot_q = self.prot_proj(prot_pooled).unsqueeze(1)
mol_q = self.mol_proj(mol_pooled).unsqueeze(1)
prot_pad_mask = prot_attention_mask == 0
mol_pad_mask = mol_attention_mask == 0
h_prot2mol, _ = self.cross_attn_prot2mol(
prot_q, mol_seq_proj, mol_seq_proj, key_padding_mask=mol_pad_mask
)
h_mol2prot, _ = self.cross_attn_mol2prot(
mol_q, prot_seq_proj, prot_seq_proj, key_padding_mask=prot_pad_mask
)
h_prot2mol = self.ln_prot2mol(h_prot2mol.squeeze(1))
h_mol2prot = self.ln_mol2prot(h_mol2prot.squeeze(1))
fused = torch.cat([h_prot2mol, h_mol2prot], dim=-1)
return self.mlp(fused).squeeze(-1)
def load_model(checkpoint_path, device="cpu", *, allow_unsafe_legacy=False):
"""Load the failed v1 pickle only after explicit risk acknowledgement."""
if not allow_unsafe_legacy:
raise RuntimeError(
"MitoInteract v1 is deprecated and full_model.pt requires unsafe pickle "
"deserialization. See AUDIT.md. Pass allow_unsafe_legacy=True only for "
"controlled artifact forensics, never for untrusted files."
)
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
config = checkpoint["config"]
model = MitoInteract(
esm_model_name=config["esm_model"],
mol_model_name=config["mol_model"],
protein_dim=config["protein_dim"],
mol_dim=config["mol_dim"],
proj_dim=config["proj_dim"],
n_heads=config["n_heads"],
dropout=config["dropout"],
freeze_encoders=True,
)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
return model, config
def predict_binding(
model, protein_seq, smiles, device="cpu", *, allow_invalid_legacy_output=False
):
"""Return the undocumented v1 score only after explicit acknowledgement."""
if not allow_invalid_legacy_output:
raise RuntimeError(
"v1 output has undocumented target semantics and failed validation. "
"It cannot be interpreted as pKd or converted to Kd. See AUDIT.md."
)
prot_tokenizer = AutoTokenizer.from_pretrained(model.esm.name_or_path)
mol_tokenizer = AutoTokenizer.from_pretrained(model.mol_encoder.name_or_path)
prot_enc = prot_tokenizer(
protein_seq, return_tensors="pt", padding=True, truncation=True, max_length=512
)
mol_enc = mol_tokenizer(
smiles, return_tensors="pt", padding=True, truncation=True, max_length=200
)
model = model.to(device)
with torch.no_grad():
legacy_score = model(
prot_enc["input_ids"].to(device),
prot_enc["attention_mask"].to(device),
mol_enc["input_ids"].to(device),
mol_enc["attention_mask"].to(device),
)
return {
"legacy_score": legacy_score.item(),
"warning": "Do not interpret this value as pKd, Kd, or validated affinity.",
}
|