add pre-train

config/config_layer_6_mol.json

+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 595,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "type_vocab_size": 2,
+  "vocab_size": 40235
+}
\ No newline at end of file

dataset.py

@@ -8,7 +8,9 @@ from sklearn.preprocessing import OneHotEncoder
 from subword_nmt.apply_bpe import BPE
 import codecs
 from tqdm import tqdm
-import csv
+import math
+import random
+from torch.nn.utils.rnn import pad_sequence
 
 
 # vocab_path = './ESPF/protein_codes_uniprot.txt'
@@ -107,6 +109,24 @@ def seq2emb_encoder(input_seq, max_len, vocab):
     return ids, input_mask
 
 
+def seq2emb_encoder_simple(input_seq, max_len, vocab):
+    try:
+        ids = np.asarray([vocab[i] for i in input_seq])
+    except:
+        ids = np.array([0])
+
+    # l = len(ids)
+    #
+    # if l < max_len:
+    #     ids = np.pad(ids, (0, max_len - l), 'constant', constant_values=0)
+    #     input_mask = np.array(([1] * l) + ([0] * (max_len - l)))
+    # else:
+    #     ids = ids[:max_len]
+    #     input_mask = np.array([1] * max_len)
+
+    return ids
+
+
 class Data_Encoder(data.Dataset):
     def __init__(self, train_file, tokenizer_config):
         'Initialization'
@@ -196,6 +216,62 @@ class Data_Encoder_mol(data.Dataset):
         # return len(d), len(p)
 
 
+class Data_Encoder_LM(data.Dataset):
+    def __init__(self, train_file, tokenizer_config):
+        'Initialization'
+        # load data
+        # with open(train_file["sps"], 'r') as f:
+        #     self.sps = f.readlines()
+        with open(train_file['seq'], 'r') as f:
+            self.seq = f.readlines()
+        with open(train_file["smile"], 'r') as f:
+            self.smile = f.readlines()
+        with open(train_file["affinity"], 'r') as f:
+            self.affinity = f.readlines()
+        # define tokenizer
+        self.begin_id = tokenizer_config["begin_id"]
+        self.sep_id = tokenizer_config["separate_id"]
+        self.max_len = tokenizer_config["max_len"]
+        self.vocab = load_vocab(tokenizer_config["vocab_file"])
+
+        bpe_codes_drug = codecs.open(tokenizer_config["vocab_pair"])
+        self.dbpe = BPE(bpe_codes_drug, merges=-1, separator='')
+
+        bpe_codes_prot = codecs.open(tokenizer_config["vocab_pair_p"])
+        self.pbpe = BPE(bpe_codes_prot, merges=-1, separator='')
+
+    def __len__(self):
+        'Denotes the total number of samples'
+        return len(self.smile)
+
+
+    def __getitem__(self, index):
+        'Generates one sample of data'
+        # Select sample
+        # Load data and get label
+        # tokenization
+        d = self.dbpe.process_line(self.smile[index].strip()).split()
+        p = self.pbpe.process_line(self.seq[index].strip()).split()
+        # mask_d, mask_d_posi = self.random_mask(d)
+        # mask_p, mask_p_posi = self.random_mask(p)
+        y = np.float64(self.affinity[index].strip())
+        #
+        # input_seq = [self.begin_id] + mask_d + [self.sep_id] + mask_p + [self.sep_id]
+        # mask_posi = np.concatenate((np.zeros(1), mask_d_posi, np.zeros(1), mask_p_posi, np.zeros(1)))
+        # token_type_ids = np.concatenate((np.zeros((len(d) + 2), dtype=np.int), np.ones((len(p) + 1), dtype=np.int)))
+        # if len(input_seq) > self.max_len:
+        #     input_seq = input_seq[:self.max_len-1] + [self.sep_id]
+        #     token_type_ids = token_type_ids[:self.max_len]
+        #     mask_posi = mask_posi[:self.max_len]
+        # else:
+        #     mask_posi = np.pad(mask_posi, (0, self.max_len - len(input_seq)), 'constant', constant_values=0)
+        #     token_type_ids = np.pad(token_type_ids, (0, self.max_len - len(input_seq)), 'constant', constant_values=0)
+        # input, input_mask = seq2emb_encoder(input_seq, self.max_len, self.vocab)
+        # return torch.from_numpy(input).long(), torch.from_numpy(token_type_ids).long(), torch.from_numpy(input_mask).long(), y, torch.from_numpy(mask_posi).long()
+        return " ".join(d), " ".join(p), y
+        # return len(d), len(p)
+
+
 def get_task(task_name):
     tokenizer_config = {"vocab_file": './config/vocab.txt',
                         "vocab_pair": './config/drug_codes_chembl.txt',
@@ -244,7 +320,7 @@ def get_task(task_name):
 
         return df, tokenizer_config
 
-    elif task_name.lower() == 'train_mol':
+    elif task_name.lower() in ['train_mol', "pre-train"]:
         df_train = {"sps": './data/train/train_sps',
                     'seq': './data/train/train_protein_seq',
                     "smile": './data/train/train_smile',
@@ -277,6 +353,63 @@ def get_task(task_name):
 
         return df_test, tokenizer_config
 
+def random_mask(input_seq, mask_proportion=0.15):
+    input = [i.split() for i in input_seq]
+    mask_len = [math.ceil(len(i)*mask_proportion) for i in input]
+    # mask_posi = np.arange(len(input_seq))
+    # mask_token_posi = random.sample(mask_posi, mask_len)
+    mask_token_posi = [np.random.choice(len(i), j) for i, j in zip(input, mask_len)]
+    # mask_vec = np.zeros(len(input_seq))
+    for i, posi in enumerate(mask_token_posi):
+        for j in posi:
+            choice = random.random()
+            if choice < 0.8:
+                input[i][j] = "[MASK]"
+            # mask_vec[i] = 1
+    # elif choice >= 0.8 and choice < 0.9:
+
+    return input
+
+
+class Tokenizer(object):
+    def __init__(self, tokenizer_config):
+        self.begin_id = tokenizer_config["begin_id"]
+        self.sep_id = tokenizer_config["separate_id"]
+        self.max_len = tokenizer_config["max_len"]
+        self.vocab = load_vocab(tokenizer_config["vocab_file"])
+
+    def convert_token_to_ids(self, d, p):
+        mask_d = random_mask(d)
+        mask_p = random_mask(p)
+        input_seq = [[self.begin_id] + i + [self.sep_id] + j + [self.sep_id] for i, j in zip(mask_d, mask_p)]
+        input_seq_ori = [[self.begin_id] + i.split() + [self.sep_id] + j.split() + [self.sep_id] for i, j in zip(d, p)]
+        # mask_posi = np.concatenate((np.zeros(1), mask_d_posi, np.zeros(1), mask_p_posi, np.zeros(1)))
+        # token_type_ids = [[np.concatenate((np.zeros((len(d) + 2), dtype=np.int), np.ones((len(p) + 1), dtype=np.int)))] for d, p in zip(mask_d, mask_p)]
+        for i, seq in enumerate(input_seq):
+            if len(seq) > self.max_len:
+                input_seq[i] = seq[:self.max_len-1] + [self.sep_id]
+                input_seq_ori[i] = seq[:self.max_len-1] + [self.sep_id]
+            # token_type_ids = token_type_ids[:self.max_len]
+            # mask_posi = mask_posi[:self.max_len]
+        # else:
+            # mask_posi = np.pad(mask_posi, (0, self.max_len - len(input_seq)), 'constant', constant_values=0)
+            # token_type_ids = np.pad(token_type_ids, (0, self.max_len - len(input_seq)), 'constant', constant_values=0)
+        all_seq = []
+        all_seq_ori = []
+        # all_mask = []
+        for seq, ori in zip(input_seq, input_seq_ori):
+            input = seq2emb_encoder_simple(seq, self.max_len, self.vocab)
+            input_ori = seq2emb_encoder_simple(ori, self.max_len, self.vocab)
+            all_seq.append(torch.from_numpy(input).long())
+            all_seq_ori.append(torch.from_numpy(input_ori).long())
+
+        input = pad_sequence(all_seq, batch_first=True)
+        input_ori = pad_sequence(all_seq_ori, batch_first=True)
+        input_mask = input != 0
+        # input_mask = pad_sequence(all_mask)
+        # return torch.from_numpy(input).long(), torch.from_numpy(input_mask).long(), torch.from_numpy(token_type_ids).long()
+        return input, input_mask, input_ori
+
 
 
 if __name__ == "__main__":

eval.py

@@ -23,7 +23,7 @@ def eval(pred_path, label_path):
         label = [float(i.strip()) for i in label]
     remse, r_mat = eval_result(pred, label)
     r = r_mat[0, 1]
-    save_path = pred_path.replace('test.txt', 'eval_results')
+    save_path = pred_path.replace('test_mol.txt', 'eval_results')
     with open(save_path, 'w') as f:
         f.write('RMSE : {} ; Pearson Correlation Coefficient : {}'.format(remse, r))
     print('RMSE : {} ; Pearson Correlation Coefficient : {}'.format(remse, r))
@@ -34,6 +34,6 @@ if __name__ == '__main__':
     #     pred_dir = f.readline()
     #     pred_dir = pred_dir.split()[5].split('/')[-1]
     # pred_result = './predict/{}/test.txt'.format(pred_dir)
-    pred_result = './predict/lr-1e-5-batch-32-e-10-layer3-0503-z-1-step-82370/test.txt'
-    test_label_path = './data/test_ic50'
+    pred_result = './predict/train_mol_lr-1e-5-batch-32-e-30-layer3-0609-e-26/test_mol.txt'
+    test_label_path = './data/test/test_ic50'
     eval(pred_result, test_label_path)

modeling_bert.py

@@ -2017,3 +2017,181 @@ class BertAffinityModel(BertPreTrainedModel):
         #     cross_attentions=encoder_outputs.cross_attentions,
         # )
         return pred_affinity
+
+
+class BertAffinityModel_MaskLM(BertPreTrainedModel):
+    """
+
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+
+    To behave as an decoder the model needs to be initialized with the :obj:`is_decoder` argument of the configuration
+    set to :obj:`True`. To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
+    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
+    input to the forward pass.
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.mlp = Multilayer_perceptron(config)
+        # self.pooler = BertPooler(config) if add_pooling_layer else None
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPoolingAndCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.config.is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # past_key_values_length
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        if attention_mask is None:
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values_length,
+        )
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        logits = self.lm_head(sequence_output)
+        # pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        # if not return_dict:
+            # return (sequence_output, pooled_output) + encoder_outputs[1:]
+        # print(sequence_output.size())
+        # bert_pred = sequence_output[:,0,:]
+        # pred_affinity = self.mlp.forward(bert_pred)
+
+        # return BaseModelOutputWithPoolingAndCrossAttentions(
+        #     last_hidden_state=sequence_output,
+        #     pooler_output=pooled_output,
+        #     past_key_values=encoder_outputs.past_key_values,
+        #     hidden_states=encoder_outputs.hidden_states,
+        #     attentions=encoder_outputs.attentions,
+        #     cross_attentions=encoder_outputs.cross_attentions,
+        # )
+        return logits
+

run_interaction.py

@@ -41,12 +41,12 @@ def train(args, model, dataset):
         for i, (input, token_type_ids, input_mask, affinity) in enumerate(data_generator):
             # use cuda
             # input model
-            if torch.cuda.is_available():
-                pred_affinity = model(input_ids=input.cuda(), token_type_ids=token_type_ids.cuda(), attention_mask=input_mask.cuda())
-                loss = loss_fct(pred_affinity, affinity.cuda().unsqueeze(-1))
-            else:
-                pred_affinity = model(input_ids=input, token_type_ids=token_type_ids, attention_mask=input_mask)
-                loss = loss_fct(pred_affinity, affinity.unsqueeze(-1))
+            # if torch.cuda.is_available():
+            pred_affinity = model(input_ids=input.cuda(), token_type_ids=token_type_ids.cuda(), attention_mask=input_mask.cuda())
+            loss = loss_fct(pred_affinity, affinity.cuda().unsqueeze(-1))
+            # else:
+            #     pred_affinity = model(input_ids=input, token_type_ids=token_type_ids, attention_mask=input_mask)
+            #     loss = loss_fct(pred_affinity, affinity.unsqueeze(-1))
             step += 1
             writer.add_scalar('loss', loss, global_step=step)
             # Update gradient
@@ -74,10 +74,10 @@ def test(args, model, dataset):
     data_generator = DataLoader(dataset, **data_loder_para)
 
     with torch.no_grad():
-        if torch.cuda.is_available():
-            model.load_state_dict(torch.load(args.init), strict=True)
-        else:
-            model.load_state_dict(torch.load(args.init, map_location=torch.device('cpu')), strict=True)
+        # if torch.cuda.is_available():
+        model.load_state_dict(torch.load(args.init), strict=True)
+        # else:
+        #     model.load_state_dict(torch.load(args.init, map_location=torch.device('cpu')), strict=True)
         model.eval()
         if not os.path.exists(args.output):
             os.mkdir(args.output)
@@ -85,12 +85,12 @@ def test(args, model, dataset):
         print('begin predicting')
         with open(result, 'w') as f:
             for i, (input, token_type_ids, input_mask, affinity) in enumerate(tqdm(data_generator)):
-                if torch.cuda.is_available():
-                    model.cuda()
-                    pred_affinity = model(input_ids=input.cuda(), token_type_ids=token_type_ids.cuda(),
+                # if torch.cuda.is_available():
+                model.cuda()
+                pred_affinity = model(input_ids=input.cuda(), token_type_ids=token_type_ids.cuda(),
                                           attention_mask=input_mask.cuda())
-                else:
-                    pred_affinity = model(input_ids=input, token_type_ids=token_type_ids, attention_mask=input_mask)
+                # else:
+                #     pred_affinity = model(input_ids=input, token_type_ids=token_type_ids, attention_mask=input_mask)
                 pred_affinity = pred_affinity.cpu().numpy().squeeze(-1)
                 for res in pred_affinity:
                     f.write(str(res) + '\n')
@@ -153,12 +153,12 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     # local test
-
-    # args.task = 'train_mol'
-    # args.savedir = 'local_test_train'
-    # args.epochs = 10
-    # args.lr = 1e-5
-    # args.config = './config/config_layer_3_mol.json'
+    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+    args.task = 'train_mol'
+    args.savedir = 'local_test_train'
+    args.epochs = 10
+    args.lr = 1e-5
+    args.config = './config/config_layer_3_mol.json'
 
 
 

run_pretraining.py

+from argparse import ArgumentParser
+from dataset import Data_Encoder, get_task, Data_Encoder_mol, Data_Encoder_LM, Tokenizer
+import torch
+from torch.utils.data import DataLoader
+from configuration_bert import BertConfig
+from modeling_bert import BertAffinityModel, BertAffinityModel_MaskLM
+from torch.utils.tensorboard import SummaryWriter
+import os
+from tqdm import tqdm
+# torch.set_default_tensor_type(torch.DoubleTensor)
+
+
+
+
+def train(args, model, dataset, tokenizer):
+    data_loder_para = {'batch_size': args.batch_size,
+                       'shuffle': True,
+                       'num_workers': args.workers,
+                       }
+    data_generator = DataLoader(dataset, **data_loder_para)
+
+
+    model.train()
+    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
+    # loss_fct = torch.nn.MSELoss()
+    loss_fct = torch.nn.CrossEntropyLoss()
+    writer = SummaryWriter('./log/' + args.savedir)
+    num_step = args.epochs * len(data_generator)
+    step = 0
+    save_step = num_step // 10
+    # detect GPU
+    if torch.cuda.is_available():
+        model.cuda()
+    # print(model)
+    print('epoch num : {}'.format(args.epochs))
+    print('step num : {}'.format(num_step))
+    print('batch size : {}'.format(args.batch_size))
+    print('learning rate : {}'.format(args.lr))
+    print('begin training')
+    # training
+    for epoch in range(args.epochs):
+        for i, (drug, protein, affinity) in enumerate(data_generator):
+            input, input_mask, input_ori = tokenizer.convert_token_to_ids(drug, protein)
+            # pred_affinity = model(input_ids=input.cuda(), token_type_ids=token_type_ids.cuda(), attention_mask=input_mask.cuda())
+            logits = model(input_ids=input.cuda(), attention_mask=input_mask.cuda())
+
+            # loss = 0
+            pred_logits = logits[input == 1]
+            label = input_ori[input == 1]
+            loss = loss_fct(pred_logits, label.cuda())
+            # else:
+            #     pred_affinity = model(input_ids=input, token_type_ids=token_type_ids, attention_mask=input_mask)
+            #     loss = loss_fct(pred_affinity, affinity.unsqueeze(-1))
+            step += 1
+            writer.add_scalar('loss', loss, global_step=step)
+            # Update gradient
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+
+            #                 if (i % 100 == 0):
+            print('Training at Epoch ' + str(epoch + 1) + ' step ' + str(step) + ' with loss ' + str(
+                loss.cpu().detach().numpy()))
+            # save
+            if epoch >= 1 and step % save_step == 0:
+                save_path = './model/' + args.savedir + '/'
+                if not os.path.exists(save_path):
+                    os.mkdir(save_path)
+                torch.save(model.state_dict(), save_path + 'epoch-{}-step-{}-loss-{}.pth'.format(epoch, step, loss))
+    print('training  over')
+    writer.close()
+
+def test(args, model, dataset):
+    data_loder_para = {'batch_size': args.batch_size,
+                       'shuffle': False,
+                       'num_workers': args.workers,
+                       }
+    data_generator = DataLoader(dataset, **data_loder_para)
+
+    with torch.no_grad():
+        # if torch.cuda.is_available():
+        model.load_state_dict(torch.load(args.init), strict=True)
+        # else:
+        #     model.load_state_dict(torch.load(args.init, map_location=torch.device('cpu')), strict=True)
+        model.eval()
+        if not os.path.exists(args.output):
+            os.mkdir(args.output)
+        result = args.output + '/' + '{}.txt'.format(args.task)
+        print('begin predicting')
+        with open(result, 'w') as f:
+            for i, (input, token_type_ids, input_mask, affinity) in enumerate(tqdm(data_generator)):
+                # if torch.cuda.is_available():
+                model.cuda()
+                pred_affinity = model(input_ids=input.cuda(), token_type_ids=token_type_ids.cuda(),
+                                          attention_mask=input_mask.cuda())
+                # else:
+                #     pred_affinity = model(input_ids=input, token_type_ids=token_type_ids, attention_mask=input_mask)
+                pred_affinity = pred_affinity.cpu().numpy().squeeze(-1)
+                for res in pred_affinity:
+                    f.write(str(res) + '\n')
+
+    # if args.do_eval:
+    #     os.system('python eval.py')
+
+
+def main(args):
+    # load data
+    data_file, tokenizer_config = get_task(args.task)
+    # dataset = Data_Encoder(data_file, tokenizer_config)
+    dataset = Data_Encoder_LM(data_file, tokenizer_config)
+    tokenizer = Tokenizer(tokenizer_config)
+    # creat model
+    print('------------------creat model---------------------------')
+    config = BertConfig.from_pretrained(args.config)
+    model = BertAffinityModel_MaskLM(config)
+
+    # if torch.cuda.device_count() > 1:
+    #     print("Let's use", torch.cuda.device_count(), "GPUs!")
+    #     model = torch.nn.DataParallel(model, dim=0)
+
+    print('model name : BertAffinity')
+    print('task name : {}'.format(args.task))
+
+    if args.task in ['pre-train']:
+        train(args, model, dataset, tokenizer)
+
+    # elif args.task in ['test', 'test_mol']:
+    #     test(args, model, dataset)
+
+
+
+
+if __name__ == '__main__':
+    # get parameter
+    parser = ArgumentParser(description='BertAffinity')
+    parser.add_argument('-b', '--batch-size', default=8, type=int,
+                        metavar='N',
+                        help='mini-batch size (default: 16), this is the total '
+                             'batch size of all GPUs on the current node when '
+                             'using Data Parallel or Distributed Data Parallel')
+    parser.add_argument('-j', '--workers', default=0, type=int, metavar='N',
+                        help='number of data loading workers (default: 0)')
+    parser.add_argument('--epochs', default=50, type=int, metavar='N',
+                        help='number of total epochs to run')
+    parser.add_argument('--task', default='train', type=str, metavar='TASK',
+                        help='Task name. Could be train, test, channel, ER, GPCR, kinase or else.')
+    parser.add_argument('--lr', '--learning-rate', default=1e-5, type=float,
+                        metavar='LR', help='initial learning rate', dest='lr')
+    parser.add_argument('--config', default='./config/config.json', type=str, help='model config file path')
+    # parser.add_argument('--log', default='training_log', type=str, help='training log')
+    parser.add_argument('--savedir', default='train', type=str, help='log and model save path')
+    # parser.add_argument('--device', default='0', type=str, help='name of GPU')
+    parser.add_argument('--init', default='model', type=str, help='init checkpoint')
+    parser.add_argument('--output', default='predict', type=str, help='result save path')
+    # parser.add_argument('--shuffle', default=True, type=str, help='shuffle data')
+    parser.add_argument('--do_eval', default=False, type=bool, help='do eval')
+
+    args = parser.parse_args()
+
+    # local test
+    os.environ["CUDA_VISIBLE_DEVICES"] = "5"
+    args.task = 'pre-train'
+    args.savedir = 'mask-LM-quick'
+    args.epochs = 30
+    args.lr = 1e-5
+    args.config = './config/config_layer_6_mol.json'
+
+
+
+    # args.task = 'test'
+    # args.init = './model/lr-1e-5-batch-32-e-10-layer6-0428/epoch-8-step-74133-loss-0.6730387237803921.pth'
+    # args.output = './predict/test'
+    # args.config = './config/config_layer_6.json'
+    main(args)
+