def main():
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__)
args = get_args()
log_info.print_args(args)
device, n_gpu = initialization.init_cuda_from_args(args, logger=logger)
initialization.init_seed(args, n_gpu=n_gpu, logger=logger)
initialization.init_train_batch_size(args)
initialization.init_output_dir(args)
initialization.save_args(args)
task = get_task(args.task_name, args.data_dir)
tokenizer = shared_model_setup.create_tokenizer(
bert_model_name=args.bert_model,
bert_load_mode=args.bert_load_mode,
do_lower_case=args.do_lower_case,
bert_vocab_path=args.bert_vocab_path,
)
all_state = shared_model_setup.load_overall_state(args.bert_load_path, relaxed=True)
model = glue_model_setup.create_model(
task_type=task.processor.TASK_TYPE,
bert_model_name=args.bert_model,
bert_load_mode=args.bert_load_mode,
bert_load_args=args.bert_load_args,
all_state=all_state,
num_labels=len(task.processor.get_labels()),
device=device,
n_gpu=n_gpu,
fp16=args.fp16,
local_rank=args.local_rank,
bert_config_json_path=args.bert_config_json_path,
)
if args.do_train:
if args.print_trainable_params:
log_info.print_trainable_params(model)
train_examples = task.get_train_examples()
t_total = shared_model_setup.get_opt_train_steps(
num_train_examples=len(train_examples),
args=args,
)
optimizer = shared_model_setup.create_optimizer(
model=model,
learning_rate=args.learning_rate,
t_total=t_total,
loss_scale=args.loss_scale,
fp16=args.fp16,
warmup_proportion=args.warmup_proportion,
state_dict=all_state["optimizer"] if args.bert_load_mode == "state_all" else None,
)
else:
train_examples = None
t_total = 0
optimizer = None
runner = GlueTaskRunner(
model=model,
optimizer=optimizer,
tokenizer=tokenizer,
label_list=task.get_labels(),
device=device,
rparams=RunnerParameters(
max_seq_length=args.max_seq_length,
local_rank=args.local_rank, n_gpu=n_gpu, fp16=args.fp16,
learning_rate=args.learning_rate, gradient_accumulation_steps=args.gradient_accumulation_steps,
t_total=t_total, warmup_proportion=args.warmup_proportion,
num_train_epochs=args.num_train_epochs,
train_batch_size=args.train_batch_size, eval_batch_size=args.eval_batch_size,
)
)
if args.do_train:
assert at_most_one_of([args.do_val_history, args.train_save_every])
if args.do_val_history:
val_examples = task.get_dev_examples()
results = runner.run_train_val(
train_examples=train_examples,
val_examples=val_examples,
task_name=task.name,
)
metrics_str = json.dumps(results, indent=2)
with open(os.path.join(args.output_dir, "val_metrics_history.json"), "w") as f:
f.write(metrics_str)
elif args.train_save_every:
train_dataloader = runner.get_train_dataloader(train_examples, verbose=not args.not_verbose)
for epoch in range(int(args.num_train_epochs)):
for step, _, _ in runner.run_train_epoch_context(train_dataloader):
if step % args.train_save_every == args.train_save_every - 1 \
or step == len(train_dataloader) - 1:
glue_model_setup.save_bert(
model=model, optimizer=optimizer, args=args,
save_path=os.path.join(
args.output_dir, f"all_state___epoch{epoch:04d}___batch{step:06d}.p"
),
save_mode=args.bert_save_mode,
verbose=not args.not_verbose,
)
else:
runner.run_train(train_examples)
if args.do_save:
# Save a trained model
glue_model_setup.save_bert(
model=model, optimizer=optimizer, args=args,
save_path=os.path.join(args.output_dir, "all_state.p"),
save_mode=args.bert_save_mode,
)
if args.do_val:
val_examples = task.get_dev_examples()
results = runner.run_val(val_examples, task_name=task.name, verbose=not args.not_verbose)
df = pd.DataFrame(results["logits"])
df.to_csv(os.path.join(args.output_dir, "val_preds.csv"), header=False, index=False)
metrics_str = json.dumps({"loss": results["loss"], "metrics": results["metrics"]}, indent=2)
print(metrics_str)
with open(os.path.join(args.output_dir, "val_metrics.json"), "w") as f:
f.write(metrics_str)
# HACK for MNLI-mismatched
if task.name == "mnli":
mm_val_examples = MnliMismatchedProcessor().get_dev_examples(task.data_dir)
mm_results = runner.run_val(mm_val_examples, task_name=task.name, verbose=not args.not_verbose)
df = pd.DataFrame(results["logits"])
df.to_csv(os.path.join(args.output_dir, "mm_val_preds.csv"), header=False, index=False)
combined_metrics = {}
for k, v in results["metrics"]:
combined_metrics[k] = v
for k, v in mm_results["metrics"]:
combined_metrics["mm-"+k] = v
combined_metrics_str = json.dumps({
"loss": results["loss"],
"metrics": combined_metrics,
}, indent=2)
with open(os.path.join(args.output_dir, "val_metrics.json"), "w") as f:
f.write(combined_metrics_str)
if args.do_test:
test_examples = task.get_test_examples()
logits = runner.run_test(test_examples, verbose=not args.not_verbose)
df = pd.DataFrame(logits)
df.to_csv(os.path.join(args.output_dir, "test_preds.csv"), header=False, index=False)
# HACK for MNLI-mismatched
if task.name == "mnli":
test_examples = MnliMismatchedProcessor().get_test_examples(task.data_dir)
logits = runner.run_test(test_examples)
df = pd.DataFrame(logits)
df.to_csv(os.path.join(args.output_dir, "mm_test_preds.csv"), header=False, index=False)
def main():
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__)
args = get_args()
log_info.print_args(args)
device, n_gpu = initialization.init_cuda_from_args(args, logger=logger)
initialization.init_seed(args, n_gpu=n_gpu, logger=logger)
initialization.init_train_batch_size(args)
initialization.init_output_dir(args)
initialization.save_args(args)
task = get_task(args.task_name, args.data_dir)
# prepare examples, load model as encoder
tokenizer = shared_model_setup.create_tokenizer(
bert_model_name=args.bert_model,
bert_load_mode=args.bert_load_mode,
do_lower_case=args.do_lower_case,
bert_vocab_path=args.bert_vocab_path,
)
all_state = shared_model_setup.load_overall_state(args.bert_load_path,
relaxed=True)
# Load Model...
if args.bert_load_mode == "state_model_only":
state_dict = all_state['model']
bert_as_encoder = BertModel.from_state_dict(
config_file=args.bert_config_json_path, state_dict=state_dict)
else:
assert args.bert_load_mode == "from_pretrained"
cache_dir = PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(
args.local_rank)
bert_as_encoder = BertModel.from_pretrained(
pretrained_model_name_or_path=args.bert_model, cache_dir=cache_dir)
bert_as_encoder.to(device)
runner_param = RunnerParameters(
max_seq_length=args.max_seq_length,
local_rank=args.local_rank,
n_gpu=n_gpu,
fp16=args.fp16,
learning_rate=args.learning_rate,
gradient_accumulation_steps=args.gradient_accumulation_steps,
t_total=None,
warmup_proportion=args.warmup_proportion,
num_train_epochs=args.num_train_epochs,
train_batch_size=args.train_batch_size,
eval_batch_size=args.eval_batch_size,
)
runner = EmbeddingTaskRunner(bert_model=bert_as_encoder,
optimizer=None,
tokenizer=tokenizer,
label_list=task.get_labels(),
device=device,
rparams=runner_param)
# Run training set encoding...
print("Run training set encoding ... ")
train_examples = task.get_train_examples()
train_dataset = runner.run_encoding(train_examples,
verbose=True,
mode='train')
print("saving embeddings ... ")
torch.save(train_dataset, os.path.join(args.output_dir, "train.dataset"))
# Run development set encoding ...
eval_examples = task.get_dev_examples()
eval_dataset = runner.run_encoding(eval_examples,
verbose=True,
mode='eval')
print("saving embeddings ... ")
torch.save(eval_dataset, os.path.join(args.output_dir, 'dev.dataset'))
# Run test set encoding ...
test_examples = task.get_test_examples()
test_dataset = runner.run_encoding(test_examples,
verbose=True,
mode='test')
print("saving embeddings ... ")
torch.save(test_dataset, os.path.join(args.output_dir, "test.dataset"))
# HACK for MNLI mis-matched set ...
if args.task_name == 'mnli':
print("=== Start embedding task for MNLI mis-matched ===")
mm_eval_examples = MnliMismatchedProcessor().get_dev_examples(
task.data_dir)
mm_eval_dataset = runner.run_encoding(mm_eval_examples,
verbose=True,
mode='eval')
print("=== Saving eval dataset ===")
torch.save(mm_eval_dataset,
os.path.join(args.output_dir, "mm_dev.dataset"))
print("=== Saved ===")
mm_test_examples = MnliMismatchedProcessor().get_test_examples(
task.data_dir)
mm_test_dataset = runner.run_encoding(mm_test_examples,
verbose=True,
mode='test')
print("=== Saving tensor dataset ===")
torch.save(mm_test_dataset,
os.path.join(args.output_dir, "mm_test.dataset"))
print("=== Saved ===")
def test_net(teloader, model, args, num_labels, print_freq=100):
all_logits, all_labels, eval_loss = run_eval(teloader, model, print_freq,
num_labels)
task_name = args.task_name.lower()
result = compute_task_metrics(task_name, all_logits, all_labels)
print(' * Net Evaluation Acc is', result)
result_metric = {"eval_loss": eval_loss, "metrics": result}
if task_name == "mnli":
mm_val_examples = MnliMismatchedProcessor().get_dev_examples(
args.dataroot)
mm_val_label_list = MnliMismatchedProcessor().get_labels()
mm_val_label_map = {
label: i
for i, label in enumerate(mm_val_label_list)
}
mm_val_input_ids = []
mm_val_input_masks = []
mm_val_segment_ids = []
mm_val_label_ids = []
tokenizer = AutoTokenizer.from_pretrained(args.model)
for (ex_index, example) in enumerate(mm_val_examples):
mm_val_input_id, mm_val_input_mask, mm_val_segment_id, mm_val_label_id = \
convert_example_to_feature(example, tokenizer, args.max_seq_length, mm_val_label_map)
mm_val_input_ids.append(mm_val_input_id)
mm_val_input_masks.append(mm_val_input_mask)
mm_val_segment_ids.append(mm_val_segment_id)
mm_val_label_ids.append(mm_val_label_id)
mm_val_input_ids = torch.tensor(mm_val_input_ids)
mm_val_input_masks = torch.tensor(mm_val_input_masks)
mm_val_segment_ids = torch.tensor(mm_val_segment_ids)
mm_val_label_ids = torch.tensor(mm_val_label_ids)
mm_val_teset = torch.utils.data.TensorDataset(mm_val_input_ids,
mm_val_input_masks,
mm_val_segment_ids,
mm_val_label_ids)
mm_val_sampler = torch.utils.data.SequentialSampler(mm_val_teset)
mm_val_teloader = torch.utils.data.DataLoader(
mm_val_teset,
batch_size=args.batch_size,
sampler=mm_val_sampler,
num_workers=args.workers)
mm_val_logits, mm_val_labels, mm_eval_loss = run_eval(
mm_val_teloader, model, print_freq, num_labels)
print(' * Mis-matched task:')
mm_result = compute_task_metrics(task_name, mm_val_logits,
mm_val_labels)
print(' * Net Evaluation Acc is', mm_result)
result_metric["mm_eval_loss"] = mm_eval_loss
result_metric["mm_metrics"] = mm_result
return result_metric