run.py

from __future__ import absolute_import, division, print_function

import argparse
import glob
import logging
import os
import pickle
import random
import re
import shutil
import csv
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler
from torch.utils.data.distributed import DistributedSampler
from sklearn.metrics import f1_score,precision_score,recall_score,accuracy_score
from sklearn.neighbors import KNeighborsClassifier

import json
from scipy import spatial
try:
    from torch.utils.tensorboard import SummaryWriter
except:
    from tensorboardX import SummaryWriter
import warnings
warnings.filterwarnings("ignore")
from tqdm import tqdm, trange
import multiprocessing
from model import Model,Model_train_3
cpu_cont = multiprocessing.cpu_count()
from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
                          BertConfig, BertForMaskedLM, BertTokenizer,
                          GPT2Config, GPT2LMHeadModel, GPT2Tokenizer,
                          OpenAIGPTConfig, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer,
                          RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer,
                          DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer)

logger = logging.getLogger(__name__)

MODEL_CLASSES = {
    'gpt2': (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer),
    'openai-gpt': (OpenAIGPTConfig, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    'bert': (BertConfig, BertForMaskedLM, BertTokenizer),
    'roberta': (RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer),
    'distilbert': (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer)
}

class InputFeatures(object):
    def __init__(self,
                 input_tokens,
                 input_ids,
                 idx,
                 label,
                 epoch  
    ):
        self.input_tokens = input_tokens
        self.input_ids = input_ids
        self.idx=str(idx)
        self.label=label
        self.epoch = epoch
        
def convert_examples_to_features(js,tokenizer,args,epoch):

    code=' '.join(js[2].split())
    code_tokens=tokenizer.tokenize(code)[:args.block_size-2]
    source_tokens =[tokenizer.cls_token]+code_tokens+[tokenizer.sep_token]
    source_ids =  tokenizer.convert_tokens_to_ids(source_tokens)
    padding_length = args.block_size - len(source_ids)
    source_ids+=[tokenizer.pad_token_id]*padding_length

    label = int(js[0])

    return InputFeatures(source_tokens,source_ids,js[1],label,epoch)

def calculate(examles):
    cnt_0 = 0
    cnt_1 = 0
    for example in examles:
        if example.label == 0:
            cnt_0 += 1
        if example.label == 1:
            cnt_1 += 1

    print('*' * 100)
    print('Positive samples:',cnt_1)
    print('Negative samples:',cnt_0)
    print('*' * 100)

class TextDataset(Dataset):
    def __init__(self, tokenizer, args, file_path=None, type = None, 
          indexes_la_first_fintune =None,indexes_rn_step1 =None,epoch = None,dataset_again = None):
        self.examples = []
        self.indexes_la = []
        index = 0
        if type == 'iterative_finetune':
            self.examples = dataset_again.examples
        with open(file_path, "r", encoding='ISO-8859-1') as f:
            for line in f.readlines():
                line = line.strip().split('<CODESPLIT>')
                example = convert_examples_to_features(line,tokenizer,args,epoch)
       
                # the Second(Third) Finetune just need labeled samples and RN samples
                if type == 'second_finetune' or type == 'third_finetune' or type == 'iterative_finetune':
                    if index in indexes_la_first_fintune:
                        example.label = 1
                        self.examples.append(example)
                    if index in indexes_rn_step1:
                        example.label = 0
                        self.examples.append(example)
                elif type == 'iterative_finetune_remain':
                    if index not in indexes_la_first_fintune and index not in indexes_rn_step1:
                        example.label = -1
                        self.examples.append(example)

                # the First Finetune,the Step1,the Evaluate and the Write need all the samples
                else:
                    self.examples.append(example)

                if int(example.label) == 1:
                    self.indexes_la.append(index)
                index += 1
                
        # choose the labeled samples randomly if First Fintune
        if type == 'first_finetune':
            random.shuffle(self.indexes_la)
            self.indexes_la = self.indexes_la[:round(len(self.indexes_la)*args.label_ratio)]
            # write indexes of labeled samples to the txt
            path = os.path.join(args.labels_file,'indexes_la.txt')
            writeindex(path,self.indexes_la)
            for i in range(len(self.examples)):
                if i not in self.indexes_la:
                    self.examples[i].label = 0

        # get the labeled samples based on the First Fintune if Step1
        if type == 'step1':
            for i in range(len(self.examples)):
                if i not in indexes_la_first_fintune:
                    self.examples[i].label = 0

        calculate(self.examples)

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):       
        return torch.tensor(self.examples[i].input_ids),torch.tensor(self.examples[i].label),torch.tensor(self.examples[i].epoch)

    def getlabel(self):
        return self.indexes_la   
    
def set_seed(seed=42):
    random.seed(seed)
    os.environ['PYHTONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

def writeindex(file_path,indexes):
    f=open(file_path,"w")
    for index in indexes: 
        f.write(str(index)+'\n')
    f.close()

def readindex(file_path):
    f=open(file_path, encoding='gbk')
    indexes=[]
    for line in f:
        indexes.append(int(line.strip()))
    return indexes

def Select(model, train_dataset, args):

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    train_sampler = SequentialSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.eval_batch_size)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    """ find RN samples """
    labels=[]
    features=[]
    for batch in tqdm(train_dataloader,total=len(train_dataloader)):
        inputs = batch[0].to(args.device)
        label=batch[1].to(args.device) 
        with torch.no_grad():
            _ ,feature = model(inputs)
            labels.append(label.cpu().numpy())
            features.append(feature.cpu().numpy())

    labels_1d = np.concatenate(labels).flatten()
    features_2d = np.concatenate(features)

    # get the indexes of labeled samples and unlabeled samples 
    indexes_la = np.where(labels_1d == 1)
    indexes_la_1d = np.array(indexes_la).flatten()
    indexes_un = np.where(labels_1d == 0)
    indexes_un_1d = np.array(indexes_un).flatten()
    indexes_rn = []

    # get the features of labeled samples and unlabeled samples 
    train_xs_la = features_2d[indexes_la]
    train_xs_un = features_2d[indexes_un]

    knn_k = int(round(0.3 * (len(train_xs_la))))
    knn = KNeighborsClassifier(n_neighbors=knn_k)
    knn.fit(train_xs_la, indexes_la_1d)
    distances_sum ,indexes_knn = knn.kneighbors(train_xs_un, knn_k, True)
    distances_sum = distances_sum.sum(axis = 1)

    train_xs_la_near_means = []
    for indexes in indexes_knn:
        indexes_1d = (np.array(indexes).flatten()).tolist()
        train_xs_la_near = train_xs_la[indexes_1d]
        train_xs_la_near_mean = train_xs_la_near.mean(axis = 0)
        train_xs_la_near_means.append(train_xs_la_near_mean)
    distances_mean = np.linalg.norm(train_xs_un - train_xs_la_near_means, axis=1,ord = 1)

    standards = sorted(distances_mean)
    standards.reverse()
    ratio_k = 1
    ratio = args.label_ratio * ratio_k * ((1-args.label_ratio)*args.P_num + args.N_num) / (args.label_ratio*args.P_num)
    standard_mean = standards[round(len(indexes_la_1d)*ratio)]

    standards = sorted(distances_sum)
    standards.reverse()
    ratio_k = 1
    ratio = args.label_ratio * ratio_k * ((1-args.label_ratio)*args.P_num + args.N_num) / (args.label_ratio*args.P_num)
    standard_sum = standards[round(len(indexes_la_1d)*ratio)]

    # get RN samples' indexes
    for index in range(len(indexes_un_1d)):
        if distances_mean[index] > standard_mean and distances_sum[index] > standard_sum:
            indexes_rn.append(indexes_un_1d[index])

    return indexes_rn
    
def train(args, train_dataset, model, tokenizer):
    """ Train the model """ 
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)

    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, 
                                  batch_size=args.train_batch_size,num_workers=4,pin_memory=True)
    args.max_steps=args.epoch*len( train_dataloader)
    args.save_steps=len( train_dataloader)
    args.warmup_steps=len( train_dataloader)
    args.logging_steps=len( train_dataloader)
    args.num_train_epochs=args.epoch
    model.to(args.device)

    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.max_steps*0.1,
                                                num_training_steps=args.max_steps)
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)

    checkpoint_last = os.path.join(args.output_dir, 'checkpoint-last')
    scheduler_last = os.path.join(checkpoint_last, 'scheduler.pt')
    optimizer_last = os.path.join(checkpoint_last, 'optimizer.pt')
    if os.path.exists(scheduler_last):
        scheduler.load_state_dict(torch.load(scheduler_last))
    if os.path.exists(optimizer_last):
        optimizer.load_state_dict(torch.load(optimizer_last))
    
    global_step = args.start_step
    tr_loss, logging_loss,avg_loss,tr_nb,tr_num,train_loss = 0.0, 0.0,0.0,0,0,0
    best_mrr=0.0
    best_acc=0.0
    model.zero_grad()
 
    for idx in range(args.start_epoch, int(args.num_train_epochs)): 
        bar = tqdm(train_dataloader,total=len(train_dataloader))
        tr_num=0
        train_loss=0
        for step, batch in enumerate(bar):
            # inputs = batch[0].to(args.device)   
            inputs = batch[0].to(args.device)             
            labels=batch[1].to(args.device) 
            model.train()
            loss, _ , _  = model(inputs,labels)
            if args.n_gpu > 1:
                loss = loss.mean()  
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
                torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
            else:
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

            tr_loss += loss.item()
            tr_num+=1
            train_loss+=loss.item()
            if avg_loss==0:
                avg_loss=tr_loss
            avg_loss=round(train_loss/tr_num,5)
            bar.set_description("epoch {} loss {}".format(idx,avg_loss))

            if (step + 1) % args.gradient_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                scheduler.step()  
                global_step += 1
                output_flag=True
                avg_loss=round(np.exp((tr_loss - logging_loss) /(global_step- tr_nb)),4)
                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    logging_loss = tr_loss
                    tr_nb=global_step
            
            if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    
                    if args.local_rank == -1 and args.evaluate_during_training:  
                        results = evaluate(args, model, tokenizer)
                        for key, value in results.items():
                            logger.info("  %s = %s", key, round(value,4))                    
                        
                    if results['eval_acc']>best_acc:
                        best_acc=results['eval_acc']
                        logger.info("  "+"*"*20)  
                        logger.info("  Best acc:%s",round(best_acc,4))
                        logger.info("  "+"*"*20)                          
                        
                        checkpoint_prefix = 'checkpoint-best-acc'
                        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))                        
                        if not os.path.exists(output_dir):
                            os.makedirs(output_dir)                        
                        model_to_save = model.module if hasattr(model,'module') else model
                        output_dir = os.path.join(output_dir, '{}'.format('model.bin')) 
                        torch.save(model_to_save.state_dict(), output_dir)
                        logger.info("Saving model checkpoint to %s", output_dir)
        
        logger.info("Fintuned model on the test set:")
        evaluate(args,model,tokenizer,type='eval')
        evaluate(args,model,tokenizer,type='test')

def train_iterative(args, train_dataset, model, epoch_now = None):
    """ Train the model """ 
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)

    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, 
                                  batch_size=args.train_batch_size,num_workers=4,pin_memory=True)
    # args.max_steps=args.epoch*len( train_dataloader)
    args.max_steps=1*len( train_dataloader)

    args.save_steps=len( train_dataloader)
    args.warmup_steps=len( train_dataloader)
    args.logging_steps=len( train_dataloader)
    # args.num_train_epochs=args.epoch
    args.num_train_epochs=1

    model.to(args.device)

    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.max_steps*0.1,
                                                num_training_steps=args.max_steps)
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)

    checkpoint_last = os.path.join(args.output_dir, 'checkpoint-last')
    scheduler_last = os.path.join(checkpoint_last, 'scheduler.pt')
    optimizer_last = os.path.join(checkpoint_last, 'optimizer.pt')
    if os.path.exists(scheduler_last):
        scheduler.load_state_dict(torch.load(scheduler_last))
    if os.path.exists(optimizer_last):
        optimizer.load_state_dict(torch.load(optimizer_last))
    
    global_step = args.start_step
    tr_loss, logging_loss,avg_loss,tr_nb,tr_num,train_loss = 0.0, 0.0,0.0,0,0,0
    best_mrr=0.0
    best_acc=0.0
    model.zero_grad()

    logits = []
    for idx in range(args.start_epoch, int(args.num_train_epochs)): 
        bar = tqdm(train_dataloader,total=len(train_dataloader))
        tr_num=0
        train_loss=0
        for step, batch in enumerate(bar):
            inputs = batch[0].to(args.device)        
            labels=batch[1].to(args.device) 
            epochs_now = batch[2].to(args.device) 
            model.train()
            loss, logit , _  = model(inputs,labels,epochs_now = epochs_now,epoch_max = args.epoch,type = 'do_train_iterative')
            logits.append(logit.cpu().detach().numpy())
            if args.n_gpu > 1:
                loss = loss.mean()  
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
                torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
            else:
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

            tr_loss += loss.item()
            tr_num+=1
            train_loss+=loss.item()
            if avg_loss==0:
                avg_loss=tr_loss
            avg_loss=round(train_loss/tr_num,5)
            bar.set_description("epoch {} loss {}".format(epoch_now,avg_loss))

            if (step + 1) % args.gradient_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                scheduler.step()  
                global_step += 1
                output_flag=True
                avg_loss=round(np.exp((tr_loss - logging_loss) /(global_step- tr_nb)),4)
                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    logging_loss = tr_loss
                    tr_nb=global_step
    
    logits=np.concatenate(logits,0)
    probs_la = []
    probs_rn = []
    for i in logits:
        if i >=0.5:
            probs_la.append(i)
        if i <0.5:
            probs_rn.append(i)
    probs_la = sorted(probs_la)
    prob_la = probs_la[round(0.20*len(probs_la))] 

    probs_rn = sorted(probs_rn)
    prob_rn = probs_rn[round(0.80*len(probs_rn))] 
    print('prob_la',prob_la)
    print('prob_rn',prob_rn)
    return prob_la,prob_rn

def evaluate(args, model, tokenizer,type = 'eval'):

    eval_output_dir = args.output_dir
    if type == 'test':
        eval_dataset = TextDataset(tokenizer, args,args.test_data_file, epoch=-1)
    elif type == 'eval':
        eval_dataset = TextDataset(tokenizer, args,args.eval_data_file, epoch=-1)

    if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
        os.makedirs(eval_output_dir)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)

    eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
    eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size,num_workers=4,pin_memory=True)

    if args.n_gpu > 1 and args.evaluate_during_training is False:
        model = torch.nn.DataParallel(model)

    logger.info("***** Running evaluation *****")
    logger.info("  Evaluation type = %s", type)
    logger.info("  Num examples = %d", len(eval_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
    eval_loss = 0.0
    nb_eval_steps = 0
    model.eval()
    logits=[] 
    labels=[]
    for batch in eval_dataloader:
        inputs = batch[0].to(args.device)        
        label=batch[1].to(args.device) 
        with torch.no_grad():
            lm_loss,logit,context_embeddings = model(inputs,label)
            eval_loss += lm_loss.mean().item()
            logits.append(logit.cpu().numpy())
            labels.append(label.cpu().numpy())
        nb_eval_steps += 1

    labels_1d = np.concatenate(labels).flatten()
    logits_1d = np.concatenate(logits).flatten()
    logits=np.concatenate(logits,0)
    labels=np.concatenate(labels,0)
    preds=logits[:,0]>0.5
    eval_acc=np.mean(labels==preds)
    eval_loss = eval_loss / nb_eval_steps
    perplexity = torch.tensor(eval_loss)
    result = {
        "eval_loss": float(perplexity),
        "eval_acc":round(eval_acc,4),
    }

    preds = np.zeros(labels_1d.shape[0])
    for index in range(labels_1d.shape[0]):
        if logits_1d[index] < 0.5:
            preds[index] = 0
        else:
            preds[index] = 1

    acc = accuracy_score(labels_1d, preds)
    pre = precision_score(labels_1d, preds)
    rec = recall_score(labels_1d, preds)
    f1 = f1_score(labels_1d, preds)

    print("acc:",acc)
    print("pre:",pre)
    print("rec:",rec)
    print("f1:",f1)		

    return result    

def iterative(args, model, tokenizer,eval_dataset=None,indexes_la=None,indexes_rn=None,indexes_un=None,prob_la=None,prob_rn=None):

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)

    eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
    eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size,num_workers=4,pin_memory=True)

    if args.n_gpu > 1 and args.evaluate_during_training is False:
        model = torch.nn.DataParallel(model)

    logger.info("***** Running evaluation *****")
    logger.info("  Num examples = %d", len(eval_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)

    model.eval()
    logits=[] 

    for batch in eval_dataloader:
        inputs = batch[0].to(args.device)        
        with torch.no_grad():
            logit,_ = model(inputs)
            logits.append(logit.cpu().numpy())
    logits=np.concatenate(logits,0)

    # RP and RN increase simultaneously according to the threshold
    cnt_rp = 0
    cnt_rn = 0
    labels= []
    indexes_rn_add = []
    indexes_la_add = []
    logits_sort = sorted(logits)
    logits_sort.reverse()
    for i in range(len(logits)):
        if logits[i]>=0.5:
            labels.append(1)
        else:
            labels.append(0)

    for i in range(len(labels)):
        if labels[i] == 0 and logits[i]<=prob_rn:
            indexes_rn.append(indexes_un[i])
            indexes_rn_add.append(indexes_un[i])
            cnt_rn += 1

    ratio = args.label_ratio*args.P_num / ((1-args.label_ratio)*args.P_num + args.N_num)

    cnt_rp = round(cnt_rn * ratio)
    standard = logits_sort[cnt_rp]

    for i in range(len(logits)):
        if logits[i] >= standard:
            indexes_la.append(indexes_un[i])
            indexes_la_add.append(indexes_un[i])

    indexes_un = []
    for i in range(args.P_num+args.N_num):
        if i not in indexes_la and i not in indexes_rn:
            indexes_un.append(i)

    return indexes_la,indexes_rn,indexes_un,indexes_la_add,indexes_rn_add
          
def main():
    parser = argparse.ArgumentParser()

    """ Required parameters """ 
    parser.add_argument("--train_data_file", default=None, type=str, required=True,
                        help="The input training data file (a text file).")
    parser.add_argument("--eval_data_file", default=None, type=str, required=True,
                        help="The input valid data file (a text file).")
    parser.add_argument("--test_data_file", default=None, type=str, required=True,
                        help="The input test data file (a text file).")
    parser.add_argument("--output_dir", default=None, type=str, required=True,
                        help="The output directory where the model predictions and checkpoints will be written.")

    """ Other parameters """ 
    parser.add_argument("--labels_file", default=None, type=str,
                        help="The labels file.")
    parser.add_argument("--label_ratio", default=None, type=float,
                        help="The label ratio")
    parser.add_argument("--do_train_1", action='store_true',
                        help="Whether to run first training.")   
    parser.add_argument("--do_step1", action='store_true',
                        help="Whether to do step1.") 
    parser.add_argument("--do_train_2", action='store_true',
                        help="Whether to run second training.")  
    parser.add_argument("--do_train_3", action='store_true',
                        help="Whether to run third training.")  
    parser.add_argument("--do_train_iterative", action='store_true',
                        help="Whether to run iterative training.")  
    parser.add_argument("--do_test", action='store_true',
                        help="Whether to test.")  
    parser.add_argument("--P_num", default=-1, type=int,
                        help="num of P")
    parser.add_argument("--N_num", default=-1, type=int,
                        help="num of N")
    
    parser.add_argument("--model_type", default="bert", type=str,
                        help="The model architecture to be fine-tuned.")
    parser.add_argument("--model_name_or_path", default=None, type=str,
                        help="The model checkpoint for weights initialization.")
    parser.add_argument("--mlm", action='store_true',
                        help="Train with masked-language modeling loss instead of language modeling.")
    parser.add_argument("--mlm_probability", type=float, default=0.15,
                        help="Ratio of tokens to mask for masked language modeling loss")
    parser.add_argument("--config_name", default="", type=str,
                        help="Optional pretrained config name or path if not the same as model_name_or_path")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Optional pretrained tokenizer name or path if not the same as model_name_or_path")
    parser.add_argument("--cache_dir", default="", type=str,
                        help="Optional directory to store the pre-trained models downloaded from s3 (instread of the default one)")
    parser.add_argument("--block_size", default=-1, type=int,
                        help="Optional input sequence length after tokenization."
                             "The training dataset will be truncated in block of this size for training."
                             "Default to the model max input length for single sentence inputs (take into account special tokens).")
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")
    parser.add_argument("--evaluate_during_training", action='store_true',
                        help="Run evaluation during training at each logging step.")
    parser.add_argument("--train_batch_size", default=4, type=int,
                        help="Batch size per GPU/CPU for training.")
    parser.add_argument("--eval_batch_size", default=4, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument("--learning_rate", default=5e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument("--weight_decay", default=0.0, type=float,
                        help="Weight deay if we apply some.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
                        help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument("--num_train_epochs", default=1.0, type=float,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--max_steps", default=-1, type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
    parser.add_argument("--warmup_steps", default=0, type=int,
                        help="Linear warmup over warmup_steps.")
    parser.add_argument('--logging_steps', type=int, default=50,
                        help="Log every X updates steps.")
    parser.add_argument('--save_steps', type=int, default=50,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument('--save_total_limit', type=int, default=None,
                        help='Limit the total amount of checkpoints, delete the older checkpoints in the output_dir, does not delete by default')
    parser.add_argument("--eval_all_checkpoints", action='store_true',
                        help="Evaluate all checkpoints starting with the same prefix as model_name_or_path ending and ending with step number")
    parser.add_argument("--no_cuda", action='store_true',
                        help="Avoid using CUDA when available")
    parser.add_argument('--overwrite_output_dir', action='store_true',
                        help="Overwrite the content of the output directory")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--epoch', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--fp16', action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
    parser.add_argument("--local_rank", type=int, default=-1,
                        help="For distributed training: local_rank")
    parser.add_argument('--server_ip', type=str, default='', help="For distant debugging.")
    parser.add_argument('--server_port', type=str, default='', help="For distant debugging.")

    args = parser.parse_args()

    """ basic settings """
    if args.server_ip and args.server_port:
        import ptvsd
        print("Waiting for debugger attach")
        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
        ptvsd.wait_for_attach()
   
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else: 
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend='nccl')
        args.n_gpu = 1
    args.device = device
    args.per_gpu_train_batch_size=args.train_batch_size//args.n_gpu
    args.per_gpu_eval_batch_size=args.eval_batch_size//args.n_gpu

    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S',
                        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
 
    set_seed(args.seed)
  
    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  
    args.start_epoch = 0
    args.start_step = 0
    checkpoint_last = os.path.join(args.output_dir, 'checkpoint-last')

    if os.path.exists(checkpoint_last) and os.listdir(checkpoint_last):
        args.model_name_or_path = os.path.join(checkpoint_last, 'pytorch_model.bin')
        args.config_name = os.path.join(checkpoint_last, 'config.json')
        idx_file = os.path.join(checkpoint_last, 'idx_file.txt')
        with open(idx_file, encoding='utf-8') as idxf:
            args.start_epoch = int(idxf.readlines()[0].strip()) + 1
        step_file = os.path.join(checkpoint_last, 'step_file.txt')
        if os.path.exists(step_file):
            with open(step_file, encoding='utf-8') as stepf:
                args.start_step = int(stepf.readlines()[0].strip())
        logger.info("reload model from {}, resume from {} epoch".format(checkpoint_last, args.start_epoch))

    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                          cache_dir=args.cache_dir if args.cache_dir else None)
    config.num_labels=1
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name,
                                                do_lower_case=args.do_lower_case,
                                                cache_dir=args.cache_dir if args.cache_dir else None)
    if args.block_size <= 0:
        args.block_size = tokenizer.max_len_single_sentence 
    args.block_size = min(args.block_size, tokenizer.max_len_single_sentence)
    if args.model_name_or_path:
        model = model_class.from_pretrained(args.model_name_or_path,
                                            from_tf=bool('.ckpt' in args.model_name_or_path),
                                            config=config,
                                            cache_dir=args.cache_dir if args.cache_dir else None,
                                            )    
    else:
        model = model_class(config)
    if args.do_train_1 or args.do_step1 or args.do_train_2 or args.do_train_iterative:
        model=Model(model,config,tokenizer,args)
    if args.do_train_3 or args.do_test:
        model=Model_train_3(model,config,tokenizer,args)

    if args.local_rank == 0:
        torch.distributed.barrier()  

    logger.info("Training/evaluation parameters %s", args)

    """ First finetune """
    if args.do_train_1:
        if args.local_rank not in [-1, 0]:
            torch.distributed.barrier()  
        
        if args.local_rank == 0:
            torch.distributed.barrier()

        # get the labeled train dataset, indexes of labeled data 
        train_dataset = TextDataset(tokenizer, args,args.train_data_file,'first_finetune',epoch=-1)
        indexes_la = train_dataset.getlabel()

        # train the model based on the labeled train dataset
        train(args, train_dataset, model, tokenizer)

    """ select RN samples based on the Step1 method """
    if args.do_step1:
        
        # load the first fintuned model 
        logger.info("select RN samples")
        checkpoint_prefix = 'checkpoint-best-acc/model.bin'
        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))  
        model.load_state_dict(torch.load(output_dir))                  
        model.to(args.device)

        # load the labeled train dataset, indexes of labeled data 
        path_la = os.path.join(args.labels_file,'indexes_la.txt')
        indexes_la = readindex(path_la)
        train_dataset = TextDataset(tokenizer, args,args.train_data_file,'step1',indexes_la,epoch=-1)

        # execute step1 method
        indexes_rn = Select(model,train_dataset,args)

        # write the RN samples' indexes to the txt
        path = os.path.join(args.labels_file,'indexes_rn.txt')
        writeindex(path,indexes_rn)

    """ Second finetune """
    if args.do_train_2:

        if args.local_rank not in [-1, 0]:
            torch.distributed.barrier()  
        if args.local_rank == 0:
            torch.distributed.barrier()

        # load the labeled train dataset, indexes of labeled samples, indexes of RN samples
        path_la = os.path.join(args.labels_file,'indexes_la.txt')
        path_rn = os.path.join(args.labels_file,'indexes_rn.txt')
        indexes_la = readindex(path_la)
        indexes_rn = readindex(path_rn)
        train_dataset_again = TextDataset(tokenizer, args,args.train_data_file,'second_finetune',indexes_la,indexes_rn,epoch=-1)
        
        # train the model based on the labeled train dataset
        train(args, train_dataset_again, model, tokenizer)

    if args.do_train_iterative:
        if args.local_rank not in [-1, 0]:
            torch.distributed.barrier()  
        if args.local_rank == 0:
            torch.distributed.barrier()

        checkpoint_prefix = 'checkpoint-best-acc/model.bin'
        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))  
        model.load_state_dict(torch.load(output_dir))                  
        model.to(args.device)

        # load the labeled train dataset, indexes of labeled samples, indexes of RN samples
        path_la = os.path.join(args.labels_file,'indexes_la.txt')
        path_rn = os.path.join(args.labels_file,'indexes_rn.txt')
        indexes_la = readindex(path_la)
        indexes_rn = readindex(path_rn)
        indexes_un = []
        for i in range(args.P_num+args.N_num):
            if i not in indexes_la and i not in indexes_rn:
                indexes_un.append(i)

        train_dataset_again = TextDataset(tokenizer, args,args.train_data_file,'third_finetune',indexes_la,indexes_rn,epoch = 0)
        train_dataset_remain = TextDataset(tokenizer, args,args.train_data_file,'iterative_finetune_remain',indexes_la,indexes_rn,epoch = -1)

        for i in range(args.epoch):
            prob_la,prob_rn = train_iterative(args, train_dataset_again, model,epoch_now = i)
            indexes_la,indexes_rn,indexes_un,indexes_la_add,indexes_rn_add = iterative(args, model, tokenizer,
                      eval_dataset=train_dataset_remain,indexes_la=indexes_la,indexes_rn=indexes_rn,indexes_un=indexes_un,
                    prob_la = prob_la, prob_rn = prob_rn)
            
            if len(indexes_un) == 0:
                break
            
            file_la = 'indexes_la_' + str(i)  + '.txt'
            file_rn = 'indexes_rn_' + str(i)  + '.txt'

            path = os.path.join(args.labels_file,file_la)
            writeindex(path,indexes_la)
            path = os.path.join(args.labels_file,file_rn)
            writeindex(path,indexes_rn)

            train_dataset_again = TextDataset(tokenizer, args,args.train_data_file,'iterative_finetune',
                    indexes_la_add,indexes_rn_add,epoch = i+1,dataset_again=train_dataset_again)
            train_dataset_remain = TextDataset(tokenizer, args,args.train_data_file,'iterative_finetune_remain',
                    indexes_la,indexes_rn,epoch=-1)
            
            # test the iterative finetuned model on the test dataset
            logger.info("Fintuned model on the test set:")
            evaluate(args,model,tokenizer,type='eval')
            evaluate(args,model,tokenizer,type='test')

        file_la = 'indexes_la_last.txt'
        file_rn = 'indexes_rn_last.txt'
        path = os.path.join(args.labels_file,file_la)
        writeindex(path,indexes_la)
        path = os.path.join(args.labels_file,file_rn)
        writeindex(path,indexes_rn)

        checkpoint_prefix = 'checkpoint-iterative'
        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))                        
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)                        
        model_to_save = model.module if hasattr(model,'module') else model
        output_dir = os.path.join(output_dir, '{}'.format('model.bin')) 
        torch.save(model_to_save.state_dict(), output_dir)
        logger.info("Saving model checkpoint to %s", output_dir)

    if args.do_train_3:
        path_la = os.path.join(args.labels_file,'indexes_la_last.txt')
        path_rn = os.path.join(args.labels_file,'indexes_rn_last.txt')

        indexes_la = readindex(path_la)
        indexes_rn = readindex(path_rn)
        train_dataset = TextDataset(tokenizer, args,args.train_data_file,'third_finetune',
                indexes_la,indexes_rn,epoch = 0)
        if args.local_rank not in [-1, 0]:
            torch.distributed.barrier()  
        if args.local_rank == 0:
            torch.distributed.barrier()
        
        train(args,train_dataset,model,tokenizer)

    if args.do_test:
        if args.local_rank not in [-1, 0]:
            torch.distributed.barrier()  
        if args.local_rank == 0:
            torch.distributed.barrier()

        checkpoint_prefix = 'model.bin'
        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))  
        model.load_state_dict(torch.load(output_dir))                  
        model.to(args.device)
        evaluate(args,model,tokenizer,type='test')


if __name__ == "__main__":

    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    main()