def create_pipeline( data_dir, model, head, loss_fn, max_query_length, max_seq_length, doc_stride, batch_size, version_2_with_negative, num_gpus=1, batches_per_step=1, mode="train", ): data_layer = nemo_nlp.BertQuestionAnsweringDataLayer( mode=mode, version_2_with_negative=version_2_with_negative, batch_size=batch_size, tokenizer=tokenizer, data_dir=data_dir, max_query_length=max_query_length, max_seq_length=max_seq_length, doc_stride=doc_stride, ) input_data = data_layer() hidden_states = model( input_ids=input_data.input_ids, token_type_ids=input_data.input_type_ids, attention_mask=input_data.input_mask, ) qa_output = head(hidden_states=hidden_states) loss_output = loss_fn( logits=qa_output, start_positions=input_data.start_positions, end_positions=input_data.end_positions, ) steps_per_epoch = len(data_layer) // (batch_size * num_gpus * batches_per_step) return ( loss_output.loss, steps_per_epoch, [ loss_output.start_logits, loss_output.end_logits, input_data.unique_ids, ], data_layer, )
def test_squad_v1(self): version_2_with_negative = False pretrained_bert_model = 'bert-base-uncased' batch_size = 3 data_dir = os.path.abspath( os.path.join(os.path.dirname(__file__), '../data/nlp/squad/v1.1')) max_query_length = 64 max_seq_length = 384 doc_stride = 128 max_steps = 100 lr_warmup_proportion = 0 eval_step_freq = 50 lr = 3e-6 do_lower_case = True n_best_size = 5 max_answer_length = 20 null_score_diff_threshold = 0.0 tokenizer = nemo_nlp.NemoBertTokenizer(pretrained_bert_model) neural_factory = nemo.core.NeuralModuleFactory( backend=nemo.core.Backend.PyTorch, local_rank=None, create_tb_writer=False, ) model = nemo_nlp.huggingface.BERT( pretrained_model_name=pretrained_bert_model) hidden_size = model.local_parameters["hidden_size"] qa_head = nemo_nlp.TokenClassifier( hidden_size=hidden_size, num_classes=2, num_layers=1, log_softmax=False, ) squad_loss = nemo_nlp.QuestionAnsweringLoss() data_layer = nemo_nlp.BertQuestionAnsweringDataLayer( mode='train', version_2_with_negative=version_2_with_negative, batch_size=batch_size, tokenizer=tokenizer, data_dir=data_dir, max_query_length=max_query_length, max_seq_length=max_seq_length, doc_stride=doc_stride, ) ( input_ids, input_type_ids, input_mask, start_positions, end_positions, _, ) = data_layer() hidden_states = model( input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask, ) qa_output = qa_head(hidden_states=hidden_states) loss, _, _ = squad_loss( logits=qa_output, start_positions=start_positions, end_positions=end_positions, ) data_layer_eval = nemo_nlp.BertQuestionAnsweringDataLayer( mode='dev', version_2_with_negative=version_2_with_negative, batch_size=batch_size, tokenizer=tokenizer, data_dir=data_dir, max_query_length=max_query_length, max_seq_length=max_seq_length, doc_stride=doc_stride, ) ( input_ids_eval, input_type_ids_eval, input_mask_eval, start_positions_eval, end_positions_eval, unique_ids_eval, ) = data_layer_eval() hidden_states_eval = model( input_ids=input_ids_eval, token_type_ids=input_type_ids_eval, attention_mask=input_mask_eval, ) qa_output_eval = qa_head(hidden_states=hidden_states_eval) _, start_logits_eval, end_logits_eval = squad_loss( logits=qa_output_eval, start_positions=start_positions_eval, end_positions=end_positions_eval, ) eval_output = [start_logits_eval, end_logits_eval, unique_ids_eval] callback_train = nemo.core.SimpleLossLoggerCallback( tensors=[loss], print_func=lambda x: print("Loss: {:.3f}".format(x[0].item())), get_tb_values=lambda x: [["loss", x[0]]], step_freq=10, tb_writer=neural_factory.tb_writer, ) callbacks_eval = nemo.core.EvaluatorCallback( eval_tensors=eval_output, user_iter_callback=lambda x, y: eval_iter_callback(x, y), user_epochs_done_callback=lambda x: eval_epochs_done_callback( x, eval_data_layer=data_layer_eval, do_lower_case=do_lower_case, n_best_size=n_best_size, max_answer_length=max_answer_length, version_2_with_negative=version_2_with_negative, null_score_diff_threshold=null_score_diff_threshold, ), tb_writer=neural_factory.tb_writer, eval_step=eval_step_freq, ) lr_policy_fn = get_lr_policy( 'WarmupAnnealing', total_steps=max_steps, warmup_ratio=lr_warmup_proportion, ) neural_factory.train( tensors_to_optimize=[loss], callbacks=[callback_train, callbacks_eval], lr_policy=lr_policy_fn, optimizer='adam_w', optimization_params={ "max_steps": max_steps, "lr": lr }, )