| |
| """ |
| Created on Thu Sep 15 16:22:05 2022 |
| |
| @author: ZNDX002 |
| """ |
| from model import ModelCLR |
| import yaml |
| import os |
| import torch |
| import numpy as np |
| import re |
| from torch_geometric.data import Data, Batch |
| from dataloader.dataset_wrapper import MolToGraph |
| from rdkit import Chem |
|
|
| class ModelInference(object): |
| def __init__(self, config_path, pretrain_model_path, device): |
| assert config_path is not None, "config_path is None" |
| assert pretrain_model_path is not None, "pretrain_model_path is None" |
|
|
| if device is None: |
| self.device = torch.device( |
| "cuda" if torch.cuda.is_available() else "cpu") |
| else: |
| self.device = torch.device(device) |
|
|
| self.config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) |
| self.model = ModelCLR(**self.config["model_config"]).to(self.device) |
| state_dict = torch.load(pretrain_model_path,map_location=self.device) |
| self.model.load_state_dict(state_dict) |
| self.model.eval() |
|
|
|
|
| def smiles_encode(self, smiles_str): |
| with torch.no_grad(): |
| if isinstance(smiles_str, str): |
| |
| v_d = MolToGraph(smiles_str) |
| v_d = v_d.to(self.device) |
| smiles_tensor = self.model.smiles_encoder(v_d) |
| smiles_tensor=self.model.smi_esa(smiles_tensor,v_d.batch) |
| smiles_tensor = self.model.smi_proj(smiles_tensor) |
| smiles_tensor = smiles_tensor/smiles_tensor.norm(dim=-1, keepdim=True) |
| return smiles_tensor |
| else: |
| |
| graphs=[] |
| for smi in smiles_str: |
| v_d = MolToGraph(smi) |
| graphs.append(v_d) |
| v_ds = Batch.from_data_list(graphs) |
| v_ds = v_ds.to(self.device) |
| smiles_tensor = self.model.smiles_encoder(v_ds) |
| smiles_tensor=self.model.smi_esa(smiles_tensor,v_ds.batch) |
| smiles_tensor = self.model.smi_proj(smiles_tensor) |
| smiles_tensor = smiles_tensor/smiles_tensor.norm(dim=-1, keepdim=True) |
| return smiles_tensor |
|
|
| def ms2_encode(self, ms2_list): |
| with torch.no_grad(): |
| if not isinstance(ms2_list, list): |
| |
| spec_mz = ms2_list.mz |
| spec_intens = ms2_list.intensities |
| num_peak = len(spec_mz) |
| spec_mz = np.around(spec_mz, decimals=4) |
| spec_mz = np.pad(spec_mz, (0, 300 - len(spec_mz)), mode='constant', constant_values=0) |
| spec_intens = np.pad(spec_intens, (0, 300 - len(spec_intens)), mode='constant', constant_values=0) |
| spec_mz= torch.tensor(spec_mz).float().unsqueeze(0) |
| spec_intens= torch.tensor(spec_intens).float().unsqueeze(0) |
| num_peak = torch.LongTensor(num_peak).unsqueeze(0) |
| spec_tensor,spec_mask = self.model.ms_encoder(spec_mz,spec_intens,num_peak) |
| spec_tensor=self.model.spec_esa(spec_tensor,spec_mask) |
| spec_tensor = self.model.spec_proj(spec_tensor) |
| spec_tensor = spec_tensor/spec_tensor.norm(dim=-1, keepdim=True) |
| return spec_tensor |
| else: |
| |
| spec_mzs = [spec.mz for spec in ms2_list] |
| spec_intens = [spec.intensities for spec in ms2_list] |
| num_peaks = [len(i) for i in spec_mzs] |
| spec_mzs = [np.around(spec_mz, decimals=4) for spec_mz in spec_mzs] |
| num_peaks = torch.LongTensor(num_peaks) |
| mzs = [torch.from_numpy(spec_mz).float() for spec_mz in spec_mzs] |
| intens = [torch.from_numpy(spec_intens).float() for spec_intens in spec_intens] |
| mzs_tensors = torch.nn.utils.rnn.pad_sequence( |
| mzs, batch_first=True, padding_value=0 |
| ) |
| intens_tensors = torch.nn.utils.rnn.pad_sequence( |
| intens, batch_first=True, padding_value=0 |
| ) |
| mzs_tensors=mzs_tensors.to(self.device) |
| intens_tensors=intens_tensors.to(self.device) |
| num_peaks=num_peaks.to(self.device) |
|
|
| spec_tensor,spec_mask = self.model.ms_encoder(mzs_tensors,intens_tensors,num_peaks) |
| spec_tensor=self.model.spec_esa(spec_tensor,spec_mask) |
| spec_tensor = self.model.spec_proj(spec_tensor) |
| spec_tensor = spec_tensor/spec_tensor.norm(dim=-1, keepdim=True) |
| return spec_tensor |
|
|
| def get_cos_distance(self, input_1, input_2): |
| with torch.no_grad(): |
| return input_1 @ input_2.t() |
|
|
