|
import random |
|
import torch |
|
import logging |
|
import multiprocessing |
|
import numpy as np |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def add_args(parser): |
|
parser.add_argument("--task", type=str, required=True, |
|
choices=['summarize', 'concode', 'translate', 'refine', 'defect', 'clone', 'multi_task']) |
|
parser.add_argument("--sub_task", type=str, default='') |
|
parser.add_argument("--lang", type=str, default='') |
|
parser.add_argument("--eval_task", type=str, default='') |
|
parser.add_argument("--model_type", default="codet5", type=str, choices=['roberta', 'bart', 'codet5']) |
|
parser.add_argument("--add_lang_ids", action='store_true') |
|
parser.add_argument("--data_num", default=-1, type=int) |
|
parser.add_argument("--start_epoch", default=0, type=int) |
|
parser.add_argument("--num_train_epochs", default=100, type=int) |
|
parser.add_argument("--patience", default=5, type=int) |
|
parser.add_argument("--cache_path", type=str, required=True) |
|
parser.add_argument("--summary_dir", type=str, required=True) |
|
parser.add_argument("--data_dir", type=str, required=True) |
|
parser.add_argument("--res_dir", type=str, required=True) |
|
parser.add_argument("--res_fn", type=str, default='') |
|
parser.add_argument("--add_task_prefix", action='store_true', help="Whether to add task prefix for t5 and codet5") |
|
parser.add_argument("--save_last_checkpoints", action='store_true') |
|
parser.add_argument("--always_save_model", action='store_true') |
|
parser.add_argument("--do_eval_bleu", action='store_true', help="Whether to evaluate bleu on dev set.") |
|
|
|
|
|
parser.add_argument("--model_name_or_path", default="roberta-base", type=str, |
|
help="Path to pre-trained model: e.g. roberta-base") |
|
parser.add_argument("--output_dir", default=None, type=str, required=True, |
|
help="The output directory where the model predictions and checkpoints will be written.") |
|
parser.add_argument("--load_model_path", default=None, type=str, |
|
help="Path to trained model: Should contain the .bin files") |
|
|
|
parser.add_argument("--train_filename", default=None, type=str, |
|
help="The train filename. Should contain the .jsonl files for this task.") |
|
parser.add_argument("--dev_filename", default=None, type=str, |
|
help="The dev filename. Should contain the .jsonl files for this task.") |
|
parser.add_argument("--test_filename", default=None, type=str, |
|
help="The test filename. Should contain the .jsonl files for this task.") |
|
|
|
parser.add_argument("--config_name", default="", type=str, |
|
help="Pretrained config name or path if not the same as model_name") |
|
parser.add_argument("--tokenizer_name", default="roberta-base", type=str, |
|
help="Pretrained tokenizer name or path if not the same as model_name") |
|
parser.add_argument("--max_source_length", default=64, type=int, |
|
help="The maximum total source sequence length after tokenization. Sequences longer " |
|
"than this will be truncated, sequences shorter will be padded.") |
|
parser.add_argument("--max_target_length", default=32, type=int, |
|
help="The maximum total target sequence length after tokenization. Sequences longer " |
|
"than this will be truncated, sequences shorter will be padded.") |
|
|
|
parser.add_argument("--do_train", action='store_true', |
|
help="Whether to run eval on the train set.") |
|
parser.add_argument("--do_eval", action='store_true', |
|
help="Whether to run eval on the dev set.") |
|
parser.add_argument("--do_test", action='store_true', |
|
help="Whether to run eval on the dev set.") |
|
parser.add_argument("--do_lower_case", action='store_true', |
|
help="Set this flag if you are using an uncased model.") |
|
parser.add_argument("--no_cuda", action='store_true', |
|
help="Avoid using CUDA when available") |
|
|
|
parser.add_argument("--train_batch_size", default=8, type=int, |
|
help="Batch size per GPU/CPU for training.") |
|
parser.add_argument("--eval_batch_size", default=8, 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("--beam_size", default=10, type=int, |
|
help="beam size for beam search") |
|
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("--save_steps", default=-1, type=int, ) |
|
parser.add_argument("--log_steps", default=-1, type=int, ) |
|
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("--eval_steps", default=-1, type=int, |
|
help="") |
|
parser.add_argument("--train_steps", default=-1, type=int, |
|
help="") |
|
parser.add_argument("--warmup_steps", default=100, type=int, |
|
help="Linear warmup over warmup_steps.") |
|
parser.add_argument("--local_rank", type=int, default=-1, |
|
help="For distributed training: local_rank") |
|
parser.add_argument('--seed', type=int, default=1234, |
|
help="random seed for initialization") |
|
args = parser.parse_args() |
|
|
|
if args.task in ['summarize']: |
|
args.lang = args.sub_task |
|
elif args.task in ['refine', 'concode', 'clone']: |
|
args.lang = 'java' |
|
elif args.task == 'defect': |
|
args.lang = 'c' |
|
elif args.task == 'translate': |
|
args.lang = 'c_sharp' if args.sub_task == 'java-cs' else 'java' |
|
return args |
|
|
|
|
|
def set_dist(args): |
|
|
|
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 |
|
cpu_cont = multiprocessing.cpu_count() |
|
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 |
|
args.cpu_cont = cpu_cont |
|
|
|
|
|
def set_seed(args): |
|
"""set random seed.""" |
|
random.seed(args.seed) |
|
np.random.seed(args.seed) |
|
torch.manual_seed(args.seed) |
|
if args.n_gpu > 0: |
|
torch.cuda.manual_seed_all(args.seed) |
|
|
