import os import torch import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm import pandas as pd from typing import List from rdkit import Chem from rdkit.Chem import AllChem from transformers import PretrainedConfig from transformers import PreTrainedModel from transformers import AutoModel from torch_geometric.nn import GCNConv from torch_geometric.data import Data from torch_geometric.loader import DataLoader from torch_scatter import scatter class SmilesDataset(torch.utils.data.Dataset): def __init__(self, smiles): self.smiles_list = smiles self.data_list = [] def __len__(self): return len(self.data_list) def __getitem__(self, idx): return self.data_list[idx] def get_data(self, smiles): self.smiles_list = smiles # self.data_list = [] # bonds = {BT.SINGLE: 0, BT.DOUBLE: 1, BT.TRIPLE: 2, BT.AROMATIC: 3} types = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'S': 4} for i in range(len(self.smiles_list)): # 将 SMILES 表示转换为 RDKit 的分子对象 # print(self.smiles_list[i]) mol = Chem.MolFromSmiles(self.smiles_list[i]) # 从smiles编码中获取结构信息 if mol is None: print("无法创建Mol对象", self.smiles_list[i]) else: mol3d = Chem.AddHs( mol) # 在rdkit中,分子在默认情况下是不显示氢的,但氢原子对于真实的几何构象计算有很大的影响,所以在计算3D构象前,需要使用Chem.AddHs()方法加上氢原子 if mol3d is None: print("无法创建mol3d对象", self.smiles_list[i]) else: AllChem.EmbedMolecule(mol3d, randomSeed=1) # 生成3D构象 N = mol3d.GetNumAtoms() # 获取原子坐标信息 if mol3d.GetNumConformers() > 0: conformer = mol3d.GetConformer() pos = conformer.GetPositions() pos = torch.tensor(pos, dtype=torch.float) type_idx = [] # atomic_number = [] # aromatic = [] # sp = [] # sp2 = [] # sp3 = [] for atom in mol3d.GetAtoms(): type_idx.append(types[atom.GetSymbol()]) # atomic_number.append(atom.GetAtomicNum()) # aromatic.append(1 if atom.GetIsAromatic() else 0) # hybridization = atom.GetHybridization() # sp.append(1 if hybridization == HybridizationType.SP else 0) # sp2.append(1 if hybridization == HybridizationType.SP2 else 0) # sp3.append(1 if hybridization == HybridizationType.SP3 else 0) # z = torch.tensor(atomic_number, dtype=torch.long) row, col, edge_type = [], [], [] for bond in mol3d.GetBonds(): start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() row += [start, end] col += [end, start] # edge_type += 2 * [bonds[bond.GetBondType()]] edge_index = torch.tensor([row, col], dtype=torch.long) # edge_type = torch.tensor(edge_type, dtype=torch.long) # edge_attr = F.one_hot(edge_type, num_classes=len(bonds)).to(torch.float) perm = (edge_index[0] * N + edge_index[1]).argsort() edge_index = edge_index[:, perm] # edge_type = edge_type[perm] # edge_attr = edge_attr[perm] # # row, col = edge_index # hs = (z == 1).to(torch.float) x = torch.tensor(type_idx).to(torch.float) # y = self.y_list[i] data = Data(x=x, pos=pos, edge_index=edge_index, smiles=self.smiles_list[i]) self.data_list.append(data) else: print("无法创建comfor", self.smiles_list[i]) return self.data_list """ MLP Layer used after graph vector representation """ class MLPReadout(nn.Module): def __init__(self, input_dim, output_dim, L=2): # L=nb_hidden_layers super().__init__() list_FC_layers = [nn.Linear(input_dim // 2 ** l, input_dim // 2 ** (l + 1), bias=True) for l in range(L)] list_FC_layers.append(nn.Linear(input_dim // 2 ** L, output_dim, bias=True)) self.FC_layers = nn.ModuleList(list_FC_layers) self.L = L def forward(self, x): y = x for l in range(self.L): y = self.FC_layers[l](y) y = F.relu(y) y = self.FC_layers[self.L](y) return y class GCNNet(torch.nn.Module): def __init__(self, input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1): super(GCNNet, self).__init__() self.embedding = torch.nn.Embedding(emb_input, hidden_size, padding_idx=0) self.input_feature = input_feature self.n_layers = n_layers # 2层GCN self.num_classes = num_classes self.conv1 = GCNConv(hidden_size, hidden_size) self.conv2 = GCNConv(hidden_size, 32) self.mlp = MLPReadout(32, num_classes) def forward_features(self, data): x, edge_index, batch = data.x.long(), data.edge_index, data.batch x = self.embedding(x.reshape(-1)) for i in range(self.n_layers): x = F.relu(self.conv1(x, edge_index)) x = F.relu(self.conv2(x, edge_index)) x = scatter(x, batch, dim=-2, reduce='mean') x = self.mlp(x) return x.squeeze(-1) class GCNConfig(PretrainedConfig): model_type = "gcn" def __init__( self, input_feature: int=64, emb_input: int=20, hidden_size: int=64, n_layers: int=6, num_classes: int=1, smiles: List[str] = None, processor_class: str = "SmilesProcessor", **kwargs, ): self.input_feature = input_feature # the dimension of input feature self.emb_input = emb_input # the embedding dimension of input feature self.hidden_size = hidden_size # the hidden size of GCN self.n_layers = n_layers # the number of GCN layers self.num_classes = num_classes # the number of output classes self.smiles = smiles # process smiles self.processor_class = processor_class super().__init__(**kwargs) class GCNModel(PreTrainedModel): config_class = GCNConfig def __init__(self, config): super().__init__(config) self.model = GCNNet( input_feature=config.input_feature, emb_input=config.emb_input, hidden_size=config.hidden_size, n_layers=config.n_layers, num_classes=config.num_classes, ) self.process = SmilesDataset( smiles=config.smiles, ) self.gcn_model = None self.dataset = None self.output = None self.data_loader = None self.pred_data = None def forward(self, tensor): return self.model.forward_features(tensor) # def process_smiles(self, smiles): # return self.process.get_data(smiles) def predict_smiles(self, smiles, device: str='cpu', result_dir: str='./', **kwargs): batch_size = kwargs.pop('batch_size', 1) shuffle = kwargs.pop('shuffle', False) drop_last = kwargs.pop('drop_last', False) num_workers = kwargs.pop('num_workers', 0) self.gcn_model = AutoModel.from_pretrained("Huhujingjing/custom-gcn", trust_remote_code=True).to(device) self.gcn_model.eval() self.dataset = self.process.get_data(smiles) self.output = "" self.output += ("predicted samples num: {}\n".format(len(self.dataset))) self.output +=("predicted samples:{}\n".format(self.dataset[0])) self.data_loader = DataLoader(self.dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers ) self.pred_data = { 'smiles': [], 'pred': [] } for batch in tqdm(self.data_loader): batch = batch.to(device) with torch.no_grad(): self.pred_data['smiles'] += batch['smiles'] self.pred_data['pred'] += self.gcn_model(batch).cpu().tolist() pred = torch.tensor(self.pred_data['pred']).reshape(-1) if device == 'cuda': pred = pred.cpu().tolist() self.pred_data['pred'] = pred pred_df = pd.DataFrame(self.pred_data) pred_df['pred'] = pred_df['pred'].apply(lambda x: round(x, 2)) self.output +=('-' * 40 + '\n'+'predicted result: \n'+'{}\n'.format(pred_df)) self.output +=('-' * 40) pred_df.to_csv(os.path.join(result_dir, 'gcn.csv'), index=False) self.output +=('\nsave predicted result to {}\n'.format(os.path.join(result_dir, 'gcn.csv'))) return self.output if __name__ == "__main__": gcn_config = GCNConfig(input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1, smiles=["C", "CC", "CCC"], processor_class="SmilesProcessor") gcnd = GCNModel(gcn_config) gcnd.model.load_state_dict(torch.load(r'G:\Trans_MXM\gcn_model\gcn.pt')) gcnd.save_pretrained("custom-gcn")