| """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.", |
| } |
|
|
