MLM_SFP.py

import random
import time
import numpy as np
import torch 
import torch.optim as optim
from torchvision import transforms, datasets
import copy
from Bio import SeqIO
import argparse
from utils.bert import get_config, BertModel, set_learned_params, BertForMaskedLM, visualize_attention, show_base_PCA, fix_params
from module import Train_Module
from dataload import DATA, MyDataset 
import datetime
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.metrics.cluster import adjusted_rand_score
import os
import time
from sklearn.metrics import normalized_mutual_info_score, adjusted_rand_score, completeness_score, homogeneity_score
import torch.nn.functional as F
from sklearn.cluster import MiniBatchKMeans, KMeans, AgglomerativeClustering, SpectralClustering 
import itertools  

import alignment_C as Aln_C

random.seed(10)
torch.manual_seed(1234)
np.random.seed(1234)

parser = argparse.ArgumentParser(description='RNABERT')
parser.add_argument('--mag',  type=int, default=1,
                    help='enumerate')
parser.add_argument('--epoch', '-e', type=int, default=200,
                    help='Number of sweeps over the dataset to train')
parser.add_argument('--batch', '-b', type=int, default=20,
                    help='Number of batch size')
parser.add_argument('--maskrate', '-m', type=float, default=0.0,
                    help='mask rate')
parser.add_argument('--pretraining', '-pre', type=str, help='use pretrained weight')
parser.add_argument('--outputweight', type=str, help='output path for weights')
parser.add_argument('--algorithm', type=str, default="global", help='algorithm method')
parser.add_argument('--data_mlm', '-d', type=str, nargs='*', help='data for mlm training')
parser.add_argument('--data_mul', type=str, nargs='*', help='data for mul training')
parser.add_argument('--data_alignment', type=str, nargs='*', help='data for alignment test')
parser.add_argument('--data_clustering', type=str, nargs='*', help='data for clustering test')
parser.add_argument('--data_showbase', type=str, nargs='*', help='data for base embedding')
parser.add_argument('--data_embedding', type=str, nargs='*', help='data for base embedding')
parser.add_argument('--embedding_output', type=str, nargs='*', help='output file for base embedding')
parser.add_argument('--show_aln', action='store_true')

args = parser.parse_args()
batch_size = args.batch
current_time = datetime.datetime.now()

print("start...")
class TRAIN:
    """The class for controlling the training process of SFP"""
    def __init__(self, config):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.module = Train_Module(config)
    
    def model_device(self, model):
        print("device: ", self.device)
        print('-----start-------')
        model.to(self.device)
        if self.device == 'cuda':
            model = torch.nn.DataParallel(model) # make parallel
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        return model

    def train_MLM_SFP(self, model, optimizer, dl_MLM_SFP, num_epochs, task_type):
        for epoch in range(num_epochs):
            model.train()
            epoch_mlm_loss = 0.0
            epoch_ssl_loss = 0.0
            epoch_mlm_correct = 0.0
            epoch_ssl_correct = 0.0
            epoch_sfp_loss=0.0
            epoch_sfp_correct = 0.0
            epoch_mul_loss = 0.0

            iteration = 1
            t_epoch_start = time.time()
            t_iter_start = time.time()
            data_num = 0
            for batch in dl_MLM_SFP:
                optimizer.zero_grad()
                if task_type == "MLM" or task_type == "SFP":
                    low_seq_0, masked_seq_0, family_0, seq_len_0, low_seq_1, masked_seq_1, family_1, seq_len_1 = batch
                elif task_type == "MUL":
                    low_seq_0, masked_seq_0, family_0, seq_len_0, low_seq_1, masked_seq_1, family_1, seq_len_1, common_index_0, common_index_1 = batch

                masked_seq_0 = masked_seq_0.to(self.device)
                low_seq_0 = low_seq_0.to(self.device)
                masked_seq_1 = masked_seq_1.to(self.device)
                low_seq_1 = low_seq_1.to(self.device)

                masked_seq = torch.cat((masked_seq_0, masked_seq_1), axis=0) 
                prediction_scores, prediction_scores_ss, encoded_layers =  model(masked_seq)
                prediction_scores0, prediction_scores1 = torch.split(prediction_scores, int(prediction_scores.shape[0]/2))
                prediction_scores_ss0, prediction_scores_ss1 = torch.split(prediction_scores_ss, int(prediction_scores_ss.shape[0]/2))
                encoded_layers0, encoded_layers1 = torch.split(encoded_layers, int(encoded_layers.shape[0]/2))

                loss = 0
                # MLM LOSS
                mlm_loss_0, mlm_correct_0 = self.module.train_MLM(low_seq_0, masked_seq_0, prediction_scores0)
                mlm_loss_1, mlm_correct_1 = self.module.train_MLM(low_seq_1, masked_seq_1, prediction_scores1)
                mlm_loss = (mlm_loss_0 + mlm_loss_1)/2
                mlm_loss = torch.tensor(0.0) if  torch.isnan(mlm_loss) else mlm_loss 
                mlm_correct = (mlm_correct_0 + mlm_correct_1)/2
                epoch_mlm_loss += mlm_loss.item() * batch_size
                epoch_mlm_correct += mlm_correct
                if task_type == "MLM":    
                    loss += mlm_loss

                # SFP LOSS
                if task_type == "SFP":    
                    z0_list, z1_list =  self.module.em(encoded_layers0, seq_len_0), self.module.em(encoded_layers1, seq_len_1)
                    sfp_loss, sfp_correct = self.module.train_SFP(low_seq_0, seq_len_0, low_seq_1, seq_len_1, family_0, family_1, z0_list, z1_list)
                    sfp_loss = torch.tensor(0.0) if  torch.isnan(sfp_loss) else sfp_loss 
                    epoch_sfp_loss += sfp_loss.item()* batch_size
                    epoch_sfp_correct += sfp_correct
                    loss += sfp_loss

                # MULTIPLE LOSS
                if task_type == "MUL":
                    common_index_0 = common_index_0.to(self.device)
                    common_index_1 = common_index_1.to(self.device)
                    z0_list, z1_list =  self.module.em(encoded_layers0, seq_len_0), self.module.em(encoded_layers1, seq_len_1)
                    mul_loss = self.module.train_MUL(z0_list, z1_list, common_index_0, common_index_1, seq_len_0, seq_len_1)
                    mul_loss = torch.tensor(0.0) if  torch.isnan(mul_loss) else mul_loss 
                    epoch_mul_loss += mul_loss.item()
                    loss +=  mul_loss

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                torch.cuda.empty_cache()

            t_epoch_finish = time.time()
            epoch_mlm_loss = epoch_mlm_loss / len(dl_MLM_SFP.dataset)
            epoch_mlm_correct = epoch_mlm_correct / len(dl_MLM_SFP)
            epoch_sfp_loss = epoch_sfp_loss  / len(dl_MLM_SFP.dataset)
            epoch_sfp_correct = epoch_sfp_correct / len(dl_MLM_SFP.dataset)
            epoch_mul_loss = epoch_mul_loss
            print('Epoch {}/{} | MLM Loss: {:.4f} MLM Acc: {:.4f}| SFP Loss: {:.4f} SFP Acc: {:.4f}| MUL Loss: {:.4f}| time: {:.4f} sec.'.format(epoch+1, num_epochs,
                                                                        epoch_mlm_loss, epoch_mlm_correct, epoch_sfp_loss, epoch_sfp_correct, epoch_mul_loss, time.time() - t_epoch_start))
            t_epoch_start = time.time()
        if args.outputweight:
            torch.save(model.state_dict(), args.outputweight + '{0:%m_%d_%H_%M}'.format(current_time))
            torch.save(model.state_dict(), args.outputweight)
        return model

    # make feature vector 
    def make_feature(self, model, dataloader, seqs):
        model.eval()
        torch.backends.cudnn.benchmark = True
        batch_size = dataloader.batch_size
        encoding = []
        for batch in dataloader:
            data, label, seq_len= batch
            inputs = data.to(self.device)
            prediction_scores, prediction_scores_ss, encoded_layers =  model(inputs)
            encoding.append(encoded_layers.cpu().detach().numpy())
        encoding = np.concatenate(encoding, 0)

        embedding = []
        for e, seq in zip(encoding, seqs):
            embedding.append(e[:len(seq)].tolist())

        return embedding 

    def validateOnCompleteTestData(self, test_loader, simirality_matrix):
        # accuracy and rand index
        nmi = normalized_mutual_info_score
        ari = adjusted_rand_score
        homo = homogeneity_score
        com = completeness_score 
        true_labels = np.concatenate([d[1].cpu().numpy() for i,d in enumerate(test_loader)], 0)
        
        # km = KMeans(n_clusters=len(np.unique(true_labels)), n_init=20, n_jobs=4)
        # y_pred = km.fit_predict(simirality_matrix)

        # ac = AgglomerativeClustering(n_clusters=len(np.unique(true_labels)), affinity='precomputed', linkage='average')
        # ac = AgglomerativeClustering(n_clusters=None,affinity='precomputed', linkage='average', distance_threshold=0.45)
        # y_pred = ac.fit_predict(1+ (-1 * simirality_matrix))

        # y_pred = y_pred.tolist()
        # true_labels = true_labels.tolist()
        # import collections
        # c = collections.Counter(y_pred)
        # y_pred_new = []
        # true_labels_new = []
        # for i, j in zip(y_pred, true_labels):
        #     if c[i] >= 2:
        #         y_pred_new.append(i)
        #         true_labels_new.append(j)
        # print(len(y_pred_new))
        # y_pred = np.array(y_pred_new)
        # true_labels = np.array(true_labels_new)

        sc=SpectralClustering(n_clusters=len(np.unique(true_labels)))
        y_pred=sc.fit(simirality_matrix).labels_

        print(' '*8 + '|==>  nmi: %.4f ,  ari: %.4f,  com: %.4f,  homo: %.4f  <==|'
                      % (nmi(true_labels, y_pred), ari(true_labels, y_pred), com(true_labels, y_pred), homo(true_labels, y_pred)))
        return ari(true_labels, y_pred)

    def align(self, model, dl):
        model.eval()
        pred_match = 0
        ref_match = 0
        TP = 0
        for batch in dl:
            low_seq_0, masked_seq_0, family_0, seq_len_0, low_seq_1, masked_seq_1, family_1, seq_len_1, common_index_0, common_index_1 = batch
            low_seq_0 = low_seq_0.to(self.device)
            low_seq_1 = low_seq_1.to(self.device)
            low_seq = torch.cat((low_seq_0, low_seq_1), axis=0)
            
            start = time.time()
            prediction_scores, prediction_scores_ss, encoded_layers =  model(low_seq)
            elapsed_time = time.time() - start
            # print ("elapsed_time:{0}".format(elapsed_time) + "[sec]")

            prediction_scores0, prediction_scores1 = torch.split(prediction_scores, int(prediction_scores.shape[0]/2))
            encoded_layers0, encoded_layers1 = torch.split(encoded_layers, int(encoded_layers.shape[0]/2))
            z0_list, z1_list =  self.module.em(encoded_layers0, seq_len_0), self.module.em(encoded_layers1, seq_len_1)
            len_TP, len_pred_match, len_ref_match = self.module.test_align(low_seq_0, low_seq_1, z0_list, z1_list, common_index_0, common_index_1, seq_len_0, seq_len_1, args.show_aln)
            TP += len_TP
            pred_match += len_pred_match
            ref_match += len_ref_match

        PPV = TP /pred_match 
        sens = TP / ref_match
        f1 = 2 * PPV * sens/(PPV + sens)
        if args.show_aln == False:
            print("alignment accuracy : ", f1, "sens : ", sens, "PPV : ", PPV)
        return f1 

    def test(self, ds, test_loader, model):
        model.eval()
        data_num = len(test_loader.dataset)
        simirality_matrix = []
        for i in range(data_num):
            single_seq = MyDataset("CLU", np.tile(ds.low_seq[i],(data_num,1)), np.tile(ds.low_seq[i],(data_num,1)),np.tile(ds.family[i],(data_num,1)), np.tile(ds.seq_len[i], data_num)) 
            single_seq = torch.utils.data.DataLoader(single_seq, batch_size, shuffle=False)
            low = []
            for data0, data1 in zip( test_loader, single_seq):
                x0, label0, seq_len_0 = data0
                x1, label1, seq_len_1 = data1
                x0, label0 = x0.to("cuda"),label0.to("cuda"),
                x1, label1 = x1.to("cuda"),label1.to("cuda"),
                x = torch.cat((x0, x1), axis=0) 
                prediction_scores, prediction_scores_ss, encoded_layers =  model(x)
                encoded_layers0, encoded_layers1 = torch.split(encoded_layers, int(encoded_layers.shape[0]/2))
                z0_list, z1_list =  self.module.em(encoded_layers0, seq_len_0), self.module.em(encoded_layers1, seq_len_1)
                _, logits = self.module.match(z0_list, z1_list)
                low.append(torch.squeeze(logits).to('cpu').detach().numpy().copy())
            simirality_matrix.append(np.concatenate(low, 0))
        currentAcc = self.validateOnCompleteTestData(test_loader, np.array(simirality_matrix))
        return currentAcc

def objective():
    config.hidden_size = config.num_attention_heads * config.multiple    
    train = TRAIN(config)
    model = BertModel(config)
    model = BertForMaskedLM(config, model)
    if args.data_mlm:
        config.adam_lr = 2e-4
    # if args.data_sfp:
    #     model = fix_params(model)
    #     config.adam_lr = config.adam_lr * 0.5
    if args.data_mul:
        # model = fix_params(model)
        config.adam_lr = 1e-4
    model = train.model_device(model)
    if args.pretraining:
        model.load_state_dict(torch.load(args.pretraining))
    optimizer = optim.AdamW([{'params': model.parameters(), 'lr': config.adam_lr}])
    return model , optimizer, train, config

config = get_config(file_path = "./RNA_bert_config.json")
data = DATA(args, config)
model, optimizer, train, config = objective()

#now start training
if args.data_mlm:
    dl_MLM = data.load_data_MLM_SFP(args.data_mlm)
    model = train.train_MLM_SFP(model, optimizer, dl_MLM, args.epoch, "MLM")
# elif args.data_sfp:
#     dl_SFP = data.load_data_MLM_SFP(args.data_sfp)
#     model = train.train_MLM_SFP(model, optimizer, dl_SFP, args.epoch, "SFP")
if args.data_mul:
    dl_MUL = data.load_data_MUL(args.data_mul, "MUL")
    model = train.train_MLM_SFP(model, optimizer, dl_MUL, args.epoch, "MUL")

if args.data_alignment: 
    dl_alignment = data.load_data_MUL(args.data_alignment, "MUL")
    alignment_accuracy = train.align(model, dl_alignment)
elif args.data_clustering:
    _, _, ds, test_dl = data.load_data_CLU(args.data_clustering) 
    train.test(ds, test_dl, model)

if args.data_showbase:
    seqs, label, SS,  ds, test_dl  = data.load_data_SHOW(args.data_showbase) 
    features = train.make_feature(model, test_dl)
    features = features.reshape(-1, features.shape[2])
    show_base_PCA(features, label.reshape(-1), SS)

if args.data_embedding:
    seqs, label, test_dl  = data.load_data_EMB(args.data_embedding) 
    features = train.make_feature(model, test_dl, seqs)
    for i, data_set in enumerate(args.embedding_output):
        with open(data_set, 'w') as f:
            for d in features:
                f.write(str(d) + '\n')