generate.py

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
        self.device = torch.device('cpu')
        # 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()

        # 加载模型
        file_name = os.path.join(opt.model_path, f'model_{opt.epoch}.pt')
        if opt.model_choice == 'transformer':
            self.model = EncoderDecoder.load_from_file(file_name)
            self.model.to(self.device)
            self.model.eval()
        elif opt.model_choice == 'seq2seq':
            self.model = Model.load_from_file(file_name, evaluation_mode=True)
            self.model.network.encoder.to(self.device)
            self.model.network.decoder.to(self.device)

    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}')
        device = torch.device('cpu')
        print(f"-------device:---------")
        print(device)
        # 构造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)
    print("--------parser-------------")
    print(parser)
    opts.generate_opts(parser)
    opt = parser.parse_args()
    opt.test_file_name = prepare_input(opt)
    print("opt输出如下")
    print(opt)
    runner = GenerateRunner(opt)
    print()
    runner.generate(opt)


if __name__ == "__main__":
    run_main()