Spaces:
Sleeping
Sleeping
| """ | |
| biointeract.py — Full BioInteract model (v2: dual-channel + focal loss ready). | |
| Connects drug encoder (GINE + Morgan FP), target encoder (ESM-2), | |
| cross-attention interaction module, and prediction head. | |
| v2 changes: | |
| - Dual-channel drug encoding: GNN atoms + Morgan fingerprint global vector | |
| - Graph augmentation (DropNode/DropEdge) for cold-start robustness | |
| - Morgan fingerprint branch fused at prediction head | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from torch_geometric.data import Batch | |
| from .drug_encoder import DrugEncoder, MorganFPEncoder, GraphAugmentation | |
| from .target_encoder import TargetEncoder | |
| from .interaction import CrossAttentionInteraction, GatedPooling | |
| class BioInteract(nn.Module): | |
| """ | |
| BioInteract: Interpretable Drug-Target Interaction Prediction | |
| via Residue-Level Cross-Attention with Biological Prior Knowledge. | |
| Forward pass: | |
| 1. Encode drug molecular graph → per-atom representations | |
| 2. Encode protein (ESM-2 + domain features) → per-residue representations | |
| 3. Cross-attention interaction → atom-residue interaction map | |
| 4. Gated pooling → fixed-size vector | |
| 5. Prediction head → binding probability / affinity | |
| """ | |
| def __init__(self, config: dict): | |
| super().__init__() | |
| drug_cfg = config.get('drug_encoder', {}) | |
| target_cfg = config.get('target_encoder', {}) | |
| inter_cfg = config.get('interaction', {}) | |
| pred_cfg = config.get('predictor', {}) | |
| hidden_dim = drug_cfg.get('hidden_dim', 256) | |
| # --- configurable Morgan FP --- | |
| self.use_morgan_fp = drug_cfg.get('use_morgan_fp', False) | |
| # --- graph augmentation (training only) --- | |
| self.graph_aug = GraphAugmentation( | |
| p_node=drug_cfg.get('drop_node', 0.0), | |
| p_edge=drug_cfg.get('drop_edge', 0.0), | |
| ) | |
| # --- encoders --- | |
| self.drug_encoder = DrugEncoder( | |
| num_atom_features=drug_cfg.get('num_atom_features', 52), | |
| edge_dim=drug_cfg.get('edge_dim', 16), | |
| hidden_dim=hidden_dim, | |
| num_layers=drug_cfg.get('num_layers', 3), | |
| dropout=drug_cfg.get('dropout', 0.2), | |
| jk=drug_cfg.get('jk', 'last'), | |
| ) | |
| # Morgan fingerprint branch (optional) | |
| if self.use_morgan_fp: | |
| morgan_nbits = drug_cfg.get('morgan_nbits', 1024) | |
| self.morgan_encoder = MorganFPEncoder( | |
| input_dim=morgan_nbits, | |
| hidden_dim=hidden_dim, | |
| dropout=drug_cfg.get('dropout', 0.2), | |
| ) | |
| self.target_encoder = TargetEncoder( | |
| esm2_dim=target_cfg.get('esm2_dim', 1280), | |
| projection_dim=hidden_dim, | |
| physchem_dim=4, | |
| domain_embed_dim=target_cfg.get('domain_embed_dim', 32), | |
| num_domain_types=target_cfg.get('num_domain_types', 50), | |
| use_domain_features=target_cfg.get('use_domain_features', True), | |
| ) | |
| # --- interaction --- | |
| self.interaction = CrossAttentionInteraction( | |
| hidden_dim=hidden_dim, | |
| num_heads=inter_cfg.get('num_heads', 8), | |
| dropout=inter_cfg.get('dropout', 0.1), | |
| ) | |
| # --- pooling --- | |
| self.drug_pooling = GatedPooling(hidden_dim) | |
| self.prot_pooling = GatedPooling(hidden_dim) | |
| # --- prediction head --- | |
| # fusion dimension depends on whether Morgan FP is used | |
| if self.use_morgan_fp: | |
| pred_input_dim = hidden_dim * 3 # drug_gnn + protein + morgan | |
| else: | |
| pred_input_dim = hidden_dim * 2 # drug_gnn + protein | |
| pred_hidden = pred_cfg.get('hidden_dims', [256, 128]) | |
| pred_dropout = pred_cfg.get('dropout', 0.3) | |
| self.task = pred_cfg.get('task', 'classification') | |
| layers = [] | |
| in_dim = pred_input_dim | |
| for h_dim in pred_hidden: | |
| layers.extend([ | |
| nn.Linear(in_dim, h_dim), | |
| nn.ReLU(), | |
| nn.Dropout(pred_dropout), | |
| ]) | |
| in_dim = h_dim | |
| layers.append(nn.Linear(in_dim, 1)) | |
| self.predictor = nn.Sequential(*layers) | |
| def forward(self, | |
| drug_batch: Batch, | |
| esm2_embedding: torch.Tensor, | |
| physicochemical: torch.Tensor, | |
| domain_labels: torch.Tensor, | |
| protein_mask: torch.Tensor, | |
| morgan_fp: torch.Tensor = None, | |
| return_attention: bool = False): | |
| """ | |
| Args: | |
| drug_batch: PyG Batch of molecular graphs | |
| esm2_embedding: (B, L, esm2_dim) | |
| physicochemical: (B, L, 4) | |
| domain_labels: (B, L) | |
| protein_mask: (B, L) boolean | |
| morgan_fp: (B, 1024) Morgan fingerprint vectors | |
| return_attention: if True, also return attention maps | |
| Returns: | |
| prediction: (B, 1) — binding score | |
| attention_data: dict (only if return_attention=True) | |
| """ | |
| batch_size = esm2_embedding.size(0) | |
| # 0. graph augmentation (training only) | |
| drug_batch = self.graph_aug(drug_batch) | |
| # 1. encode drug → per-atom representations | |
| atom_repr, batch_index = self.drug_encoder(drug_batch) | |
| # 2. encode protein → per-residue representations | |
| residue_repr = self.target_encoder( | |
| esm2_embedding, physicochemical, domain_labels, protein_mask | |
| ) | |
| # 3. reshape drug atoms into (B, max_atoms, D) for cross-attention | |
| drug_repr_padded, drug_mask = self._pad_drug_atoms( | |
| atom_repr, batch_index, batch_size | |
| ) | |
| # 4. cross-attention interaction | |
| drug_updated, prot_updated, interaction_map = self.interaction( | |
| drug_repr_padded, drug_mask, residue_repr, protein_mask | |
| ) | |
| # 5. gated pooling to fixed-size vectors | |
| drug_pooled = self.drug_pooling(drug_updated, drug_mask) # (B, D) | |
| prot_pooled = self.prot_pooling(prot_updated, protein_mask) # (B, D) | |
| # 6. Fusion — conditionally include Morgan FP | |
| if self.use_morgan_fp and morgan_fp is not None: | |
| morgan_repr = self.morgan_encoder(morgan_fp) # (B, D) | |
| fused = torch.cat([drug_pooled, prot_pooled, morgan_repr], dim=-1) # (B, 3D) | |
| else: | |
| fused = torch.cat([drug_pooled, prot_pooled], dim=-1) # (B, 2D) | |
| # 7. predict | |
| prediction = self.predictor(fused) # (B, 1) | |
| # NOTE: for classification, we return raw logits here. | |
| # Apply sigmoid only during inference (not training with AMP). | |
| # Use BCEWithLogitsLoss for training. | |
| if return_attention: | |
| return prediction, { | |
| 'interaction_map': interaction_map, | |
| 'drug_mask': drug_mask, | |
| 'protein_mask': protein_mask, | |
| } | |
| return prediction | |
| def predict_proba(self, *args, **kwargs): | |
| """Return probabilities (sigmoid applied) for inference.""" | |
| logits = self.forward(*args, **kwargs) | |
| if isinstance(logits, tuple): | |
| return torch.sigmoid(logits[0]), logits[1] | |
| return torch.sigmoid(logits) | |
| def _pad_drug_atoms(self, atom_repr, batch_index, batch_size): | |
| """ | |
| Convert scattered atom representations to padded batch tensor. | |
| Args: | |
| atom_repr: (total_atoms, D) — all atoms from all graphs | |
| batch_index: (total_atoms,) — which graph each atom belongs to | |
| batch_size: int | |
| Returns: | |
| padded: (B, max_atoms, D) | |
| mask: (B, max_atoms) boolean | |
| """ | |
| device = atom_repr.device | |
| D = atom_repr.size(-1) | |
| # count atoms per graph | |
| counts = torch.bincount(batch_index, minlength=batch_size) | |
| max_atoms = counts.max().item() | |
| padded = torch.zeros(batch_size, max_atoms, D, device=device) | |
| mask = torch.zeros(batch_size, max_atoms, dtype=torch.bool, device=device) | |
| for i in range(batch_size): | |
| atom_indices = (batch_index == i).nonzero(as_tuple=True)[0] | |
| n = atom_indices.size(0) | |
| padded[i, :n] = atom_repr[atom_indices] | |
| mask[i, :n] = True | |
| return padded, mask | |