custom-gcn / modeling_gcn.py

Upload model

64fa2ae 11 months ago

10.1 kB

	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")