def main():
logger.info('Launching the MT-DNN training')
opt = vars(args)
# update data dir
opt['data_dir'] = data_dir
batch_size = args.batch_size
train_data_list = []
tasks = {}
tasks_class = {}
nclass_list = []
decoder_opts = []
dropout_list = []
for dataset in args.train_datasets:
prefix = dataset.split('_')[0]
if prefix in tasks: continue
assert prefix in DATA_META
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
nclass = DATA_META[prefix]
task_id = len(tasks)
if args.mtl_opt > 0:
task_id = tasks_class[nclass] if nclass in tasks_class else len(
tasks_class)
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
dopt = generate_decoder_opt(prefix, opt['answer_opt'])
if task_id < len(decoder_opts):
decoder_opts[task_id] = min(decoder_opts[task_id], dopt)
else:
decoder_opts.append(dopt)
if prefix not in tasks:
tasks[prefix] = len(tasks)
if args.mtl_opt < 1: nclass_list.append(nclass)
if (nclass not in tasks_class):
tasks_class[nclass] = len(tasks_class)
if args.mtl_opt > 0: nclass_list.append(nclass)
dropout_p = args.dropout_p
if tasks_config and prefix in tasks_config:
dropout_p = tasks_config[prefix]
dropout_list.append(dropout_p)
train_path = os.path.join(data_dir, '{}_train.json'.format(dataset))
logger.info('Loading {} as task {}'.format(train_path, task_id))
train_data = BatchGen(BatchGen.load(train_path,
True,
pairwise=pw_task,
maxlen=args.max_seq_len),
batch_size=batch_size,
dropout_w=args.dropout_w,
gpu=args.cuda,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
train_data_list.append(train_data)
opt['answer_opt'] = decoder_opts
opt['tasks_dropout_p'] = dropout_list
args.label_size = ','.join([str(l) for l in nclass_list])
logger.info(args.label_size)
dev_data_list = []
test_data_list = []
for dataset in args.test_datasets:
prefix = dataset.split('_')[0]
if args.mtl_opt > 0:
task_id = tasks_class[DATA_META[prefix]]
else:
task_id = tasks[prefix]
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
dev_path = os.path.join(data_dir, '{}_dev.json'.format(dataset))
dev_data = None
if os.path.exists(dev_path):
dev_data = BatchGen(BatchGen.load(dev_path,
False,
pairwise=pw_task,
maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda,
is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
dev_data_list.append(dev_data)
test_path = os.path.join(data_dir, '{}_test.json'.format(dataset))
test_data = None
if os.path.exists(test_path):
test_data = BatchGen(BatchGen.load(test_path,
False,
pairwise=pw_task,
maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda,
is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
test_data_list.append(test_data)
logger.info('#' * 20)
logger.info(opt)
logger.info('#' * 20)
all_iters = [iter(item) for item in train_data_list]
all_lens = [len(bg) for bg in train_data_list]
num_all_batches = args.epochs * sum(all_lens)
if len(train_data_list) > 1 and args.ratio > 0:
num_all_batches = int(args.epochs * (len(train_data_list[0]) *
(1 + args.ratio)))
model_path = args.init_checkpoint
state_dict = None
if os.path.exists(model_path):
state_dict = torch.load(model_path)
config = state_dict['config']
config['attention_probs_dropout_prob'] = args.bert_dropout_p
config['hidden_dropout_prob'] = args.bert_dropout_p
opt.update(config)
else:
logger.error('#' * 20)
logger.error(
'Could not find the init model!\n The parameters will be initialized randomly!'
)
logger.error('#' * 20)
config = BertConfig(vocab_size_or_config_json_file=30522).to_dict()
opt.update(config)
model = MTDNNModel(opt,
state_dict=state_dict,
num_train_step=num_all_batches)
####model meta str
headline = '############# Model Arch of MT-DNN #############'
###print network
logger.info('\n{}\n{}\n'.format(headline, model.network))
# dump config
config_file = os.path.join(output_dir, 'config.json')
with open(config_file, 'w', encoding='utf-8') as writer:
writer.write('{}\n'.format(json.dumps(opt)))
writer.write('\n{}\n{}\n'.format(headline, model.network))
logger.info("Total number of params: {}".format(model.total_param))
if args.freeze_layers > 0:
model.network.freeze_layers(args.freeze_layers)
if args.cuda:
model.cuda()
for epoch in range(0, args.epochs):
logger.warning('At epoch {}'.format(epoch))
for train_data in train_data_list:
train_data.reset()
start = datetime.now()
all_indices = []
if len(train_data_list) > 1 and args.ratio > 0:
main_indices = [0] * len(train_data_list[0])
extra_indices = []
for i in range(1, len(train_data_list)):
extra_indices += [i] * len(train_data_list[i])
random_picks = int(
min(len(train_data_list[0]) * args.ratio, len(extra_indices)))
extra_indices = np.random.choice(extra_indices,
random_picks,
replace=False)
if args.mix_opt > 0:
extra_indices = extra_indices.tolist()
random.shuffle(extra_indices)
all_indices = extra_indices + main_indices
else:
all_indices = main_indices + extra_indices.tolist()
else:
for i in range(1, len(train_data_list)):
all_indices += [i] * len(train_data_list[i])
if args.mix_opt > 0:
random.shuffle(all_indices)
all_indices += [0] * len(train_data_list[0])
if args.mix_opt < 1:
random.shuffle(all_indices)
for i in range(len(all_indices)):
task_id = all_indices[i]
batch_meta, batch_data = next(all_iters[task_id])
model.update(batch_meta, batch_data)
if (model.updates
) % args.log_per_updates == 0 or model.updates == 1:
logger.info(
'Task [{0:2}] updates[{1:6}] train loss[{2:.5f}] remaining[{3}]'
.format(
task_id, model.updates, model.train_loss.avg,
str((datetime.now() - start) / (i + 1) *
(len(all_indices) - i - 1)).split('.')[0]))
for idx, dataset in enumerate(args.test_datasets):
prefix = dataset.split('_')[0]
label_dict = GLOBAL_MAP.get(prefix, None)
dev_data = dev_data_list[idx]
if dev_data is not None:
dev_metrics, dev_predictions, scores, golds, dev_ids = eval_model(
model, dev_data, dataset=prefix, use_cuda=args.cuda)
for key, val in dev_metrics.items():
logger.warning(
"Task {0} -- epoch {1} -- Dev {2}: {3:.3f}".format(
dataset, epoch, key, val))
score_file = os.path.join(
output_dir, '{}_dev_scores_{}.json'.format(dataset, epoch))
results = {
'metrics': dev_metrics,
'predictions': dev_predictions,
'uids': dev_ids,
'scores': scores
}
dump(score_file, results)
official_score_file = os.path.join(
output_dir, '{}_dev_scores_{}.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
# test eval
test_data = test_data_list[idx]
if test_data is not None:
# For eval_model function, with_label = True specifies that evaluation metrics will be reported for test data -
# this was presumably disabled by authors as it is bad practice in hyperparameter tuning, however it is the most convenient
# way to get test scores. To avoid bias, hyperparameter decisions are made based on dev evaluation metrics, and test metrics
# are only recorded for the final versions of models.
test_metrics, test_predictions, scores, golds, test_ids = eval_model(
model,
test_data,
dataset=prefix,
use_cuda=args.cuda,
with_label=True)
score_file = os.path.join(
output_dir,
'{}_test_scores_{}.json'.format(dataset, epoch))
results = {
'metrics': test_metrics,
'predictions': test_predictions,
'uids': test_ids,
'scores': scores
}
dump(score_file, results)
official_score_file = os.path.join(
output_dir, '{}_test_scores_{}.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
logger.info('[new test scores saved.]')
model_file = os.path.join(output_dir, 'model_{}.pt'.format(epoch))
model.save(model_file)
def main():
logger.info('Launching the MT-DNN training')
opt = vars(args)
# update data dir
opt['data_dir'] = data_dir
batch_size = args.batch_size
train_data_list = []
tasks = {}
tasks_class = {}
nclass_list = []
decoder_opts = []
dropout_list = []
for dataset in args.train_datasets:
prefix = dataset.split('_')[0]
if prefix in tasks: continue
assert prefix in DATA_META
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
nclass = DATA_META[prefix]
task_id = len(tasks)
if args.mtl_opt > 0:
task_id = tasks_class[nclass] if nclass in tasks_class else len(
tasks_class)
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
dopt = generate_decoder_opt(prefix, opt['answer_opt'])
if task_id < len(decoder_opts):
decoder_opts[task_id] = min(decoder_opts[task_id], dopt)
else:
decoder_opts.append(dopt)
if prefix not in tasks:
tasks[prefix] = len(tasks)
if args.mtl_opt < 1: nclass_list.append(nclass)
if (nclass not in tasks_class):
tasks_class[nclass] = len(tasks_class)
if args.mtl_opt > 0: nclass_list.append(nclass)
dropout_p = args.dropout_p
if tasks_config and prefix in tasks_config:
dropout_p = tasks_config[prefix]
dropout_list.append(dropout_p)
train_data_ratio_string = str(
args.train_data_ratio) + "p" if args.train_data_ratio < 100 else ""
train_path = os.path.join(
data_dir, '{0}_train{1}.json'.format(dataset,
train_data_ratio_string))
logger.info('Loading {} as task {}'.format(train_path, task_id))
train_data = BatchGen(BatchGen.load(train_path,
True,
pairwise=pw_task,
maxlen=args.max_seq_len),
batch_size=batch_size,
dropout_w=args.dropout_w,
gpu=args.cuda,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
train_data_list.append(train_data)
opt['answer_opt'] = decoder_opts
opt['tasks_dropout_p'] = dropout_list
args.label_size = ','.join([str(l) for l in nclass_list])
logger.info(args.label_size)
dev_data_list = []
test_data_list = []
for dataset in args.test_datasets:
prefix = dataset.split('_')[0]
task_id = tasks_class[
DATA_META[prefix]] if args.mtl_opt > 0 else tasks[prefix]
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
dev_path = os.path.join(data_dir, '{}_dev.json'.format(dataset))
dev_data = None
if os.path.exists(dev_path):
dev_data = BatchGen(BatchGen.load(dev_path,
False,
pairwise=pw_task,
maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda,
is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
dev_data_list.append(dev_data)
test_path = os.path.join(data_dir, '{}_test.json'.format(dataset))
test_data = None
if os.path.exists(test_path):
test_data = BatchGen(BatchGen.load(test_path,
False,
pairwise=pw_task,
maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda,
is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
test_data_list.append(test_data)
logger.info('#' * 20)
logger.info(opt)
logger.info('#' * 20)
all_iters = [iter(item) for item in train_data_list]
all_lens = [len(bg) for bg in train_data_list]
num_all_batches = args.epochs * sum(all_lens)
if len(train_data_list) > 1 and args.ratio > 0:
num_all_batches = int(args.epochs * (len(train_data_list[0]) *
(1 + args.ratio)))
model_path = args.init_checkpoint
state_dict = None
if os.path.exists(model_path):
state_dict = torch.load(model_path, map_location='cpu')
config = state_dict['config']
config['attention_probs_dropout_prob'] = args.bert_dropout_p
config['hidden_dropout_prob'] = args.bert_dropout_p
opt.update(config)
else:
logger.error('#' * 20)
logger.error('Could not find the init model!\n Exit application!')
logger.error('#' * 20)
try:
shutil.rmtree(output_dir)
except Exception as e:
print(e)
exit(1)
model = MTDNNModel(opt,
state_dict=state_dict,
num_train_step=num_all_batches)
####model meta str
headline = '############# Model Arch of MT-DNN #############'
###print network
logger.info('\n{}\n{}\n'.format(headline, model.network))
# dump config
config_file = os.path.join(output_dir, 'config.json')
with open(config_file, 'w', encoding='utf-8') as writer:
writer.write('{}\n'.format(json.dumps(opt)))
writer.write('\n{}\n{}\n'.format(headline, model.network))
logger.info("Total number of params: {}".format(model.total_param))
if args.freeze_layers > 0:
model.network.freeze_layers(args.freeze_layers)
if args.cuda:
model.cuda()
best_F1_macro = -1.0
for epoch in range(0, args.epochs):
logger.warning('At epoch {}'.format(epoch))
for train_data in train_data_list:
train_data.reset()
start = datetime.now()
all_indices = []
if len(train_data_list) > 1 and (args.ratio > 0 or
args.reduce_first_dataset_ratio > 0):
main_indices = [0] * (int(args.reduce_first_dataset_ratio * len(
train_data_list[0])) if args.reduce_first_dataset_ratio > 0
else len(train_data_list[0]))
extra_indices = []
for i in range(1, len(train_data_list)):
extra_indices += [i] * len(train_data_list[i])
if args.ratio > 0:
random_picks = int(
min(
len(train_data_list[0]) * args.ratio,
len(extra_indices)))
extra_indices = np.random.choice(extra_indices,
random_picks,
replace=False).tolist()
if args.mix_opt > 0:
extra_indices = extra_indices
random.shuffle(extra_indices)
all_indices = extra_indices + main_indices
else:
all_indices = main_indices + extra_indices
logger.info(
"Main batches loaded (first dataset in list): {}".format(
len(main_indices)))
logger.info(
"Extra batches loaded (all except first dataset in list): {}".
format(len(extra_indices)))
else: # shuffle the index of the train sets whose batches will be trained on in the order: e.g. if train_set[1] is large, it will get trained on more often
for i in range(1, len(train_data_list)):
all_indices += [i] * len(train_data_list[i])
if args.mix_opt > 0:
random.shuffle(all_indices)
all_indices += [0] * len(train_data_list[0])
if args.mix_opt < 1:
random.shuffle(all_indices)
for i in range(len(all_indices)):
task_id = all_indices[i]
batch_meta, batch_data = next(all_iters[task_id])
model.update(batch_meta, batch_data)
if (model.updates
) % args.log_per_updates == 0 or model.updates == 1:
logger.info(
'Task [{0:2}] updates[{1:6}] train loss[{2:.5f}] remaining[{3}]'
.format(
task_id, model.updates, model.train_loss.avg,
str((datetime.now() - start) / (i + 1) *
(len(all_indices) - i - 1)).split('.')[0]))
temp_dev_F1s = []
dev_dump_list = []
test_dump_list = []
for idx, dataset in enumerate(args.test_datasets):
prefix = dataset.split('_')[0]
label_dict = GLOBAL_MAP.get(prefix, None)
dev_data = dev_data_list[idx]
if dev_data is not None:
dev_metrics, dev_predictions, scores, golds, dev_ids, premises, hypotheses = eval_model(
model, dev_data, dataset=prefix, use_cuda=args.cuda)
for key, val in dev_metrics.items():
if not isinstance(val, dict):
logger.warning(
"Task {0} -- epoch {1} -- Dev {2}: {3:.3f}".format(
dataset, epoch, key, val))
score_file = os.path.join(
output_dir, '{}_dev_scores_{}.json'.format(dataset, epoch))
results = {
'metrics': dev_metrics,
'predictions': dev_predictions,
'uids': dev_ids,
'scores': scores,
'golds': golds,
'premises': premises,
'hypotheses': hypotheses
}
dump(score_file, results)
official_score_file = os.path.join(
output_dir, '{}_dev_scores_{}.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
# for checkpoint
temp_dev_F1s.append(dev_metrics['F1_macro'])
dev_dump_list.append({
"output_dir": output_dir,
"dev_metrics": dev_metrics,
"dev_predictions": dev_predictions,
"golds": golds,
"opt": opt,
"dataset": dataset
})
# test eval
test_data = test_data_list[idx]
if test_data is not None:
test_metrics, test_predictions, scores, golds, test_ids, premises, hypotheses = eval_model(
model,
test_data,
dataset=prefix,
use_cuda=args.cuda,
with_label=True)
score_file = os.path.join(
output_dir,
'{}_test_scores_{}.json'.format(dataset, epoch))
results = {
'metrics': test_metrics,
'predictions': test_predictions,
'uids': test_ids,
'scores': scores,
'golds': golds,
'premises': premises,
'hypotheses': hypotheses
}
dump(score_file, results)
official_score_file = os.path.join(
output_dir, '{}_test_scores_{}.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
logger.info('[new test scores saved.]')
# for checkpoint
test_dump_list.append({
"output_dir": output_dir,
"test_metrics": test_metrics,
"test_predictions": test_predictions,
"golds": golds,
"opt": opt,
"dataset": dataset
})
# save checkpoint
if np.average(temp_dev_F1s) > best_F1_macro:
print("Save new model! Current best F1 macro over all dev sets: " +
"{0:.2f}".format(best_F1_macro) + ". New: " +
"{0:.2f}".format(np.average(temp_dev_F1s)))
best_F1_macro = np.average(temp_dev_F1s)
# override current dump file
for l in dev_dump_list:
dump_result_files(l['dataset'])(l['output_dir'], epoch,
l['dev_metrics'],
str(l['dev_predictions']),
str(l['golds']), "dev",
l['opt'], l['dataset'])
for l in test_dump_list:
dump_result_files(l['dataset'])(l['output_dir'], epoch,
l['test_metrics'],
str(l['test_predictions']),
str(l['golds']), "test",
l['opt'], l['dataset'])
# save model
model_file = os.path.join(output_dir, 'model.pt')
model.save(model_file)
def main():
logger.info('Launching the MT-DNN training')
opt = vars(args)
# update data dir
opt['data_dir'] = data_dir
batch_size = args.batch_size
train_data_list = []
tasks = {}
tasks_class = {}
nclass_list = []
dropout_list = []
for dataset in args.train_datasets:
prefix = dataset.split('_')[0]
if prefix in tasks: continue
assert prefix in DATA_META
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
nclass = DATA_META[prefix]
task_id = len(tasks)
if args.mtl_opt > 0:
task_id = tasks_class[nclass] if nclass in tasks_class else len(tasks_class)
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
if prefix not in tasks:
tasks[prefix] = len(tasks)
if args.mtl_opt < 1: nclass_list.append(nclass)
if (nclass not in tasks_class):
tasks_class[nclass] = len(tasks_class)
if args.mtl_opt > 0: nclass_list.append(nclass)
dropout_p = args.dropout_p
if tasks_config and prefix in tasks_config:
dropout_p = tasks_config[prefix]
dropout_list.append(dropout_p)
train_path = os.path.join(data_dir, '{}_train.json'.format(dataset))
logger.info('Loading {} as task {}'.format(train_path, task_id))
train_data = BatchGen(BatchGen.load(train_path, True, pairwise=pw_task, maxlen=args.max_seq_len, filter_long_parses=(prefix!='mednli')),
batch_size=batch_size,
dropout_w=args.dropout_w,
gpu=args.cuda,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type,
use_parse=args.use_parse,
use_generic_features=args.use_generic_features,
use_domain_features=args.use_domain_features,
feature_pkl_dir=args.feature_pkl_dir if args.feature_pkl_dir is not None else args.data_dir,
feature_pkl_namespace='train')
train_data_list.append(train_data)
opt['tasks_dropout_p'] = dropout_list
args.label_size = ','.join([str(l) for l in nclass_list])
logger.info(args.label_size)
dev_data_list = []
test_data_list = []
for dataset in args.test_datasets:
prefix = dataset.split('_')[0]
task_id = tasks_class[DATA_META[prefix]] if args.mtl_opt > 0 else tasks[prefix]
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
dev_path = os.path.join(data_dir, '{}_dev.json'.format(dataset))
dev_data = None
if os.path.exists(dev_path):
dev_data = BatchGen(BatchGen.load(dev_path, False, pairwise=pw_task, maxlen=args.max_seq_len, filter_long_parses=False),
batch_size=args.batch_size_eval,
gpu=args.cuda, is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type,
use_parse=args.use_parse,
use_generic_features=args.use_generic_features,
use_domain_features=args.use_domain_features,
feature_pkl_dir=args.feature_pkl_dir if args.feature_pkl_dir is not None else args.data_dir,
feature_pkl_namespace='dev')
dev_data_list.append(dev_data)
test_path = os.path.join(data_dir, '{}_test.json'.format(dataset))
test_data = None
if os.path.exists(test_path):
test_data = BatchGen(BatchGen.load(test_path, False, pairwise=pw_task, maxlen=args.max_seq_len, filter_long_parses=False),
batch_size=args.batch_size_eval,
gpu=args.cuda, is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type,
use_parse=args.use_parse,
use_generic_features=args.use_generic_features,
use_domain_features=args.use_domain_features,
feature_pkl_dir=args.feature_pkl_dir if args.feature_pkl_dir is not None else args.data_dir,
feature_pkl_namespace='test')
test_data_list.append(test_data)
logger.info('#' * 20)
logger.info(opt)
logger.info('#' * 20)
model_path = args.init_checkpoint
state_dict = None
pretrained_embeddings = pretrained_idx2token = None
if os.path.exists(model_path):
state_dict = torch.load(model_path)
state_dict.pop('optimizer', None)
config = state_dict['config']
config['attention_probs_dropout_prob'] = args.bert_dropout_p
config['hidden_dropout_prob'] = args.bert_dropout_p
for k in {'epochs', 'output_dir', 'train_datasets', 'test_datasets', 'seed', 'local_model_idx2token',
'use_parse', 'stx_parse_dim', 'glove_path', 'unk_threshold', 'use_generic_features', 'use_domain_features',
'feature_dim', 'feature_pkl_dir'}:
config.pop(k, None)
opt.update(config)
if 'treelstm.embedding.weight' in state_dict['state']:
pretrained_embeddings = state_dict['state']['treelstm.embedding.weight']
pretrained_idx2token = pickle.load(open(args.local_model_idx2token, 'rb'))
del state_dict['state']['treelstm.embedding.weight']
else:
logger.error('#' * 20)
logger.error('Could not find the init model!\n The parameters will be initialized randomly!')
logger.error('#' * 20)
config = BertConfig(vocab_size_or_config_json_file=30522).to_dict()
opt.update(config)
assert len(train_data_list) == len(dev_data_list) == len(test_data_list) == 1
embedding_matrix = token2idx = unked_words = None
if args.use_parse:
assert args.stx_parse_dim is not None
assert args.glove_path is not None
vocab = Counter()
for data in train_data_list:
for batch in data.data:
for example in batch:
for leaf in Tree.from_char_indices(example['parse_id_a']).leaves(): vocab[leaf.content.lower()] += 1
for leaf in Tree.from_char_indices(example['parse_id_b']).leaves(): vocab[leaf.content.lower()] += 1
for data in dev_data_list:
for batch in data.data:
for example in batch:
for leaf in Tree.from_char_indices(example['parse_id_a']).leaves(): vocab[leaf.content.lower()] += 1
for leaf in Tree.from_char_indices(example['parse_id_b']).leaves(): vocab[leaf.content.lower()] += 1
final_vocab = {'<unk>'}
unked_words = set()
for word, count in vocab.items():
(final_vocab if count >= args.unk_threshold else unked_words).add(word)
assert len(final_vocab) + len(unked_words) == len(vocab) + 1
vocab = final_vocab
idx2token = {}
for token in vocab:
idx2token[len(idx2token)] = token
pickle.dump(idx2token, open(os.path.join(args.output_dir, "idx2token.pkl"), 'wb'))
token2idx = {token: idx for idx, token in idx2token.items()}
embedding_matrix = load_embeddings(args.glove_path, vocab, idx2token, pretrained_embeddings=pretrained_embeddings, pretrained_idx2token=pretrained_idx2token)
num_generic_features = num_domain_features = None
if args.use_generic_features:
num_generic_features = len(train_data_list[0].data[0][0]['generic_features'])
if args.use_domain_features:
num_domain_features = len(train_data_list[0].data[0][0]['domain_features'])
all_iters =[iter(item) for item in train_data_list]
all_lens = [len(bg) for bg in train_data_list]
num_all_batches = args.epochs * sum(all_lens)
if len(train_data_list)> 1 and args.ratio > 0:
num_all_batches = int(args.epochs * (len(train_data_list[0]) * (1 + args.ratio)))
model = MTDNNModel(opt, state_dict=state_dict, num_train_step=num_all_batches,
use_parse=args.use_parse, embedding_matrix=embedding_matrix,
token2idx=token2idx, stx_parse_dim=args.stx_parse_dim, unked_words=unked_words,
use_generic_features=args.use_generic_features, num_generic_features=num_generic_features,
use_domain_features=args.use_domain_features, num_domain_features=num_domain_features, feature_dim=args.feature_dim)
####model meta str
headline = '############# Model Arch of MT-DNN #############'
###print network
logger.info('\n{}\n{}\n'.format(headline, model.network))
# dump config
config_file = os.path.join(output_dir, 'config.json')
with open(config_file, 'w', encoding='utf-8') as writer:
writer.write('{}\n'.format(json.dumps(opt)))
writer.write('\n{}\n{}\n'.format(headline, model.network))
logger.info("Total number of params: {}".format(model.total_param))
if args.freeze_layers > 0:
model.network.freeze_layers(args.freeze_layers)
highest_dev_acc = -1
if args.cuda:
model.cuda()
for epoch in range(0, args.epochs):
logger.warning('At epoch {}'.format(epoch))
for train_data in train_data_list:
train_data.reset()
start = datetime.now()
all_indices=[]
if len(train_data_list)> 1 and args.ratio > 0:
main_indices =[0] * len(train_data_list[0])
extra_indices=[]
for i in range(1, len(train_data_list)):
extra_indices += [i] * len(train_data_list[i])
random_picks=int(min(len(train_data_list[0]) * args.ratio, len(extra_indices)))
extra_indices = np.random.choice(extra_indices, random_picks, replace=False)
if args.mix_opt > 0:
extra_indices = extra_indices.tolist()
random.shuffle(extra_indices)
all_indices = extra_indices + main_indices
else:
all_indices = main_indices + extra_indices.tolist()
else:
for i in range(1, len(train_data_list)):
all_indices += [i] * len(train_data_list[i])
if args.mix_opt > 0:
random.shuffle(all_indices)
all_indices += [0] * len(train_data_list[0])
if args.mix_opt < 1:
random.shuffle(all_indices)
for i in range(len(all_indices)):
task_id = all_indices[i]
batch_meta, batch_data= next(all_iters[task_id])
model.update(batch_meta, batch_data)
if (model.updates) % args.log_per_updates == 0 or model.updates == 1:
logger.info('Task [{0:2}] updates[{1:6}] train loss[{2:.5f}] remaining[{3}]'.format(task_id,
model.updates, model.train_loss.avg,
str((datetime.now() - start) / (i + 1) * (len(all_indices) - i - 1)).split('.')[0]))
assert len(args.test_datasets) == 1
for idx, dataset in enumerate(args.test_datasets):
prefix = dataset.split('_')[0]
label_dict = GLOBAL_MAP.get(prefix, None)
dev_data = dev_data_list[idx]
if dev_data is not None:
dev_metrics, dev_predictions, scores, golds, dev_ids= eval_model(model, dev_data, dataset=prefix,
use_cuda=args.cuda)
assert len(dev_metrics) == 1
for key, val in dev_metrics.items():
logger.warning("Task {0} -- epoch {1} -- Dev {2}: {3:.3f}".format(dataset, epoch, key, val))
score_file = os.path.join(output_dir, '{}_dev_scores_{}.json'.format(dataset, epoch))
results = {'metrics': dev_metrics, 'predictions': dev_predictions, 'uids': dev_ids, 'scores': scores}
dump(score_file, results)
official_score_file = os.path.join(output_dir, '{}_dev_scores_{}.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
if list(dev_metrics.values())[0] > highest_dev_acc:
model.save(os.path.join(output_dir, 'best_model.pt'))
highest_dev_acc = list(dev_metrics.values())[0]
logger.warning(f'Best dev {highest_dev_acc}')
def main():
logger.info('Launching the MT-DNN training')
opt = vars(args)
# update data dir
opt['data_dir'] = data_dir
batch_size = args.batch_size
train_data_list = []
tasks = {}
tasks_class = {}
nclass_list = []
decoder_opts = []
dropout_list = []
for dataset in args.train_datasets:
prefix = dataset.split('_')[0]
if prefix in tasks: continue
assert prefix in DATA_META
assert prefix in DATA_TYPE
nclass = DATA_META[prefix]
task_id = len(tasks)
if args.mtl_opt > 0:
task_id = tasks_class[nclass] if nclass in tasks_class else len(tasks_class)
dopt = generate_decoder_opt(prefix, opt['answer_opt'])
if task_id < len(decoder_opts):
decoder_opts[task_id] = min(decoder_opts[task_id], dopt)
else:
decoder_opts.append(dopt)
if prefix not in tasks:
tasks[prefix] = len(tasks)
if args.mtl_opt < 1: nclass_list.append(nclass)
if (nclass not in tasks_class):
tasks_class[nclass] = len(tasks_class)
if args.mtl_opt > 0: nclass_list.append(nclass)
dropout_p = args.dropout_p
if tasks_config and prefix in tasks_config:
dropout_p = tasks_config[prefix]
dropout_list.append(dropout_p)
opt['answer_opt'] = decoder_opts
opt['tasks_dropout_p'] = dropout_list
args.label_size = ','.join([str(l) for l in nclass_list])
logger.info(args.label_size)
dev_data_list = []
test_data_list = []
stress_data_list = []
for dataset in args.test_datasets:
prefix = dataset.split('_')[0]
task_id = tasks_class[DATA_META[prefix]] if args.mtl_opt > 0 else tasks[prefix]
task_type = TASK_TYPE[prefix]
pw_task = False
if prefix in opt['pw_tasks']:
pw_task = True
assert prefix in DATA_TYPE
data_type = DATA_TYPE[prefix]
dev_path = os.path.join(data_dir, '{}_dev.json'.format(dataset))
dev_data = None
if os.path.exists(dev_path):
dev_data = BatchGen(BatchGen.load(dev_path, False, pairwise=pw_task, maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda, is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
dev_data_list.append(dev_data)
test_path = os.path.join(data_dir, '{}_test.json'.format(dataset))
test_data = None
if os.path.exists(test_path):
test_data = BatchGen(BatchGen.load(test_path, False, pairwise=pw_task, maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda, is_train=False,
task_id=task_id,
maxlen=args.max_seq_len,
pairwise=pw_task,
data_type=data_type,
task_type=task_type)
test_data_list.append(test_data)
stress_data = []
if args.stress_tests != "NONE":
for stress_test in args.stress_tests.split(','):
stress_path = os.path.join(data_dir, '{}_test_{}.json'.format(dataset, stress_test))
if os.path.exists(stress_path):
stress_data.append(BatchGen(BatchGen.load(stress_path, False, pairwise=pw_task, maxlen=args.max_seq_len),
batch_size=args.batch_size_eval,
gpu=args.cuda, is_train=False,
task_id=task_id,
maxlen=512,
pairwise=pw_task,
data_type=data_type,
task_type=task_type) )
stress_data_list.append(stress_data)
logger.info('#' * 20)
logger.info(opt)
logger.info('#' * 20)
all_lens = [len(bg) for bg in train_data_list]
num_all_batches = args.epochs * sum(all_lens)
if len(train_data_list) > 1 and args.ratio > 0:
num_all_batches = int(args.epochs * (len(train_data_list[0]) * (1 + args.ratio)))
model_path = args.init_checkpoint
state_dict = None
if os.path.exists(model_path):
state_dict = torch.load(model_path)
config = state_dict['config']
config['attention_probs_dropout_prob'] = args.bert_dropout_p
config['hidden_dropout_prob'] = args.bert_dropout_p
opt.update(config)
else:
logger.error('#' * 20)
logger.error('Could not find the init model!\n Exit application!')
logger.error('#' * 20)
model = MTDNNModel(opt, state_dict=state_dict, num_train_step=num_all_batches)
####model meta str
headline = '############# Model Arch of MT-DNN #############'
###print network
logger.info('\n{}\n{}\n'.format(headline, model.network))
# dump config
config_file = os.path.join(output_dir, 'config.json')
with open(config_file, 'w', encoding='utf-8') as writer:
writer.write('{}\n'.format(json.dumps(opt)))
writer.write('\n{}\n{}\n'.format(headline, model.network))
logger.info("Total number of params: {}".format(model.total_param))
if args.freeze_layers > 0:
model.network.freeze_layers(args.freeze_layers)
if args.cuda:
model.cuda()
for epoch in range(0, 1):
dev_dump_list = []
test_dump_list = []
stress_dump_list = []
for idx, dataset in enumerate(args.test_datasets):
prefix = dataset.split('_')[0]
label_dict = GLOBAL_MAP.get(prefix, None)
dev_data = dev_data_list[idx]
if dev_data is not None:
dev_metrics, dev_predictions, scores, golds, dev_ids, premises, hypotheses = eval_model(model, dev_data, dataset=prefix,
use_cuda=args.cuda)
for key, val in dev_metrics.items():
if not isinstance(val, dict):
logger.warning("Task {0} -- epoch {1} -- Dev {2}: {3:.3f}".format(dataset, epoch, key, val))
if args.dump_to_checkpoints == 1:
score_file = os.path.join(output_dir, '{}_dev_scores_{}_EVAL_ONLY.json'.format(dataset, epoch))
results = {'metrics': dev_metrics, 'predictions': dev_predictions, 'uids': dev_ids,
'scores': scores, 'golds': golds,
'premises': premises, 'hypotheses': hypotheses}
dump(score_file, results)
official_score_file = os.path.join(output_dir,
'{}_dev_scores_{}_EVAL_ONLY.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
# for checkpoint
dev_dump_list.append({
"output_dir": output_dir,
"dev_metrics": dev_metrics,
"dev_predictions": dev_predictions,
"golds": golds,
"opt": opt,
"dataset": dataset
})
# test eval
test_data = test_data_list[idx]
if test_data is not None:
test_metrics, test_predictions, scores, golds, test_ids, premises, hypotheses = eval_model(model, test_data, dataset=prefix,
use_cuda=args.cuda, with_label=True)
if args.dump_to_checkpoints == 1:
score_file = os.path.join(output_dir, '{}_test_scores_{}_EVAL_ONLY.json'.format(dataset, epoch))
results = {'metrics': test_metrics, 'predictions': test_predictions, 'uids': test_ids, 'scores': scores, 'golds': golds,
'premises': premises, 'hypotheses': hypotheses}
dump(score_file, results)
official_score_file = os.path.join(output_dir, '{}_test_scores_{}_EVAL_ONLY.tsv'.format(dataset, epoch))
submit(official_score_file, results, label_dict)
logger.info('[new test scores saved.]')
# for checkpoint
test_dump_list.append({
"output_dir": output_dir,
"test_metrics": test_metrics,
"test_predictions": test_predictions,
"golds": golds,
"opt": opt,
"dataset": dataset
})
# stress test eval
if args.stress_tests != "NONE":
stress_data = stress_data_list[idx]
for j, stress_test in enumerate(args.stress_tests.split(',')):
stress_metrics, stress_predictions, scores, golds, stress_ids, premises, hypotheses = \
eval_model(model, stress_data[j], dataset=prefix, use_cuda=args.cuda, with_label=True)
if args.dump_to_checkpoints == 1:
score_file = os.path.join(output_dir, '{}_test_{}_scores_{}_EVAL_ONLY.json'.format(dataset, stress_test, epoch))
results = {'metrics': stress_metrics, 'predictions': stress_predictions, 'uids': stress_ids, 'scores': scores, 'golds': golds,
'premises': premises, 'hypotheses': hypotheses}
dump(score_file, results)
official_score_file = os.path.join(output_dir, '{}_test_{}_scores_{}_EVAL_ONLY.tsv'.format(dataset, stress_test, epoch))
submit(official_score_file, results, label_dict)
logger.info('[new stress test scores for "{}" saved.]'.format(stress_test))
# for checkpoint
stress_dump_list.append({
"output_dir": output_dir,
"test_metrics": stress_metrics,
"test_predictions": stress_predictions,
"golds": golds,
"opt": opt,
"dataset": dataset,
"stress_test": stress_test
})
# save results
print("Save new results!")
for l in dev_dump_list:
dump_result_files(l['dataset'])(l['output_dir'], -1, l['dev_metrics'], str(l['dev_predictions']),
str(l['golds']), "dev", l['opt'], l['dataset'])
for l in test_dump_list:
dump_result_files(l['dataset'])(l['output_dir'], -1, l['test_metrics'], str(l['test_predictions']),
str(l['golds']), "test", l['opt'], l['dataset'])
if args.stress_tests != "NONE":
for l in stress_dump_list:
dump_result_files(l['dataset'])(l['output_dir'], -1, l['test_metrics'], str(l['test_predictions']),
str(l['golds']), l['stress_test'], l['opt'], l['dataset'])