Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import dgl | |
| from dgllife.model import GCN, GAT | |
| class MolEnc(nn.Module): | |
| def __init__(self, | |
| args, | |
| in_dim,): | |
| super().__init__() | |
| # Whether to return node-level embeddings (sum_nodes, D) instead of a | |
| # pooled graph embedding (B, D). | |
| # | |
| # Backward compatible defaults: | |
| # - Historically, this was inferred from `args.model` for FLARE-style models. | |
| # - New: allow explicit override via `args.mol_return_emb` (preferred) or | |
| # `args.return_mol_emb` (legacy-friendly alias). | |
| explicit = None | |
| if hasattr(args, "mol_return_emb"): | |
| explicit = getattr(args, "mol_return_emb") | |
| elif hasattr(args, "return_mol_emb"): | |
| explicit = getattr(args, "return_mol_emb") | |
| if explicit is not None: | |
| self.return_emb = bool(explicit) | |
| else: | |
| self.return_emb = False | |
| if args.model in ("filipContrastive", "crossAttenContrastive", "filipGlobalContrastive"): | |
| self.return_emb = True | |
| dropout = [args.gnn_dropout for _ in range(len(args.gnn_channels))] | |
| batchnorm = [True for _ in range(len(args.gnn_channels))] | |
| gnn_map = { | |
| "gcn": GCN(in_dim, args.gnn_channels, batchnorm = batchnorm, dropout = dropout), | |
| } | |
| self.GNN = gnn_map[args.gnn_type] | |
| pooling_type = getattr(args, "mol_pooling", "max") | |
| pool_map = { | |
| "max": dgl.nn.pytorch.glob.MaxPooling, | |
| "mean": dgl.nn.pytorch.glob.AvgPooling, | |
| "avg": dgl.nn.pytorch.glob.AvgPooling, | |
| "sum": dgl.nn.pytorch.glob.SumPooling, | |
| } | |
| if pooling_type not in pool_map: | |
| raise ValueError(f"Unsupported mol_pooling: {pooling_type}") | |
| self.pool = pool_map[pooling_type]() | |
| if not self.return_emb: | |
| self.fc1_graph = nn.Linear(args.gnn_channels[len(args.gnn_channels) - 1], args.gnn_hidden_dim * 2) | |
| self.fc2_graph = nn.Linear(args.gnn_hidden_dim * 2, args.final_embedding_dim) | |
| self.dropout = nn.Dropout(args.fc_dropout) | |
| self.relu = nn.ReLU() | |
| def forward(self, g) -> torch.Tensor: | |
| g1 = g | |
| f1 = g.ndata['h'] | |
| f = self.GNN(g1, f1) | |
| if self.return_emb: | |
| return f | |
| h = self.pool(g1, f) | |
| h1 = self.relu(self.fc1_graph(h)) | |
| h1 = self.dropout(h1) | |
| h1 = self.fc2_graph(h1) | |
| h1 = self.dropout(h1) | |
| return h1 | |