add files

common/utils.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class Data_Type(Enum):
+    frag = 'frag'
+    whole = 'whole'

configuration/config_default.py

@@ -0,0 +1,18 @@
+import math
+
+# Data
+DATA_DEFAULT = {
+    'max_sequence_length': 256,
+    'padding_value': 0
+}
+
+# Properties
+PROPERTIES = ['pki']
+
+
+# For Test_Property test
+LOD_MIN = 1.0
+LOD_MAX = 3.4
+
+
+

configuration/opts.py

@@ -0,0 +1,141 @@
+""" Implementation of all available options """
+from __future__ import print_function
+
+
+def train_opts(parser):
+    # Transformer or Seq2Seq
+    parser.add_argument('--model-choice', required=True, help="transformer or seq2seq")
+    # Common training options
+    group = parser.add_argument_group('Training_options')
+    group.add_argument('--batch-size', type=int, default=512,
+                       help='Batch size for training')
+    group.add_argument('--num-epoch', type=int, default=200,
+                       help='Number of training steps')
+    group.add_argument('--starting-epoch', type=int, default=1,
+                       help="Training from given starting epoch")
+    # Input output settings
+    group = parser.add_argument_group('Input-Output')
+    group.add_argument('--data-path', required=True,
+                       help="""Input data path""")
+    group.add_argument('--save-directory', default='finetune-TLR7',
+                       help="""Result save directory""")
+
+    subparsers = parser.add_subparsers()
+    transformer_parser = subparsers.add_parser('transformer')
+    train_opts_transformer(transformer_parser)
+
+    seq2seq_parser = subparsers.add_parser('seq2seq')
+    train_opts_seq2seq(seq2seq_parser)
+
+def train_opts_transformer(parser):
+    # Model architecture options
+    group = parser.add_argument_group('Model')
+    group.add_argument('--vocab-path', required=False, default='',
+                       help="vocab path for finetuning")
+    group.add_argument('--pretrain-path', default='',
+                       help="pretrain directory")
+    group.add_argument('-N', type=int, default=6,
+                       help="number of encoder and decoder")
+    group.add_argument('-H', type=int, default=8,
+                       help="heads of attention")
+    group.add_argument('-d-model', type=int, default=128,
+                       help="embedding dimension, model dimension")
+    group.add_argument('-d-ff', type=int, default=2048,
+                       help="dimension in feed forward network")
+    # Regularization
+    group.add_argument('--dropout', type=float, default=0.1,
+                       help="Dropout probability; applied in LSTM stacks.")
+    group.add_argument('--label-smoothing', type=float, default=0.0,
+                       help="""Label smoothing value epsilon.
+                       Probabilities of all non-true labels
+                       will be smoothed by epsilon / (vocab_size - 1).
+                       Set to zero to turn off label smoothing.
+                       For more detailed information, see:
+                       https://arxiv.org/abs/1512.00567""")
+    # Optimization options
+    group = parser.add_argument_group('Optimization')
+    group.add_argument('--factor', type=float, default=1.0,
+                       help="""Factor multiplied to the learning rate scheduler formula in NoamOpt. 
+                       For more information about the formula, 
+                       see paper Attention Is All You Need https://arxiv.org/pdf/1706.03762.pdf""")
+    group.add_argument('--warmup-steps', type=int, default=4000,
+                       help="""Number of warmup steps for custom decay.""")
+    group.add_argument('--adam-beta1', type=float, default=0.9,
+                       help="""The beta1 parameter for Adam optimizer""")
+    group.add_argument('--adam-beta2', type=float, default=0.98,
+                       help="""The beta2 parameter for Adam optimizer""")
+    group.add_argument('--adam-eps', type=float, default=1e-9,
+                       help="""The eps parameter for Adam optimizer""")
+
+
+def train_opts_seq2seq(parser):
+    # Model architecture options
+    group = parser.add_argument_group('Model')
+    group.add_argument("--num-layers", "-l", help="Number of RNN layers of the model",
+                        default=5, type=int)
+    group.add_argument("--layer-size", "-s", help="Size of each of the RNN layers",
+                        default=512, type=int)
+    group.add_argument("--cell-type", "-c",
+                        help="Type of cell used in RNN [gru, lstm]",
+                        default='lstm', type=str)
+    group.add_argument("--embedding-layer-size", "-e", help="Size of the embedding layer",
+                        default=256, type=int)
+    group.add_argument("--dropout", "-d", help="Amount of dropout between layers ",
+                        default=0.3, type=float)
+    group.add_argument("--bidirectional", "--bi", help="Encoder bidirectional", action="store_false")
+    group.add_argument("--bidirect-model",
+                        help="Method to use encoder hidden state for initialising decoder['concat', 'addition', 'none']",
+                        default='addition', type=str)
+    group.add_argument("--attn-model", help="Attention model ['dot', 'general', 'concat']",
+                        default='dot', type=str)
+    # Optimization options
+    group = parser.add_argument_group('Optimization')
+    group.add_argument('--learning-rate', type=float, default=0.0001,
+                       help="""Starting learning rate""")
+    group.add_argument("--clip-gradient-norm", help="Clip gradients to a given norm",
+                       default=1.0, type=float)
+
+
+def generate_opts(parser):
+    # Transformer or Seq2Seq
+    parser.add_argument('--model-choice', required=True, help="transformer or seq2seq")
+    """Input output settings"""
+    group = parser.add_argument_group('Input-Output')
+    group.add_argument('--data-path', required=True,
+                       help="""Input data path""")
+    group.add_argument('--test-file-name', required=True, help="""test file name without .csv,
+        [test, test_not_in_train, test_unseen_L-1_S01_C10_range]""")
+    group.add_argument('--save-directory', default='evaluation',
+                       help="""Result save directory""")
+    group.add_argument('--vocab-path', required=False, default='',
+                       help="vocab path for finetuning")
+    # Model to be used for generating molecules
+    group = parser.add_argument_group('Model')
+    group.add_argument('--model-path', help="""Model path""", required=True)
+    group.add_argument('--epoch', type=int, help="""Which epoch to use""", required=True)
+    # General
+    group = parser.add_argument_group('General')
+    group.add_argument('--batch-size', type=int, default=64,
+                       help='Batch size for training')
+    group.add_argument('--num-samples', type=int, default=50,
+                       help='Number of molecules to be generated')
+    group.add_argument('--decode-type',type=str, default='multinomial',help='decode strategy')
+    group.add_argument('--dev-no',type=int, default=0, help='using device')
+    group.add_argument('--overwrite',type=bool, default=False, help='whether overwrite exist file')
+
+
+def evaluation_opts(parser):
+    """Evaluation options (compute properties)"""
+    group = parser.add_argument_group('General')
+    group.add_argument('--data-path', required=True,
+                       help="""Input data path for generated molecules""")
+    group.add_argument('--num-samples', type=int, default=10,
+                       help='Number of molecules generated')
+    group = parser.add_argument_group('Evaluation')
+    group.add_argument('--range-evaluation', default='',
+                       help='[ , lower, higher]; set lower when evaluating test_unseen_L-1_S01_C10_range')
+    group = parser.add_argument_group('MMP')
+    group.add_argument('--mmpdb-path', help='mmpdb path; download from https://github.com/rdkit/mmpdb')
+    group.add_argument('--train-path', help='Training data path')
+    group.add_argument('--only-desirable', help='Only check generated molecules with desirable properties',
+                       action="store_true")

models/dataset.py

@@ -0,0 +1,175 @@
+# coding=utf-8
+
+"""
+Implementation of a SMILES dataset.
+"""
+import pandas as pd
+
+import torch
+import torch.utils.data as tud
+from torch.autograd import Variable
+
+import configuration.config_default as cfgd
+from models.transformer.module.subsequent_mask import subsequent_mask
+
+from rdkit.Chem.SaltRemover import SaltRemover
+import random
+import rdkit.Chem as rkc
+from common.utils import Data_Type
+
+class Dataset(tud.Dataset):
+    """Custom PyTorch Dataset that takes a file containing
+    Source_Mol_ID,Target_Mol_ID,Source_Mol,Target_Mol,
+    Source_Mol_LogD,Target_Mol_LogD,Delta_LogD,
+    Source_Mol_Solubility,Target_Mol_Solubility,Delta_Solubility,
+    Source_Mol_Clint,Target_Mol_Clint,Delta_Clint,
+    Transformation,Core"""
+
+    def __init__(self, data, vocabulary, tokenizer, prediction_mode=False, use_random=False, data_type=Data_Type.frag.value):
+        """
+
+        :param data: dataframe read from training, validation or test file
+        :param vocabulary: used to encode source/target tokens
+        :param tokenizer: used to tokenize source/target smiles
+        :param prediction_mode: if use target smiles or not (training or test)
+        """
+        self._vocabulary = vocabulary
+        self._tokenizer = tokenizer
+        self._data = data
+        self._prediction_mode = prediction_mode
+        self._use_random = use_random
+        self._data_type = data_type
+
+    def smiles_preprocess(self, smiles, random_type="unrestricted"):
+        """
+        Returns a random SMILES given a SMILES of a molecule.
+        :param mol: A Mol object
+        :param random_type: The type (unrestricted, restricted) of randomization performed.
+        :return : A random SMILES string of the same molecule or None if the molecule is invalid.
+        """
+        if not self._use_random:
+            return smiles
+        mol = rkc.MolFromSmiles(smiles)
+        if not mol:
+            return None
+
+        remover = SaltRemover()  ## default salt remover
+        if random_type == "unrestricted":
+            stripped = remover.StripMol(mol)
+            if stripped == None:
+                return smiles
+            ret = rkc.MolToSmiles(stripped, canonical=False, doRandom=True, isomericSmiles=False)
+            if not bool(ret):
+                return smiles
+            return ret
+        if random_type == "restricted":
+            new_atom_order = list(range(mol.GetNumAtoms()))
+            random.shuffle(new_atom_order)
+            random_mol = rkc.RenumberAtoms(mol, newOrder=new_atom_order)
+            ret = rkc.MolToSmiles(random_mol, canonical=False, isomericSmiles=False)
+            if not bool(ret):
+                return smiles
+            return ret
+        raise ValueError("Type '{}' is not valid".format(random_type))
+
+
+    def __getitem__(self, i):
+        """
+        Tokenize and encode source smile and/or target smile (if prediction_mode is True)
+        :param i:
+        :return:
+        """
+
+        row = self._data.iloc[i]
+        # tokenize and encode source smiles
+        main_cls = row['main_cls']
+        minor_cls = row['minor_cls']
+        target_name = row['target_name']
+        target_name = target_name if isinstance(target_name, str) else ''
+        value = row['Delta_Value']
+        # value = row['Delta_pki']
+        source_tokens = []
+
+        if self._data_type == Data_Type.frag.value:
+            sourceConstant = self.smiles_preprocess(row['constantSMILES'])
+            sourceVariable = self.smiles_preprocess(row['fromVarSMILES'])
+            # 先variable
+            source_tokens.extend(self._tokenizer.tokenize(sourceVariable))  ## add source variable SMILES token
+            # 接着constant
+            source_tokens.extend(self._tokenizer.tokenize(sourceConstant)) ## add source constant SMILES token
+        elif self._data_type == Data_Type.whole.value:
+            sourceSmi = self.smiles_preprocess(row['cpd1SMILES'])
+            source_tokens.extend(self._tokenizer.tokenize(sourceSmi))
+        # 再 major class eg activity
+        source_tokens.append(main_cls)
+        # 再 minor class eg Ki
+        source_tokens.append(minor_cls)
+        # 然后value
+        source_tokens.append(value)
+        # 然后target name
+        source_tokens.extend(list(target_name))
+
+        source_encoded = self._vocabulary.encode(source_tokens)
+        
+        # print(source_tokens,'\n=====\n', source_encoded)
+        # tokenize and encode target smiles if it is for training instead of evaluation
+        if not self._prediction_mode:
+            target_smi = ''
+            if self._data_type == Data_Type.frag.value:
+                target_smi = row['toVarSMILES']
+            elif self._data_type == Data_Type.whole.value:
+                target_smi = row['cpd2SMILES']
+            target_tokens = self._tokenizer.tokenize(target_smi)
+            target_encoded = self._vocabulary.encode(target_tokens)
+
+            return torch.tensor(source_encoded, dtype=torch.long), torch.tensor(
+                target_encoded, dtype=torch.long), row
+        else:
+            return torch.tensor(source_encoded, dtype=torch.long),  row
+
+    def __len__(self):
+        return len(self._data)
+
+    @classmethod
+    def collate_fn(cls, data_all):
+        # sort based on source sequence's length
+        data_all.sort(key=lambda x: len(x[0]), reverse=True)
+        is_prediction_mode = True if len(data_all[0]) == 2 else False
+        if is_prediction_mode:
+            source_encoded, data = zip(*data_all)
+            data = pd.DataFrame(data)
+        
+        else:
+            source_encoded, target_encoded, data = zip(*data_all)
+            data = pd.DataFrame(data)
+        
+        # maximum length of source sequences
+        max_length_source = max([seq.size(0) for seq in source_encoded])
+        # print('=====max len', max_length_source)
+        # padded source sequences with zeroes
+        collated_arr_source = torch.zeros(len(source_encoded), max_length_source, dtype=torch.long)
+        for i, seq in enumerate(source_encoded):
+            collated_arr_source[i, :seq.size(0)] = seq
+        # length of each source sequence
+        source_length = [seq.size(0) for seq in source_encoded]
+        source_length = torch.tensor(source_length)
+        # mask of source seqs
+        src_mask = (collated_arr_source !=0).unsqueeze(-2)
+
+        # target seq
+        if not is_prediction_mode:
+            max_length_target = max([seq.size(0) for seq in target_encoded])
+            collated_arr_target = torch.zeros(len(target_encoded), max_length_target, dtype=torch.long)
+            for i, seq in enumerate(target_encoded):
+                collated_arr_target[i, :seq.size(0)] = seq
+
+            trg_mask = (collated_arr_target != 0).unsqueeze(-2)
+            trg_mask = trg_mask & Variable(subsequent_mask(collated_arr_target.size(-1)).type_as(trg_mask))
+            trg_mask = trg_mask[:, :-1, :-1]  # save start token, skip end token
+        else:
+            trg_mask = None
+            max_length_target = None
+            collated_arr_target = None
+
+        return collated_arr_source, source_length, collated_arr_target, src_mask, trg_mask, max_length_target, data
+

models/transformer/encode_decode/clones.py

@@ -0,0 +1,8 @@
+import copy
+
+import torch.nn as nn
+
+
+def clones(module, N):
+    "Produce N identical layers."
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

models/transformer/encode_decode/decoder.py

@@ -0,0 +1,23 @@
+import torch.nn as nn
+
+from models.transformer.encode_decode.clones import clones
+from models.transformer.encode_decode.layer_norm import LayerNorm
+
+
+class Decoder(nn.Module):
+    "Generic N layer decoder with masking."
+
+    def __init__(self, layer, N):
+        super(Decoder, self).__init__()
+        self.layers = clones(layer, N)
+        self.norm = LayerNorm(layer.size)
+
+    def forward(self, x, memory, src_mask, tgt_mask):
+
+
+        memory = memory    
+                    
+        for layer in self.layers:
+            x = layer(x, memory, src_mask, tgt_mask)
+
+        return self.norm(x)

models/transformer/encode_decode/decoder_layer.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+
+from models.transformer.encode_decode.clones import clones
+from models.transformer.encode_decode.sublayer_connection import SublayerConnection
+
+class DecoderLayer(nn.Module):
+    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
+
+    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
+        super(DecoderLayer, self).__init__()
+        self.size = size
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.sublayer = clones(SublayerConnection(size, dropout), 3)
+
+    def forward(self, x, memory, src_mask, tgt_mask):
+        "Follow Figure 1 (right) for connections."
+
+        m = memory    
+
+        x = self.sublayer[0](x, lambda x: self.self_attn(
+            x, x, x, tgt_mask))
+        x = self.sublayer[1](x, lambda x: self.src_attn(
+            x, m, m, src_mask))
+
+        return self.sublayer[2](x, self.feed_forward)
+

models/transformer/encode_decode/encoder.py

@@ -0,0 +1,20 @@
+import torch.nn as nn
+
+from models.transformer.encode_decode.clones import clones
+from models.transformer.encode_decode.layer_norm import LayerNorm
+
+
+class Encoder(nn.Module):
+    "Core encoder is a stack of N layers"
+
+    def __init__(self, layer, N):
+        super(Encoder, self).__init__()
+        self.layers = clones(layer, N)
+        self.norm = LayerNorm(layer.size)
+        
+    def forward(self, x, mask):
+        "Pass the input (and mask) through each layer in turn."
+
+        for layer in self.layers:
+            x = layer(x, mask)
+        return self.norm(x)

models/transformer/encode_decode/encoder_layer.py

@@ -0,0 +1,20 @@
+import torch.nn as nn
+
+from models.transformer.encode_decode.clones import clones
+from models.transformer.encode_decode.sublayer_connection import SublayerConnection
+
+
+class EncoderLayer(nn.Module):
+    "Encoder is made up of self-attn and feed forward (defined below)"
+
+    def __init__(self, size, self_attn, feed_forward, dropout):
+        super(EncoderLayer, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.sublayer = clones(SublayerConnection(size, dropout), 2)
+        self.size = size
+
+    def forward(self, x, mask):
+        "Follow Figure 1 (left) for connections."
+        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
+        return self.sublayer[1](x, self.feed_forward)

models/transformer/encode_decode/layer_norm.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+import torch
+
+
+class LayerNorm(nn.Module):
+    "Construct a layernorm module (See citation for details)."
+
+    def __init__(self, features, eps=1e-6):
+        super(LayerNorm, self).__init__()
+        self.a_2 = nn.Parameter(torch.ones(features))
+        self.b_2 = nn.Parameter(torch.zeros(features))
+        self.eps = eps
+
+    def forward(self, x):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

models/transformer/encode_decode/model.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+import copy
+
+from models.transformer.module.positional_encoding import PositionalEncoding
+from models.transformer.module.positionwise_feedforward import PositionwiseFeedForward
+from models.transformer.module.multi_headed_attention import MultiHeadedAttention
+from models.transformer.module.embeddings import Embeddings
+from models.transformer.encode_decode.encoder import Encoder
+from models.transformer.encode_decode.decoder import Decoder
+from models.transformer.encode_decode.encoder_layer import EncoderLayer
+from models.transformer.encode_decode.decoder_layer import DecoderLayer
+from models.transformer.module.generator import Generator
+
+class EncoderDecoder(nn.Module):
+    """
+    A standard Encoder-Decoder architecture.
+    """
+
+    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
+        super(EncoderDecoder, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.src_embed = src_embed
+        self.tgt_embed = tgt_embed
+        self.generator = generator
+
+    def forward(self, src, tgt, src_mask, tgt_mask):
+        "Take in and process masked src and target sequences."
+        return self.decode(self.encode(src, src_mask), src_mask,
+                           tgt, tgt_mask)
+
+    def encode(self, src, src_mask):
+        return self.encoder(self.src_embed(src), src_mask)
+
+    def decode(self, memory, src_mask, tgt, tgt_mask):
+        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
+
+    @classmethod
+    def make_model(cls, src_vocab, tgt_vocab, N=6,
+                   d_model=256, d_ff=2048, h=8, dropout=0.1):
+        "Helper: Construct a model from hyperparameters."
+        c = copy.deepcopy
+        attn = MultiHeadedAttention(h, d_model)
+        ff = PositionwiseFeedForward(d_model, d_ff, dropout)
+        position = PositionalEncoding(d_model, dropout)
+        model = EncoderDecoder(
+            Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
+            Decoder(DecoderLayer(d_model, c(attn), c(attn),
+                                 c(ff), dropout), N),
+            nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
+            nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
+            Generator(d_model, tgt_vocab))
+
+        # This was important from their code.
+        # Initialize parameters with Glorot / fan_avg.
+        for p in model.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform(p)
+
+        return model
+
+    @classmethod
+    def load_from_file(cls, file_path):
+        # Load model
+        checkpoint = torch.load(file_path, map_location='cuda:0')
+        para_dict = checkpoint['model_parameters']
+        vocab_size = para_dict['vocab_size']
+        model = EncoderDecoder.make_model(vocab_size, vocab_size, para_dict['N'],
+                                  para_dict['d_model'], para_dict['d_ff'],
+                                  para_dict['H'], para_dict['dropout'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+        return model

models/transformer/encode_decode/sublayer_connection.py

@@ -0,0 +1,18 @@
+import torch.nn as nn
+from models.transformer.encode_decode.layer_norm import LayerNorm
+
+
+class SublayerConnection(nn.Module):
+    """
+    A residual connection followed by a layer norm.
+    Note for code simplicity the norm is first as opposed to last.
+    """
+
+    def __init__(self, size, dropout):
+        super(SublayerConnection, self).__init__()
+        self.norm = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, sublayer):
+        "Apply residual connection to any sublayer with the same size."
+        return x + self.dropout(sublayer(self.norm(x)))

models/transformer/module/decode.py

@@ -0,0 +1,34 @@
+import torch
+from torch.autograd import Variable
+
+from models.transformer.module.subsequent_mask import subsequent_mask
+
+
+def decode(model, src, src_mask, max_len, type):
+    ys = torch.ones(1)
+    ys = ys.repeat(src.shape[0], 1).view(src.shape[0], 1).type_as(src.data)
+    # ys shape [batch_size, 1]
+    encoder_outputs = model.encode(src, src_mask)
+    break_condition = torch.zeros(src.shape[0], dtype=torch.bool)
+    for i in range(max_len-1):
+        with torch.no_grad():
+            out = model.decode(encoder_outputs, src_mask, Variable(ys),
+                                      Variable(subsequent_mask(ys.size(1)).type_as(src.data)))
+
+            log_prob = model.generator(out[:, -1])
+            prob = torch.exp(log_prob)
+
+            if type == 'greedy':
+                _, next_word = torch.max(prob, dim = 1)
+                ys = torch.cat([ys, next_word.unsqueeze(-1)], dim=1)  # [batch_size, i]
+            elif type == 'multinomial':
+                next_word = torch.multinomial(prob, 1)
+                ys = torch.cat([ys, next_word], dim=1) #[batch_size, i]
+                next_word = torch.squeeze(next_word)
+
+            break_condition = (break_condition | (next_word.to('cpu')==2))
+            if all(break_condition): # end token
+                break
+
+    return ys
+

models/transformer/module/embeddings.py

@@ -0,0 +1,13 @@
+import torch.nn as nn
+import math
+
+
+class Embeddings(nn.Module):
+    def __init__(self, d_model, vocab):
+        super(Embeddings, self).__init__()
+        # weight matrix, each row present one word
+        self.lut = nn.Embedding(vocab, d_model)
+        self.d_model = d_model
+
+    def forward(self, x):
+        return self.lut(x) * math.sqrt(self.d_model)

models/transformer/module/generator.py

@@ -0,0 +1,13 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Generator(nn.Module):
+    "Define standard linear + softmax generation step."
+
+    def __init__(self, d_model, vocab):
+        super(Generator, self).__init__()
+        self.proj = nn.Linear(d_model, vocab)
+
+    def forward(self, x):
+        return F.log_softmax(self.proj(x), dim=-1)

models/transformer/module/label_smoothing.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+import torch
+from torch.autograd import Variable
+
+
+class LabelSmoothing(nn.Module):
+    "Implement label smoothing."
+
+    def __init__(self, size, padding_idx, smoothing=0.00):
+        super(LabelSmoothing, self).__init__()
+        self.criterion = nn.KLDivLoss(size_average=False)
+        self.padding_idx = padding_idx
+        self.confidence = 1.0 - smoothing
+        self.smoothing = smoothing
+        self.size = size
+        self.true_dist = None
+
+    def forward(self, x, target):
+        assert x.size(1) == self.size
+        true_dist = x.data.clone()
+        true_dist.fill_(self.smoothing / (self.size - 2))
+        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
+        true_dist[:, self.padding_idx] = 0
+        mask = torch.nonzero(target.data == self.padding_idx)
+
+        if mask.dim() > 0:
+            true_dist.index_fill_(0, mask.squeeze(), 0.0)
+        self.true_dist = true_dist
+        return self.criterion(x, Variable(true_dist, requires_grad=False))

models/transformer/module/multi_headed_attention.py

@@ -0,0 +1,55 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.transformer.encode_decode.clones import clones
+
+
+def attention(query, key, value, mask=None, dropout=None):
+    "Compute 'Scaled Dot Product Attention'"
+    d_k = query.size(-1)
+    scores = torch.matmul(query, key.transpose(-2, -1)) \
+        / math.sqrt(d_k)
+    if mask is not None:
+        scores = scores.masked_fill(mask == 0, -1e9)
+    p_attn = F.softmax(scores, dim=-1)
+    if dropout is not None:
+        p_attn = dropout(p_attn)
+    return torch.matmul(p_attn, value), p_attn
+
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, h, d_model, dropout=0.1):
+        "Take in model size and number of heads."
+        super(MultiHeadedAttention, self).__init__()
+        assert d_model % h == 0
+        # We assume d_v always equals d_k
+        self.d_k = d_model // h
+        self.h = h
+        self.linears = clones(nn.Linear(d_model, d_model), 4)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, query, key, value, mask=None):
+        "Implements Figure 2"
+        if mask is not None:
+            # Same mask applied to all h heads.
+            mask = mask.unsqueeze(1)
+        nbatches = query.size(0)
+
+        # 1) Do all the linear projections in batch from d_model => h x d_k
+        query, key, value = \
+            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
+             for l, x in zip(self.linears, (query, key, value))]
+
+        # 2) Apply attention on all the projected vectors in batch.
+        x, self.attn = attention(query, key, value, mask=mask,
+                                 dropout=self.dropout)
+
+        # 3) "Concat" using a view and apply a final linear.
+        x = x.transpose(1, 2).contiguous() \
+            .view(nbatches, -1, self.h * self.d_k)
+
+        return self.linears[-1](x)

models/transformer/module/noam_opt.py

@@ -0,0 +1,47 @@
+
+class NoamOpt:
+    "Optim wrapper that implements rate."
+
+    def __init__(self, model_size, factor, warmup, optimizer):
+        self.optimizer = optimizer
+        self._step = 0
+        self.warmup = warmup
+        self.factor = factor
+        self.model_size = model_size
+        self._rate = 0
+
+    def step(self):
+        "Update parameters and rate"
+        self._step += 1
+        rate = self.rate()
+        for p in self.optimizer.param_groups:
+            p['lr'] = rate
+        self._rate = rate
+        self.optimizer.step()
+
+    def rate(self, step=None):
+        "Implement `lrate` above"
+        if step is None:
+            step = self._step
+        return self.factor * \
+            (self.model_size ** (-0.5) *
+             min(step ** (-0.5), step * self.warmup ** (-1.5)))
+
+    def save_state_dict(self):
+        return {
+            'inner_optimizer_state_dict': self.optimizer.state_dict(),
+            'step': self._step,
+            'warmup': self.warmup,
+            'factor': self.factor,
+            'model_size': self.model_size,
+            'rate': self._rate
+        }
+
+    def load_state_dict(self, state_dict):
+        self._rate = state_dict['rate']
+        self._step = state_dict['step']
+        self.optimizer.load_state_dict(state_dict['inner_optimizer_state_dict'])
+
+
+
+

models/transformer/module/positional_encoding.py

@@ -0,0 +1,27 @@
+import torch.nn as nn
+import torch
+import math
+from torch.autograd import Variable
+
+
+class PositionalEncoding(nn.Module):
+    "Implement the PE function."
+
+    def __init__(self, d_model, dropout, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0., d_model, 2) *
+                             -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + Variable(self.pe[:, :x.size(1)],
+                         requires_grad=False)
+        return self.dropout(x)

models/transformer/module/positionwise_feedforward.py

@@ -0,0 +1,15 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PositionwiseFeedForward(nn.Module):
+    "Implements FFN equation."
+
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.w_2(self.dropout(F.relu(self.w_1(x))))

models/transformer/module/simpleloss_compute.py

@@ -0,0 +1,23 @@
+
+class SimpleLossCompute:
+    "A simple loss compute and train function."
+
+    def __init__(self, generator, loss_function, opt):
+        self.generator = generator
+        self.loss_function = loss_function
+        self.opt = opt
+
+    def __call__(self, x, y, norm):
+
+        x = self.generator(x)
+
+        loss = self.loss_function(x.contiguous().view(-1, x.size(-1)),
+                                  y.contiguous().view(-1)) / norm
+
+        if self.opt is not None:
+            loss.backward()
+            self.opt.step()
+            self.opt.optimizer.zero_grad()
+       # print("loss from simplelosscompute:",loss)
+       # print("norm from simplelosscompute:",norm)
+        return loss.data * norm

models/transformer/module/subsequent_mask.py

@@ -0,0 +1,10 @@
+import numpy as np
+import torch
+
+
+def subsequent_mask(size):
+    "Mask out subsequent positions."
+    attn_shape = (1, size, size)
+    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
+    return torch.from_numpy(subsequent_mask) == 0
+

preprocess/data_preparation.py

@@ -0,0 +1,86 @@
+import pandas as pd
+import os
+
+from sklearn.model_selection import train_test_split
+
+import utils.file as uf
+import configuration.config_default as cfgd
+import preprocess.property_change_encoder as pce
+
+SEED = 42
+# SPLIT_RATIO = 0.8
+
+
+def get_smiles_list(file_name):
+    """
+    Get smiles list for building vocabulary
+    :param file_name:
+    :return:
+    """
+    pd_data = pd.read_csv(file_name, sep=",")
+
+    print("Read %s file" % file_name)
+    # ravel('K') 是将二维数组展成一维
+    smiles_list = pd.unique(pd_data[['constantSMILES', 'fromVarSMILES', 'toVarSMILES']].values.ravel('K'))
+    print("Number of SMILES in chemical transformations: %d" % len(smiles_list))
+
+    return smiles_list
+
+def split_data(input_transformations_path,SPLIT_RATIO, LOG=None):
+    """
+    Split data into training, validation and test set, write to files
+    :param input_transformations_path:L
+    :return: dataframe
+    """
+    data = pd.read_csv(input_transformations_path, sep=",")
+    if LOG:
+        LOG.info("Read %s file" % input_transformations_path)
+
+    train, test = train_test_split(
+        data, test_size=(1-SPLIT_RATIO)/2, random_state=SEED)
+    train, validation = train_test_split(train, test_size=(1-SPLIT_RATIO)/2, random_state=SEED)
+    if LOG:
+        LOG.info("Train, Validation, Test: %d, %d, %d" % (len(train), len(validation), len(test)))
+
+    parent = uf.get_parent_dir(input_transformations_path)
+    train.to_csv(os.path.join(parent, "train.csv"), index=False)
+    validation.to_csv(os.path.join(parent, "validation.csv"), index=False)
+    test.to_csv(os.path.join(parent, "test.csv"), index=False)
+
+    return train, validation, test
+
+def save_df_property_encoded(file_name, property_change_encoder, LOG=None):
+    data = pd.read_csv(file_name, sep=",")
+    for property_name in cfgd.PROPERTIES:
+        if property_name == 'pki':
+            encoder, start_map_interval = property_change_encoder[property_name]
+            data['Delta_{}'.format(property_name)] = \
+                data['Delta_{}'.format(property_name)].apply(lambda x:
+                                                                 pce.value_in_interval(x, start_map_interval), encoder)
+                    
+        elif property_name == 'qed':
+              encoder, start_map_interval = property_change_encoder[property_name]
+              data['Delta_{}'.format(property_name)] = \
+                  data['Delta_{}'.format(property_name)].apply(lambda x:
+                                                                 pce.value_in_interval(x, start_map_interval), encoder)
+                      
+        elif property_name == 'sa':
+              encoder, start_map_interval = property_change_encoder[property_name]
+              data['Delta_{}'.format(property_name)] = \
+                  data['Delta_{}'.format(property_name)].apply(lambda x:
+                                                                 pce.value_in_interval(x, start_map_interval), encoder)
+
+    output_file = file_name.split('.csv')[0] + '_encoded.csv'
+    LOG.info("Saving encoded property change to file: {}".format(output_file))
+    data.to_csv(output_file, index=False)
+    return output_file
+
+def prop_change(source, target, threshold):
+    if source <= threshold and target > threshold:
+        return "low->high"
+    elif source > threshold and target <= threshold:
+        return "high->low"
+    elif source <= threshold and target <= threshold:
+        return "no_change"
+    elif source > threshold and target > threshold:
+        return "no_change"

preprocess/property_change_encoder.py

@@ -0,0 +1,73 @@
+import numpy as np
+import pandas as pd
+
+import configuration.config_default as cfgd
+
+STEP_pki = 1
+
+
+def encode_property_change(input_data_path, LOG=None):
+    property_change_encoder = {}
+    for property_name in cfgd.PROPERTIES:
+        if property_name == 'pki':
+            # intervals ['(3,4]', ...] 形式
+            # start_map_interval <number, interval> 用来找区间
+            intervals, start_map_interval = build_intervals(input_data_path, step=STEP_pki, LOG=LOG)
+        
+        if property_name == 'pki':
+            property_change_encoder[property_name] = intervals, start_map_interval
+        
+    return property_change_encoder
+
+
+def value_in_interval(value, start_map_interval):
+    start_vals = sorted(list(start_map_interval.keys()))
+    return start_map_interval[start_vals[np.searchsorted(start_vals, value, side='right') - 1]]
+
+
+def interval_to_onehot(interval, encoder):
+    return encoder.transform([interval]).toarray()[0]
+
+
+def build_intervals(input_transformations_path, step=STEP_pki, LOG=None):
+    df = pd.read_csv(input_transformations_path)
+    # df=input_transformations_path
+    delta_pki = df['Delta_pki'].tolist()
+    min_val, max_val = min(delta_pki), max(delta_pki)
+    if LOG:
+         LOG.info("pki min and max: {}, {}".format(min_val, max_val))
+    
+    start_map_interval = {}
+    interval_str = '({}, {}]'.format(round(-step/2, 2), round(step/2, 2))
+    intervals = [interval_str]
+    start_map_interval[-step/2] = interval_str
+
+    smallStep=step
+    bigStep=2
+    positives = step/2
+    while positives < 10:
+        if positives>2:
+            step=bigStep
+        interval_str = '({}, {}]'.format(round(positives, 2), round(positives+step, 2))
+        intervals.append(interval_str)
+        start_map_interval[positives] = interval_str
+        positives += step
+        
+        
+    interval_str = '({}, inf]'.format(round(positives, 2))
+    intervals.append(interval_str)
+    start_map_interval[positives] = interval_str
+
+    step=smallStep
+    negatives = -step/2
+    while negatives > min_val:
+        interval_str = '({}, {}]'.format(round(negatives-step, 2), round(negatives, 2))
+        intervals.append(interval_str)
+        negatives -= step
+        start_map_interval[negatives] = interval_str
+    interval_str = '(-inf, {}]'.format(round(negatives, 2))
+    intervals.append(interval_str)
+    start_map_interval[float('-inf')] = interval_str
+
+    return intervals, start_map_interval
+

preprocess/vocabulary.py

@@ -0,0 +1,145 @@
+# coding=utf-8
+
+"""
+Vocabulary helper class
+"""
+
+import re
+import numpy as np
+
+class Vocabulary:
+    """Stores the tokens and their conversion to one-hot vectors."""
+
+    def __init__(self, tokens=None, starting_id=0):
+        self._tokens = {}
+        self._current_id = starting_id
+
+        if tokens:
+            for token, idx in tokens.items():
+                self._add(token, idx)
+                self._current_id = max(self._current_id, idx + 1)
+
+    def __getitem__(self, token_or_id):
+        return self._tokens[token_or_id]
+
+    def add(self, token):
+        """Adds a token."""
+        if not isinstance(token, str):
+            raise TypeError("Token is not a string")
+        if token in self:
+            # raise ValueError("Token already present in the vocabulary")
+            print(f'=== Token "{token}"already present in the vocabulary')
+            return
+        self._add(token, self._current_id)
+        self._current_id += 1
+        return self._current_id - 1
+
+    def update(self, tokens):
+        """Adds many tokens."""
+        return [self.add(token) for token in tokens]
+
+    def __delitem__(self, token_or_id):
+        other_val = self._tokens[token_or_id]
+        del self._tokens[other_val]
+        del self._tokens[token_or_id]
+
+    def __contains__(self, token_or_id):
+        return token_or_id in self._tokens
+
+    def __eq__(self, other_vocabulary):
+        return self._tokens == other_vocabulary._tokens
+
+    def __len__(self):
+        return len(self._tokens) // 2
+
+    def encode(self, tokens):
+        """Encodes a list of tokens, encoding them in 1-hot encoded vectors."""
+        ohe_vect = np.zeros(len(tokens), dtype=np.float32)
+        for i, token in enumerate(tokens):
+            try:
+                ohe_vect[i] = self._tokens[token]
+            except KeyError:
+                ohe_vect[i] = self._tokens["default_key"]
+        return ohe_vect
+
+    def decode(self, ohe_vect):
+        """Decodes a one-hot encoded vector matrix to a list of tokens."""
+        tokens = []
+        for ohv in ohe_vect:
+            try:
+                tokens.append(self[ohv])
+            except KeyError:
+                tokens.append("default_key")
+        return tokens
+
+    def _add(self, token, idx):
+        if idx not in self._tokens:
+            self._tokens[token] = idx
+            self._tokens[idx] = token
+        else:
+            raise ValueError("IDX already present in vocabulary")
+
+    def tokens(self):
+        """Returns the tokens from the vocabulary"""
+        return [t for t in self._tokens if isinstance(t, str)]
+
+    def word2idx(self):
+        return {k: self._tokens[k] for k in self._tokens if isinstance(k, str)}
+
+
+class SMILESTokenizer:
+    """Deals with the tokenization and untokenization of SMILES."""
+
+    REGEXPS = {
+        "brackets": re.compile(r"(\[[^\]]*\])"),
+        "2_ring_nums": re.compile(r"(%\d{2})"),
+        "brcl": re.compile(r"(Br|Cl)")
+    }
+    REGEXP_ORDER = ["brackets", "2_ring_nums", "brcl"]
+
+    def tokenize(self, data, with_begin_and_end=True):
+        """Tokenizes a SMILES string."""
+        def split_by(data, regexps):
+            if not regexps:
+                return list(data)
+            regexp = self.REGEXPS[regexps[0]]
+            splitted = regexp.split(data)
+            tokens = []
+            for i, split in enumerate(splitted):
+                if i % 2 == 0:
+                    tokens += split_by(split, regexps[1:])
+                else:
+                    tokens.append(split)
+            return tokens
+
+        tokens = split_by(data, self.REGEXP_ORDER)
+        if with_begin_and_end:
+            tokens = ["^"] + tokens + ["$"]
+        return tokens
+
+    def untokenize(self, tokens):
+        """Untokenizes a SMILES string."""
+        smi = ""
+        for token in tokens:
+            if token == "$":
+                break
+            if token != "^":
+                smi += token
+        return smi
+
+
+def create_vocabulary(smiles_list, tokenizer, property_condition=None):
+    """Creates a vocabulary for the SMILES syntax."""
+    tokens = set()
+    for smi in smiles_list:
+        tokens.update(tokenizer.tokenize(smi, with_begin_and_end=False))
+
+    vocabulary = Vocabulary()
+    vocabulary.update(["*", "^", "$"] + sorted(tokens))  # pad=0, start=1, end=2
+    if property_condition is not None:
+        vocabulary.update(property_condition)
+    # for random smiles
+    if "8" not in vocabulary.tokens():
+        vocabulary.update(["8"])
+
+    return vocabulary