update readme & other scripts, add run_defect, and support bart

Author: yuewang-cuhk

run_defect.py

@@ -0,0 +1,331 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
+GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
+using a masked language modeling (MLM) loss.
+"""
+
+from __future__ import absolute_import
+import os
+import sys
+import torch
+import logging
+import argparse
+import math
+import numpy as np
+from io import open
+from tqdm import tqdm
+
+try:
+    from torch.utils.tensorboard import SummaryWriter
+except:
+    from tensorboardX import SummaryWriter
+
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler, TensorDataset
+from torch.utils.data.distributed import DistributedSampler
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+                          RobertaConfig, RobertaModel, RobertaTokenizer,
+                          BartConfig, BartForConditionalGeneration, BartTokenizer,
+                          T5Config, T5ForConditionalGeneration, T5Tokenizer)
+import multiprocessing
+import pdb
+import time
+
+from models import DefectModel
+from configs import add_args, set_seed
+from utils import get_filenames, get_elapse_time, load_and_cache_defect_data
+from models import get_model_size, load_codet5
+
+MODEL_CLASSES = {'roberta': (RobertaConfig, RobertaModel, RobertaTokenizer),
+                 't5': (T5Config, T5ForConditionalGeneration, T5Tokenizer),
+                 'codet5': (T5Config, T5ForConditionalGeneration, RobertaTokenizer),
+                 'bart': (BartConfig, BartForConditionalGeneration, BartTokenizer)}
+
+cpu_cont = multiprocessing.cpu_count()
+
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',
+                    level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def evaluate(args, model, eval_examples, eval_data, write_to_pred=False):
+    eval_sampler = SequentialSampler(eval_data)
+    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
+
+    # Eval!
+    logger.info("***** Running evaluation *****")
+    logger.info("  Num examples = %d", len(eval_examples))
+    logger.info("  Num batches = %d", len(eval_dataloader))
+    logger.info("  Batch size = %d", args.eval_batch_size)
+    eval_loss = 0.0
+    nb_eval_steps = 0
+    model.eval()
+    logits = []
+    labels = []
+    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Evaluating"):
+        inputs = batch[0].to(args.device)
+        label = batch[1].to(args.device)
+        with torch.no_grad():
+            lm_loss, logit = model(inputs, label)
+            eval_loss += lm_loss.mean().item()
+            logits.append(logit.cpu().numpy())
+            labels.append(label.cpu().numpy())
+        nb_eval_steps += 1
+    logits = np.concatenate(logits, 0)
+    labels = np.concatenate(labels, 0)
+    preds = logits[:, 1] > 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),
+    }
+
+    logger.info("***** Eval results *****")
+    for key in sorted(result.keys()):
+        logger.info("  %s = %s", key, str(round(result[key], 4)))
+
+    if write_to_pred:
+        with open(os.path.join(args.output_dir, "predictions.txt"), 'w') as f:
+            for example, pred in zip(eval_examples, preds):
+                if pred:
+                    f.write(str(example.idx) + '\t1\n')
+                else:
+                    f.write(str(example.idx) + '\t0\n')
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    t0 = time.time()
+    args = add_args(parser)
+    logger.info(args)
+
+    # Setup CUDA, GPU & distributed training
+    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:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
+        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
+
+    logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, cpu count: %d",
+                   args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), cpu_cont)
+    args.device = device
+    set_seed(args)
+
+    # Build model
+    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)
+    model = model_class.from_pretrained(args.model_name_or_path)
+
+    if args.model_type == 'codet5':
+        # reset special ids: pad_token_id = 0, bos_token_id = 1, eos_token_id = 2
+        config, model, tokenizer = load_codet5(config, model, tokenizer_class,
+                                               load_extra_ids=False, add_lang_ids=False,
+                                               tokenizer_path=args.tokenizer_path)
+    else:
+        tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name)
+
+    model = DefectModel(model, config, tokenizer, args)
+    logger.info("Finish loading model [%s] from %s", get_model_size(model), args.model_name_or_path)
+
+    if args.load_model_path is not None:
+        logger.info("Reload model from {}".format(args.load_model_path))
+        model.load_state_dict(torch.load(args.load_model_path))
+
+    model.to(device)
+
+    pool = multiprocessing.Pool(cpu_cont)
+    args.train_filename, args.dev_filename, args.test_filename = get_filenames(args.data_dir, args.task, args.sub_task)
+    fa = open(os.path.join(args.output_dir, 'summary.log'), 'a+')
+
+    if args.do_train:
+        if args.n_gpu > 1:
+            # multi-gpu training
+            model = torch.nn.DataParallel(model)
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            summary_fn = '{}/{}'.format(args.summary_dir, '/'.join(args.output_dir.split('/')[1:]))
+            tb_writer = SummaryWriter(summary_fn)
+
+        # Prepare training data loader
+        train_examples, train_data = load_and_cache_defect_data(args, args.train_filename, pool, tokenizer, 'train',
+                                                                is_sample=False)
+        if args.local_rank == -1:
+            train_sampler = RandomSampler(train_data)
+        else:
+            train_sampler = DistributedSampler(train_data)
+        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)
+
+        num_train_optimization_steps = args.num_train_epochs * len(train_dataloader)
+        save_steps = max(len(train_dataloader), 1)
+
+        # Prepare optimizer and schedule (linear warmup and decay)
+        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)
+
+        if args.warmup_steps < 1:
+            warmup_steps = num_train_optimization_steps * args.warmup_steps
+        else:
+            warmup_steps = int(args.warmup_steps)
+        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps,
+                                                    num_training_steps=num_train_optimization_steps)
+
+        # Start training
+        train_example_num = len(train_data)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %d", train_example_num)
+        logger.info("  Batch size = %d", args.train_batch_size)
+        logger.info("  Batch num = %d", math.ceil(train_example_num / args.train_batch_size))
+        logger.info("  Num epoch = %d", args.num_train_epochs)
+
+        global_step, best_acc = 0, 0
+        not_acc_inc_cnt = 0
+        is_early_stop = False
+        for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
+            bar = tqdm(train_dataloader, total=len(train_dataloader), desc="Training")
+            nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
+            model.train()
+            for step, batch in enumerate(bar):
+                batch = tuple(t.to(device) for t in batch)
+                source_ids, labels = batch
+                # pdb.set_trace()
+
+                loss, logits = model(source_ids, labels)
+
+                if args.n_gpu > 1:
+                    loss = loss.mean()  # mean() to average on multi-gpu.
+                if args.gradient_accumulation_steps > 1:
+                    loss = loss / args.gradient_accumulation_steps
+                tr_loss += loss.item()
+
+                nb_tr_examples += source_ids.size(0)
+                nb_tr_steps += 1
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                if nb_tr_steps % args.gradient_accumulation_steps == 0:
+                    # Update parameters
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    scheduler.step()
+                    global_step += 1
+                    train_loss = round(tr_loss * args.gradient_accumulation_steps / nb_tr_steps, 4)
+                    bar.set_description("[{}] Train loss {}".format(cur_epoch, round(train_loss, 3)))
+
+                if (step + 1) % save_steps == 0 and args.do_eval:
+                    logger.info("***** CUDA.empty_cache() *****")
+                    torch.cuda.empty_cache()
+
+                    eval_examples, eval_data = load_and_cache_defect_data(args, args.dev_filename, pool, tokenizer,
+                                                                          'valid', is_sample=False)
+
+                    result = evaluate(args, model, eval_examples, eval_data)
+                    eval_acc = result['eval_acc']
+
+                    if args.data_num == -1:
+                        tb_writer.add_scalar('dev_acc', round(eval_acc, 4), cur_epoch)
+
+                    # save last checkpoint
+                    last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
+                    if not os.path.exists(last_output_dir):
+                        os.makedirs(last_output_dir)
+
+                    if True or args.data_num == -1 and args.save_last_checkpoints:
+                        model_to_save = model.module if hasattr(model, 'module') else model
+                        output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
+                        torch.save(model_to_save.state_dict(), output_model_file)
+                        logger.info("Save the last model into %s", output_model_file)
+
+                    if eval_acc > best_acc:
+                        not_acc_inc_cnt = 0
+                        logger.info("  Best acc: %s", round(eval_acc, 4))
+                        logger.info("  " + "*" * 20)
+                        fa.write("[%d] Best acc changed into %.4f\n" % (cur_epoch, round(eval_acc, 4)))
+                        best_acc = eval_acc
+                        # Save best checkpoint for best ppl
+                        output_dir = os.path.join(args.output_dir, 'checkpoint-best-acc')
+                        if not os.path.exists(output_dir):
+                            os.makedirs(output_dir)
+                        if args.data_num == -1 or True:
+                            model_to_save = model.module if hasattr(model, 'module') else model
+                            output_model_file = os.path.join(output_dir, "pytorch_model.bin")
+                            torch.save(model_to_save.state_dict(), output_model_file)
+                            logger.info("Save the best ppl model into %s", output_model_file)
+                    else:
+                        not_acc_inc_cnt += 1
+                        logger.info("acc does not increase for %d epochs", not_acc_inc_cnt)
+                        if not_acc_inc_cnt > args.patience:
+                            logger.info("Early stop as acc do not increase for %d times", not_acc_inc_cnt)
+                            fa.write("[%d] Early stop as not_acc_inc_cnt=%d\n" % (cur_epoch, not_acc_inc_cnt))
+                            is_early_stop = True
+                            break
+
+                model.train()
+            if is_early_stop:
+                break
+
+            logger.info("***** CUDA.empty_cache() *****")
+            torch.cuda.empty_cache()
+
+        if args.local_rank in [-1, 0] and args.data_num == -1:
+            tb_writer.close()
+
+    if args.do_test:
+        logger.info("  " + "***** Testing *****")
+        logger.info("  Batch size = %d", args.eval_batch_size)
+
+        for criteria in ['best-acc']:  # , 'last'
+            file = os.path.join(args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
+            # logger.info("*" * 10 + "Start testing" + "*" * 10)
+            logger.info("Reload model from {}".format(file))
+            model.load_state_dict(torch.load(file))
+
+            if args.n_gpu > 1:
+                # multi-gpu training
+                model = torch.nn.DataParallel(model)
+
+            eval_examples, eval_data = load_and_cache_defect_data(args, args.test_filename, pool, tokenizer, 'test',
+                                                                  False)
+
+            result = evaluate(args, model, eval_examples, eval_data, write_to_pred=True)
+            logger.info("  test_acc=%.4f", result['eval_acc'])
+            logger.info("  " + "*" * 20)
+
+            fa.write("[%s] test-acc: %.4f\n" % (criteria, result['eval_acc']))
+            if args.res_fn:
+                with open(args.res_fn, 'a+') as f:
+                    f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
+                    f.write("[%s] acc: %.4f\n\n" % (
+                        criteria, result['eval_acc']))
+    fa.close()
+
+
+if __name__ == "__main__":
+    # print(' '.join(sys.argv[1:]))
+    main()