add LLM files

generate.py

@@ -0,0 +1,286 @@
+import numpy as np
+import pickle as pkl
+import os
+import argparse
+import pandas as pd
+from pathlib import Path
+import torch
+
+import utils.chem as uc
+import utils.torch_util as ut
+import utils.log as ul
+import utils.plot as up
+import configuration.config_default as cfgd
+import models.dataset as md
+import preprocess.vocabulary as mv
+import configuration.opts as opts
+from models.transformer.module.decode import decode
+from models.transformer.encode_decode.model import EncoderDecoder
+# from models.seq2seq.model import Model
+
+def prepare_input(opt):
+    ''' check if the input file contain pk_diff 
+    if not the pk_diff will be added 
+        '''
+    df_input = pd.read_csv(os.path.join(opt.data_path, opt.test_file_name + '.csv'), sep=",")
+    delta_pkis=['(-0.5, 0.5]','(0.5, 1.5]','(1.5, 2.5]','(2.5, 4.5]','(4.5, 6.5]','(6.5, 8.5]','(8.5, 10.5]','(10.5, inf]']
+    # TODO：这是分子到分子的generate数据处理？
+    # if "Delta_pki" not in df_input.columns:
+    #     newData=[]
+    #     for idx,irow in df_input.iterrows():
+    #         for idelta_pki in delta_pkis:
+    #             newData.append([irow['fromVarSMILES'],irow['constantSMILES'],idelta_pki])
+    #     df_new=pd.DataFrame(newData, columns=['fromVarSMILES','constantSMILES','Delta_pki'])
+    #     df_new.to_csv(os.path.join(opt.data_path, opt.test_file_name + '_prepared'+'.csv'), index=None)
+    #     return opt.test_file_name + '_prepared'
+    return opt.test_file_name
+
+class GenerateRunner():
+
+    def __init__(self, opt):
+
+        # self.save_path = os.path.join('experiments', opt.save_directory, opt.test_file_name,
+                                    #   f'evaluation_{opt.epoch}')
+        path = Path(os.path.join(opt.save_directory))
+        path.mkdir(parents=True, exist_ok=True)
+        self.save_path = os.path.join(path)
+        self.exist_flag = Path(f'{self.save_path}/generated_molecules.csv').exists()
+        self.overwrite = opt.overwrite
+        self.dev_no = opt.dev_no
+        global LOG
+        LOG = ul.get_logger(name="generate",
+                            log_path=os.path.join(self.save_path, 'generate.log'))
+        LOG.info(opt)
+        LOG.info("Save directory: {}".format(self.save_path))
+
+        # Load vocabulary
+        with open(os.path.join(opt.vocab_path, 'vocab.pkl'), "rb") as input_file:
+            vocab = pkl.load(input_file)
+        self.vocab = vocab
+        self.tokenizer = mv.SMILESTokenizer()
+
+    def initialize_dataloader(self, opt, vocab, test_file):
+        """
+        Initialize dataloader
+        :param opt:
+        :param vocab: vocabulary
+        :param test_file: test_file_name
+        :return:
+        """
+
+        # Read test
+        data = pd.read_csv(os.path.join(opt.data_path, test_file + '.csv'), sep=",")
+        dataset = md.Dataset(data=data, vocabulary=vocab, tokenizer=self.tokenizer, prediction_mode=True)
+        dataloader = torch.utils.data.DataLoader(dataset, opt.batch_size,
+                                                 shuffle=False, collate_fn=md.Dataset.collate_fn)
+        return dataloader
+
+    def generate(self, opt):
+        if not self.overwrite and self.exist_flag:
+            print('GENERATED MOL EXIST, SKIP GENERATING!')
+            return
+        # set device
+        #device = ut.allocate_gpu()
+        # torch.cuda.set_device(1)
+        # current_device = torch.cuda.current_device()
+        # print("当前使用的 CUDA 设备编号是:", current_device)
+        device = torch.device(f'cuda:{self.dev_no}')
+        # 构造loader
+        dataloader_test = self.initialize_dataloader(opt, self.vocab, opt.test_file_name)
+
+        # Load model
+        file_name = os.path.join(opt.model_path, f'model_{opt.epoch}.pt')
+        if opt.model_choice == 'transformer':
+            model = EncoderDecoder.load_from_file(file_name)
+            model.to(device)
+            model.eval()
+        elif opt.model_choice == 'seq2seq':
+            model = Model.load_from_file(file_name, evaluation_mode=True)
+            # move to GPU
+            model.network.encoder.to(device)
+            model.network.decoder.to(device)
+        # TODO: 有没有可能超长？模型崩溃调整长度，规则是2的倍数
+        max_len = cfgd.DATA_DEFAULT['max_sequence_length']
+        df_list = []
+        sampled_smiles_list = []
+        for j, batch in enumerate(ul.progress_bar(dataloader_test, total=len(dataloader_test))):
+
+            # df是dataframe 是一行的原始数据
+            src, source_length, _, src_mask, _, _, df = batch
+
+            # Move to GPU
+            src = src.to(device)
+            src_mask = src_mask.to(device)
+            smiles= self.sample(opt.model_choice, model, src, src_mask,
+                                                                       source_length,
+                                                                       opt.decode_type,
+                                                                       num_samples=opt.num_samples,
+                                                                       max_len=max_len,
+                                                                       device=device)
+
+            df_list.append(df)
+            sampled_smiles_list.extend(smiles)
+
+        # prepare dataframe
+        data_sorted = pd.concat(df_list)
+        sampled_smiles_list = np.array(sampled_smiles_list)
+
+        for i in range(opt.num_samples):
+            data_sorted['Predicted_smi_{}'.format(i + 1)] = sampled_smiles_list[:, i]
+
+        result_path = os.path.join(self.save_path, "generated_molecules.csv")
+        LOG.info("Save to {}".format(result_path))
+        data_sorted.to_csv(result_path, index=False)
+
+    def sample(self, model_choice, model, src, src_mask, source_length, decode_type, num_samples=10,
+               max_len=cfgd.DATA_DEFAULT['max_sequence_length'],
+               device=None):
+        batch_size = src.shape[0]
+        num_valid_batch = np.zeros(batch_size)  # current number of unique and valid samples out of total sampled
+        num_valid_batch_total = np.zeros(batch_size)  # current number of sampling times no matter unique or valid
+        num_valid_batch_desired = np.asarray([num_samples] * batch_size)
+        unique_set_num_samples = [set() for i in range(batch_size)]   # for each starting molecule
+        batch_index = torch.LongTensor(range(batch_size))
+        batch_index_current = torch.LongTensor(range(batch_size)).to(device)
+        # TODO:这个好像没有用到？
+        start_mols = []
+        # zeros correspondes to ****** which is valid according to RDKit
+        sequences_all = torch.ones((num_samples, batch_size, max_len))
+        sequences_all = sequences_all.type(torch.LongTensor)
+        max_trials = 100000  # Maximum trials for sampling
+        current_trials = 0
+
+        # greedy意思是只尝试一次生成，成了就有分子式，没成的话就没有
+        if decode_type == 'greedy':
+            max_trials = 1
+
+        # Set of unique starting molecules
+        if src is not None:
+            # 这里需要修改，delta_value并不是放在第一位置
+            start_ind = len(cfgd.PROPERTIES)
+            for ibatch in range(batch_size):
+                source_smi = self.tokenizer.untokenize(self.vocab.decode(src[ibatch].tolist()[start_ind:]))
+                source_smi = uc.get_canonical_smile(source_smi)
+                if source_smi:
+                    # 先添加source，用于后面去重，TODO: 但这里也不太对，因为这里已经是被mmpdb分开的，而不是一个完整的SMILES
+                    unique_set_num_samples[ibatch].add(source_smi)
+                    start_mols.append(source_smi)
+
+        with torch.no_grad():
+            if model_choice == 'seq2seq':
+                encoder_outputs, decoder_hidden = model.network.encoder(src, source_length)
+            while not all(num_valid_batch >= num_valid_batch_desired) and current_trials < max_trials:
+                current_trials += 1
+
+                # batch input for current trial
+                if src is not None:
+                    # 这个不就是全选嘛？
+                    src_current = src.index_select(0, batch_index_current)
+                if src_mask is not None:
+                    mask_current = src_mask.index_select(0, batch_index_current)
+                batch_size = src_current.shape[0]
+
+                # sample molecule
+                if model_choice == 'transformer':
+                    sequences = decode(model, src_current, mask_current, max_len, decode_type)
+                    padding = (0, max_len-sequences.shape[1],
+                               0, 0)
+                    sequences = torch.nn.functional.pad(sequences, padding)
+                elif model_choice == 'seq2seq':
+                    sequences = self.sample_seq2seq(model, mask_current, batch_index_current, decoder_hidden,
+                                                    encoder_outputs, max_len, device)
+                else:
+                    LOG.info('Specify transformer or seq2seq for model_choice')
+
+                # Check valid and unique
+                smiles = []
+                is_valid_index = []
+                batch_index_map = dict(zip(list(range(batch_size)), batch_index_current))
+                # Valid, ibatch index is different from original, need map back
+                for ibatch in range(batch_size):
+                    seq = sequences[ibatch]
+                    smi = self.tokenizer.untokenize(self.vocab.decode(seq.cpu().numpy()))
+                    smi = uc.get_canonical_smile(smi)
+                    smiles.append(smi)
+                    # valid and not same as starting molecules
+                    if uc.is_valid(smi):
+                        is_valid_index.append(ibatch)
+                    # total sampled times
+                    num_valid_batch_total[batch_index_map[ibatch]] += 1
+
+                # Check if duplicated and update num_valid_batch and unique
+                for good_index in is_valid_index:
+                    index_in_original_batch = batch_index_map[good_index]
+                    if smiles[good_index] not in unique_set_num_samples[index_in_original_batch]:
+                        unique_set_num_samples[index_in_original_batch].add(smiles[good_index])
+                        num_valid_batch[index_in_original_batch] += 1
+
+                        sequences_all[int(num_valid_batch[index_in_original_batch] - 1), index_in_original_batch, :] = \
+                            sequences[good_index]
+
+                not_completed_index = np.where(num_valid_batch < num_valid_batch_desired)[0]
+                # 选择未生成满的source样本继续生成
+                if len(not_completed_index) > 0:
+                    batch_index_current = batch_index.index_select(0, torch.LongTensor(not_completed_index)).to(device)
+
+        # Convert to SMILES
+        smiles_list = [] # [batch, topk]
+        seqs = np.asarray(sequences_all.numpy())
+        # [num_sample, batch_size, max_len]
+        batch_size = len(seqs[0])
+        for ibatch in range(batch_size):
+            topk_list = []
+            for k in range(num_samples):
+                seq = seqs[k, ibatch, :]
+                topk_list.extend([self.tokenizer.untokenize(self.vocab.decode(seq))])
+            smiles_list.append(topk_list)
+
+        return smiles_list
+
+    def sample_seq2seq(self, model, mask, batch_index_current, decoder_hidden, encoder_outputs, max_len, device):
+        # batch size will change when some of the generated molecules are valid
+        encoder_outputs_current = encoder_outputs.index_select(0, batch_index_current)
+        batch_size = encoder_outputs_current.shape[0]
+
+        # start token
+        start_token = torch.zeros(batch_size, dtype=torch.long)
+        start_token[:] = self.vocab["^"]
+        decoder_input = start_token.to(device)
+        sequences = []
+        mask = torch.squeeze(mask, 1).to(device)
+
+        # initial decoder hidden states
+        if isinstance(decoder_hidden, tuple):
+            decoder_hidden_current = (decoder_hidden[0].index_select(1, batch_index_current),
+                                      decoder_hidden[1].index_select(1, batch_index_current))
+        else:
+            decoder_hidden_current = decoder_hidden.index_select(1, batch_index_current)
+        for i in range(max_len):
+            logits, decoder_hidden_current = model.network.decoder(decoder_input.unsqueeze(1),
+                                                                  decoder_hidden_current,
+                                                                  encoder_outputs_current, mask)
+            logits = logits.squeeze(1)
+            probabilities = logits.softmax(dim=1)  # torch.Size([batch_size, vocab_size])
+            topi = torch.multinomial(probabilities, 1)  # torch.Size([batch_size, 1])
+            decoder_input = topi.view(-1).detach()
+            sequences.append(decoder_input.view(-1, 1))
+
+        sequences = torch.cat(sequences, 1)
+        return sequences
+
+def run_main():
+    """Main function."""
+    parser = argparse.ArgumentParser(
+        description='generate.py',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    opts.generate_opts(parser)
+    opt = parser.parse_args()
+    opt.test_file_name = prepare_input(opt)
+
+    runner = GenerateRunner(opt)
+    runner.generate(opt)
+
+
+if __name__ == "__main__":
+    run_main()

utils/chem.py

@@ -0,0 +1,65 @@
+"""
+RDKit util functions.
+"""
+import rdkit.Chem as rkc
+from rdkit.Chem import AllChem
+from rdkit import DataStructs
+
+def disable_rdkit_logging():
+    """
+    Disables RDKit whiny logging.
+    """
+    import rdkit.RDLogger as rkl
+    logger = rkl.logger()
+    logger.setLevel(rkl.ERROR)
+
+    import rdkit.rdBase as rkrb
+    rkrb.DisableLog('rdApp.error')
+
+
+disable_rdkit_logging()
+
+def to_fp_ECFP(smi):
+    if smi:
+        mol = rkc.MolFromSmiles(smi)
+        if mol is None:
+            return None
+        return AllChem.GetMorganFingerprint(mol, 2)
+
+def tanimoto_similarity_pool(args):
+    return tanimoto_similarity(*args)
+
+def tanimoto_similarity(smi1, smi2):
+    fp1, fp2 = None, None
+    if smi1 and type(smi1)==str and len(smi1)>0:
+        fp1 = to_fp_ECFP(smi1)
+    if smi2 and type(smi2)==str and len(smi2)>0:
+        fp2 = to_fp_ECFP(smi2)
+
+    if fp1 is not None and fp2 is not None:
+        return DataStructs.TanimotoSimilarity(fp1, fp2)
+    else:
+        return None
+
+def is_valid(smi):
+    return 1 if to_mol(smi) else 0
+
+def to_mol(smi):
+    """
+    Creates a Mol object from a SMILES string.
+    :param smi: SMILES string.
+    :return: A Mol object or None if it's not valid.
+    """
+    if isinstance(smi, str) and smi and len(smi)>0 and smi != 'nan':
+        return rkc.MolFromSmiles(smi)
+
+def get_canonical_smile(smile):
+    if smile != 'None':
+        mol = rkc.MolFromSmiles(smile)
+        if mol is not None:
+            smi = rkc.MolToSmiles(mol, canonical=True, doRandom=False, isomericSmiles=False)
+            return smi
+        else:
+            return None
+    else:
+        return None

utils/file.py

@@ -0,0 +1,29 @@
+import os
+
+def make_directory(file, is_dir=True):
+    dirs = file.split('/')[:-1] if not is_dir else file.split('/')
+    path = '/' if file.startswith('/') else ''
+    for dir in dirs:
+        path = os.path.join(path, dir)
+        if not os.path.exists(path):
+            os.makedirs(path)
+
+def get_parent_dir(file):
+    dirs = file.split('/')[:-1]
+    path = ''
+    for dir in dirs:
+        path = os.path.join(path, dir)
+    if file.startswith('/'):
+        path = '/' + path
+    return path
+
+def chunkIt(seq, num):
+    avg = len(seq) / float(num)
+    out = []
+    last = 0.0
+
+    while last < len(seq):
+        out.append(seq[int(last):int(last + avg)])
+        last += avg
+
+    return out

utils/log.py

@@ -0,0 +1,32 @@
+import logging
+import tqdm
+
+import utils.file as uf
+
+
+def get_logger(name, log_path, isMain=False, level=logging.INFO):
+    formatter = logging.Formatter(
+        fmt="%(asctime)s: %(module)s.%(funcName)s +%(lineno)s: %(levelname)-8s %(message)s",
+        datefmt="%H:%M:%S"
+    )
+
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    # # Logging to console
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+
+    # Logging to a file
+    if isMain:
+        uf.make_directory(log_path, is_dir=False)
+        file_handler = logging.FileHandler(log_path)
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    return logger
+
+
+def progress_bar(iterable, total, **kwargs):
+    return tqdm.tqdm(iterable=iterable, total=total, ascii=True, **kwargs)

utils/plot.py

@@ -0,0 +1,84 @@
+import numpy as np
+import os
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from scipy.stats import gaussian_kde
+mpl.use('Agg')
+
+
+def hist_box(data_frame, field, name="hist_box", path="./", title=None):
+
+    title = title if title else field
+    fig, axs = plt.subplots(1,2,figsize=(10,4))
+    data_frame[field].plot.hist(bins=100, title=title, ax=axs[0])
+    data_frame.boxplot(field, ax=axs[1])
+    plt.title(title)
+    plt.suptitle("")
+
+    plt.savefig(os.path.join(path, '{}.png'.format(name)), bbox_inches='tight')
+    plt.close()
+
+def hist(data_frame, field, name="hist", path="./", title=None):
+
+
+    title = title if title else field
+
+    plt.hist(data_frame[field])
+    plt.title(title)
+    plt.savefig(os.path.join(path, '{}.png'.format(name)), bbox_inches='tight')
+    plt.close()
+
+def hist_box_list(data_list, name="hist_box", path="./", title=None):
+
+    fig, axs = plt.subplots(1,2,figsize=(10,4))
+    axs[0].hist(data_list, bins=100)
+    axs[0].set_title(title)
+    axs[1].boxplot(data_list)
+    axs[1].set_title(title)
+
+    plt.savefig(os.path.join(path, '{}.png'.format(name)), bbox_inches='tight')
+    plt.close()
+
+def scatter_hist(x, y, name, path, field=None, lims=None):
+    fig, axs = plt.subplots(1, 2, figsize=(10, 4))
+    n = len(x)
+    xy = np.vstack([x+ 0.00001 * np.random.rand(n), y+ 0.00001 * np.random.rand(n)])
+    z = gaussian_kde(xy)(xy)
+    axs[0].scatter(x, y, c=z, s=3, marker='o', alpha=0.2)
+    lims = [np.min([axs[0].get_xlim(), axs[0].get_ylim()]), np.max([axs[0].get_xlim(), axs[0].get_ylim()])] if lims is None else lims
+    axs[0].plot(lims, lims, 'k-', alpha=0.75)
+    axs[0].set_aspect('equal')
+    axs[0].set_xlim(lims)
+    axs[0].set_ylim(lims)
+    xlabel = ""
+    ylabel = ""
+    if "delta" in field:
+        if "data" in field:
+            axs[0].set_xlabel(r'$\Delta LogD$ (experimental)')
+            axs[0].set_ylabel(r'$\Delta LogD$ (calculated)')
+            xlabel = 'Delta LogD (experimental)'
+            ylabel = 'Delta LogD (calculated)'
+        elif "predict" in field:
+            axs[0].set_xlabel(r'$\Delta LogD$ (desirable)')
+            axs[0].set_ylabel(r'$\Delta LogD$ (generated)')
+            xlabel = 'Delta LogD (desirable)'
+            ylabel = 'Delta LogD (generated)'
+    if "single" in field:
+        if "data" in field:
+            xlabel, ylabel = 'LogD (experimental)', 'LogD (calculated)'
+            axs[0].set_xlabel(xlabel)
+            axs[0].set_ylabel(ylabel)
+        elif "predict" in field:
+            xlabel, ylabel = 'LogD (desirable)', 'LogD (generated)'
+            axs[0].set_xlabel(xlabel)
+            axs[0].set_ylabel(ylabel)
+    bins = np.histogram(np.hstack((x, y)), bins=100)[1]  # get the bin edges
+    kwargs = dict(histtype='stepfilled', alpha=0.3, density=False, bins=bins, stacked=False)
+    axs[1].hist(x, **kwargs, color='b', label=xlabel)
+    axs[1].hist(y, **kwargs, color='r', label=ylabel)
+    plt.ylabel('Frequency')
+    plt.legend(loc='upper left')
+    plt.savefig(os.path.join(path, '{}.png'.format(name)), bbox_inches='tight')
+    plt.close()
+

utils/torch_util.py

@@ -0,0 +1,32 @@
+"""
+PyTorch related util functions
+"""
+import torch
+import os 
+def allocate_gpu(id=None):
+    '''
+    choose the free gpu in the node
+    '''
+    v = torch.empty(1)
+    if id is not None:
+        return torch.device("cuda:{}".format(str(id)))
+    else:
+        for i in range(8):
+            try:
+                dev_name = "cuda:{}".format(str(i))
+                v = v.to(dev_name)
+                print("Allocating cuda:{}.".format(i))
+
+                return v.device
+            except Exception as e:
+                pass
+        print("CUDA error: all CUDA-capable devices are busy or unavailable")
+        return v.device
+
+def allocate_gpu_multi(id=None):
+
+    os.environ['CUDA_VISIBLE_DEVICES']='1'
+    device=torch.device("cuda:1" if torch.cuda.is_available() else 'cpu')
+    os.environ['CUDA_VISIBLE_DEVICES']='0'
+    device=torch.device("cuda:1" if torch.cuda.is_available() else 'cpu')
+    return device