def main_cli(): parser = argparse.ArgumentParser('CEDR model training and validation') parser.add_argument('--model', choices=MODEL_MAP.keys(), default='vanilla_mbart') parser.add_argument('--datafiles', type=argparse.FileType('rt'), nargs='+') parser.add_argument('--qrels', type=argparse.FileType('rt')) parser.add_argument('--train_pairs', type=argparse.FileType('rt')) parser.add_argument('--valid_run', type=argparse.FileType('rt')) parser.add_argument('--initial_mbart_weights', type=argparse.FileType('rb')) parser.add_argument('--model_out_dir') args = parser.parse_args() model = MODEL_MAP[args.model]().cuda() dataset = data.read_datafiles(args.datafiles) qrels = data.read_qrels_dict(args.qrels) train_pairs = data.read_pairs_dict(args.train_pairs) valid_run = data.read_run_dict(args.valid_run) if args.initial_mbart_weights is not None: model.load(args.initial_mbart_weights.name) os.makedirs(args.model_out_dir, exist_ok=True) # we use the same qrels object for both training and validation sets main(model, dataset, train_pairs, qrels, valid_run, qrels, args.model_out_dir)
def main_cli(): parser = argparse.ArgumentParser('CEDR model training and validation') parser.add_argument('--model', choices=MODEL_MAP.keys(), default='vanilla_bert') parser.add_argument('--datafiles', nargs='+') # parser.add_argument('--datafiles', type=argparse.FileType('rt'), nargs='+') parser.add_argument('--qrels', type=argparse.FileType('rt')) parser.add_argument('--train_pairs', type=argparse.FileType('rt')) parser.add_argument('--valid_run', type=argparse.FileType('rt')) parser.add_argument('--initial_bert_weights', type=argparse.FileType('rb')) parser.add_argument('--model_out_dir') args = parser.parse_args() print('datafiles: ', args.datafiles) datafiles = [open(f, 'r', encoding='utf-8') for f in args.datafiles] model = MODEL_MAP[args.model]() # dataset = data.read_datafiles(args.datafiles) dataset = data.read_datafiles(datafiles) qrels = data.read_qrels_dict(args.qrels) train_pairs = data.read_pairs_dict(args.train_pairs) valid_run = data.read_run_dict(args.valid_run) if args.initial_bert_weights is not None: model.load(args.initial_bert_weights.name) os.makedirs(args.model_out_dir, exist_ok=True) main(model, dataset, train_pairs, qrels, valid_run, args.qrels.name, args.model_out_dir)
def predict(model, initial_bert_weights, datafiles, valid_run, query_len, doc_len): model = MODEL_MAP[model](query_len).cuda() if initial_bert_weights is not None: model.load(initial_bert_weights) data_f = open(datafiles, 'r', encoding='utf-8') valid_r = open(valid_run, 'r', encoding='utf-8') data.init_qd_len(query_len, doc_len) dataset = data.read_datafiles(data_f) run = data.read_qrels_dict(valid_r) runf = os.path.join('final.run') run_model(model, dataset, run, runf)
def main_cli(): parser = argparse.ArgumentParser('CEDR model training and validation') parser.add_argument('--model', choices=MODEL_MAP.keys(), default='vanilla_bert') parser.add_argument('--datafiles', type=argparse.FileType('rt'), nargs='+') parser.add_argument('--qrels', type=argparse.FileType('rt')) parser.add_argument('--train_pairs', type=argparse.FileType('rt')) parser.add_argument('--valid_run', type=argparse.FileType('rt')) parser.add_argument('--initial_bert_weights', type=argparse.FileType('rb')) parser.add_argument('--model_out_dir', default="../models/vbert") print(sys.argv) args = parser.parse_args(args=sys.argv[2:-1]) dataset = data.read_datafiles(args.datafiles) qrels = data.read_qrels_dict(args.qrels) train_pairs = data.read_pairs_dict(args.train_pairs) valid_run = data.read_run_dict(args.valid_run) model = MODEL_MAP[args.model]().cuda() if args.initial_bert_weights is not None: model.load(args.initial_bert_weights.name) os.makedirs(args.model_out_dir, exist_ok=True) main(model, dataset, train_pairs, qrels, valid_run, args.qrels, args.model_out_dir)
def main_cli(): parser = argparse.ArgumentParser('CEDR model training and validation') parser.add_argument('--model', choices=MODEL_MAP.keys()) parser.add_argument('--config', required=True, help='the model config file') #parser.add_argument('--toks_datafiles', required=True) # parser.add_argument('--valid_trec', type=argparse.FileType('rt')) # parser.add_argument('--train_pairs', type=argparse.FileType('rt')) # parser.add_argument('--valid_run', type=argparse.FileType('rt')) parser.add_argument('--initial_bert_weights', type=argparse.FileType('rb')) parser.add_argument( '--task_id', required=True, help= 'the task id for this training task. help for insert into our database.' ) parser.add_argument( '--model_id', required=True, help= 'the model id for this training task. help for insert into our database.' ) args = parser.parse_args() global task_id task_id = args.task_id global model_id model_id = args.model_id config = json.load(open(args.config, 'r')) data_dir = config["data_dir"] global base_metric base_metric = config["base_metric"] global all_metric all_metric = config["metrics"] global LR LR = config["learning_rate"] global MAX_EPOCH MAX_EPOCH = config["MAX_EPOCH"] global BATCH_SIZE BATCH_SIZE = config["BATCH_SIZE"] global BATCHES_PER_EPOCH BATCHES_PER_EPOCH = config["BATCHES_PER_EPOCH"] args.text1_maxlen = config["text1_maxlen"] args.text2_maxlen = config["text2_maxlen"] data.init_qd_len(args.text1_maxlen, args.text2_maxlen) args.toks_datafiles = open(data_dir + 'corpus_preprocessed.txt', 'r', encoding='utf-8') args.valid_trec = open(data_dir + 'relation_valid_trec.txt', 'r', encoding='utf-8') args.train_pairs = open(data_dir + 'relation_train.txt', 'r', encoding='utf-8') args.valid_run = open(data_dir + 'relation_valid.txt', 'r', encoding='utf-8') args.model_out_dir = config["model_out_dir"] model = MODEL_MAP[args.model](args.config).cuda() toks_dataset = data.read_toks(args.toks_datafiles) '''qrels = data.read_qrels_dict(args.qrels) train_pairs = data.read_pairs_dict(args.train_pairs) valid_run = data.read_run_dict(args.valid_run)''' train_pairs = data.read_qrels_dict(args.train_pairs) valid_run = data.read_qrels_dict(args.valid_run) if args.initial_bert_weights is not None: model.load(args.initial_bert_weights.name) os.makedirs(args.model_out_dir, exist_ok=True) #main(model, dataset, train_pairs, qrels, valid_run, args.qrels.name, args.model_out_dir) main(model, toks_dataset, train_pairs, valid_run, args.valid_trec.name, args.model_out_dir)
def main_cli(): MODEL_MAP = modeling.MODEL_MAP parser = argparse.ArgumentParser('CEDR model training and validation') parser.add_argument('--model', choices=MODEL_MAP.keys(), default='vbert') parser.add_argument('--submodel1', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel2', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel3', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel4', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel5', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel6', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel7', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--submodel8', choices=MODEL_MAP.keys(), default=None) parser.add_argument('--datafiles', type=argparse.FileType('rt'), nargs='+') parser.add_argument('--qrels', type=argparse.FileType('rt')) parser.add_argument('--train_pairs', type=argparse.FileType('rt')) parser.add_argument('--valid_run', type=argparse.FileType('rt')) parser.add_argument('--initial_bert_weights', type=str, default=None) parser.add_argument('--model_out_dir') parser.add_argument('--max_epoch', type=int, default=100, help='max epoch') parser.add_argument('--warmup_epoch', type=int, default=0, help='warmup epoch') parser.add_argument('--gpunum', type=str, default="0", help='gpu number') parser.add_argument('--random_seed', type=int, default=42, help='ranodm seed number') parser.add_argument('--freeze_bert', type=int, default=0, help='freezing bert') args = parser.parse_args() setRandomSeed(args.random_seed) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpunum print("GPU count=", torch.cuda.device_count()) if (args.model.startswith('duet')): model = MODEL_MAP[args.model](MODEL_MAP[args.submodel1](), MODEL_MAP[args.submodel2]()) elif (args.model.startswith('quad')): model = MODEL_MAP[args.model](MODEL_MAP[args.submodel1](), MODEL_MAP[args.submodel2](), MODEL_MAP[args.submodel3](), MODEL_MAP[args.submodel4]()) else: model = MODEL_MAP[args.model]().cuda() dataset = data.read_datafiles(args.datafiles) qrels = data.read_qrels_dict(args.qrels) train_pairs = data.read_pairs_dict(args.train_pairs) valid_run = data.read_run_dict(args.valid_run) ## initial if (args.initial_bert_weights is not None): wts = args.initial_bert_weights.split(',') if (len(wts) == 1): model.load(wts[0]) elif (len(wts) == 2): model.load_duet(wts[0], wts[1]) elif (len(wts) == 4): model.load_quad(wts[0], wts[1], wts[2], wts[3]) os.makedirs(args.model_out_dir, exist_ok=True) main(model, dataset, train_pairs, qrels, valid_run, args.qrels.name, args.model_out_dir, args.max_epoch, args.warmup_epoch, args.freeze_bert)