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 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():
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)