model = nemo_nlp.huggingface.BERT(
pretrained_model_name=args.pretrained_bert_model)
model.restore_from(args.bert_checkpoint)
hidden_size = model.local_parameters["hidden_size"]
# uses [CLS] token for classification (the first token)
if args.task_name == 'sts-b':
pooler = nemo_nlp.SequenceRegression(hidden_size=hidden_size)
glue_loss = MSELoss()
else:
pooler = nemo_nlp.SequenceClassifier(hidden_size=hidden_size,
num_classes=num_labels,
log_softmax=False)
glue_loss = CrossEntropyLoss()
def create_pipeline(max_seq_length=args.max_seq_length,
batch_size=args.batch_size,
local_rank=args.local_rank,
num_gpus=args.num_gpus,
evaluate=False,
processor=task_processors[0]):
data_layer = 'GlueDataLayerClassification'
if output_mode == 'regression':
data_layer = 'GlueDataLayerRegression'
data_layer = getattr(sys.modules[__name__], data_layer)
def sentence_classification(args):
# TODO: construct name of experiment based on args
"""
name = construct_name(
args.exp_name,
args.lr,
args.batch_size,
args.num_epochs,
args.weight_decay,
args.optimizer)
work_dir = name
if args.work_dir:
work_dir = os.path.join(args.work_dir, name)
"""
# Instantiate neural modules
nf = NeuralModuleFactory(
backend=nemo.core.Backend.PyTorch,
local_rank=args.local_rank,
optimization_level=args.amp_opt_level,
log_dir=args.work_dir,
create_tb_writer=True,
files_to_copy=[__file__],
add_time_to_log_dir=True)
# Pre-trained BERT
tokenizer = BertTokenizer.from_pretrained(args.pretrained_bert_model)
if args.bert_checkpoint is None:
bert = nemo_nlp.BERT(pretrained_model_name=args.pretrained_bert_model)
# save bert config for inference after fine-tuning
bert_config = bert.config.to_dict()
with open(args.work_dir + '/' + args.pretrained_bert_model + '_config.json', 'w+') as json_file:
json.dump(bert_config, json_file)
else:
if args.bert_config is not None:
with open(args.bert_config) as json_file:
bert_config = json.load(json_file)
bert = nemo_nlp.BERT(**bert_config)
bert.restore_from(args.bert_checkpoint)
# MLP
bert_hidden_size = bert.local_parameters['hidden_size']
mlp = nemo_nlp.SequenceClassifier(
hidden_size=bert_hidden_size,
num_classes=args.num_classes,
num_layers=args.num_layers,
log_softmax=False,
dropout=args.dropout)
# TODO: save mlp/all model configs (bake in to Neural Module?)
if args.mlp_checkpoint:
mlp.restore_from(args.mlp_checkpoint)
# Loss function for classification
loss_fn = CrossEntropyLoss()
# Data layers, pipelines, and callbacks
callbacks = [] # callbacks depend on files present
if args.train_file:
if args.preproc:
train_data_layer = preproc_data_layer.PreprocBertSentenceClassificationDataLayer(
input_file=args.train_file,
shuffle=True,
num_samples=args.num_samples, # lower for dev, -1 for all dataset
batch_size=args.batch_size,
num_workers=0,
local_rank=args.local_rank)
else:
train_data_layer = nemo_nlp.BertSentenceClassificationDataLayer(
input_file=args.train_file,
tokenizer=tokenizer,
max_seq_length=args.max_seq_length,
shuffle=True,
num_samples=args.num_samples, # lower for dev, -1 for all dataset
batch_size=args.batch_size,
num_workers=0,
local_rank=args.local_rank)
train_logits, train_loss, steps_per_epoch, train_labels = create_pipeline(
nf,
train_data_layer,
bert,
mlp,
loss_fn)
train_callback = nemo.core.SimpleLossLoggerCallback(
tensors=[train_loss, train_logits],
print_func=lambda x: nf.logger.info(f'Train loss: {str(np.round(x[0].item(), 3))}'),
tb_writer=nf.tb_writer,
get_tb_values=lambda x: [["train_loss", x[0]]],
step_freq=steps_per_epoch)
callbacks.append(train_callback)
if args.num_checkpoints != 0:
ckpt_callback = nemo.core.CheckpointCallback(
folder=nf.checkpoint_dir,
epoch_freq=args.save_epoch_freq,
step_freq=args.save_step_freq,
checkpoints_to_keep=args.num_checkpoints)
callbacks.append(ckpt_callback)
if args.eval_file:
if args.preproc:
eval_data_layer = preproc_data_layer.PreprocBertSentenceClassificationDataLayer(
input_file=args.eval_file,
shuffle=False,
num_samples=args.num_samples,
batch_size=args.batch_size,
num_workers=0,
local_rank=args.local_rank)
else:
eval_data_layer = nemo_nlp.BertSentenceClassificationDataLayer(
input_file=args.eval_file,
tokenizer=tokenizer,
max_seq_length=args.max_seq_length,
shuffle=False,
num_samples=args.num_samples,
batch_size=args.batch_size,
num_workers=0,
local_rank=args.local_rank)
eval_logits, eval_loss, _, eval_labels = create_pipeline(
nf,
eval_data_layer,
bert,
mlp,
loss_fn)
eval_callback = nemo.core.EvaluatorCallback(
eval_tensors=[eval_logits, eval_labels],
user_iter_callback=lambda x, y: eval_iter_callback(
x, y, eval_data_layer),
user_epochs_done_callback=lambda x: eval_epochs_done_callback(
x, f'{nf.work_dir}/graphs'),
tb_writer=nf.tb_writer,
eval_step=steps_per_epoch)
callbacks.append(eval_callback)
if args.inference_file:
if args.preproc:
inference_data_layer = preproc_data_layer.PreprocBertSentenceClassificationDataLayer(
input_file=args.inference_file,
shuffle=False,
num_samples=args.num_samples,
batch_size=args.batch_size,
num_workers=0,
local_rank=args.local_rank)
else:
inference_data_layer = nemo_nlp.BertSentenceClassificationDataLayer(
input_file=args.inference_file,
tokenizer=tokenizer,
max_seq_length=args.max_seq_length,
shuffle=False,
num_samples=args.num_samples,
batch_size=args.batch_size,
num_workers=0,
local_rank=args.local_rank)
# TODO: Finish inference
inference_callback = None
# Training, eval and inference
if args.train_file:
lr_policy_fn = get_lr_policy(
args.lr_policy,
total_steps=args.num_epochs * steps_per_epoch,
warmup_ratio=args.lr_warmup_proportion)
nf.train(
tensors_to_optimize=[train_loss],
callbacks=callbacks,
lr_policy=lr_policy_fn,
optimizer=args.optimizer_kind,
optimization_params={'num_epochs': args.num_epochs, 'lr': args.lr})