# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import numpy as np import paddle from biencoder_base_model import BiEncoder, BiEncoderNllLoss from NQdataset import DataUtil, NQdataSetForDPR from paddle.optimizer.lr import LambdaDecay from paddlenlp.transformers.bert.modeling import BertModel parser = argparse.ArgumentParser() parser.add_argument("--batch_size", required=True, type=int, default=None) parser.add_argument("--learning_rate", required=True, type=float, default=None) parser.add_argument("--save_dir", required=True, type=str, default=None) parser.add_argument("--warmup_steps", required=True, type=int) parser.add_argument("--epoches", required=True, type=int) parser.add_argument("--max_grad_norm", required=True, type=int) parser.add_argument("--train_data_path", required=True, type=str) parser.add_argument("--chunk_size", required=True, type=int) args = parser.parse_args() chunk_nums = args.batch_size // args.chunk_size data_path = args.train_data_path batch_size = args.batch_size learning_rate = args.learning_rate epoches = args.epoches def dataLoader_for_DPR(batch_size, source_data: list, epochs): index = np.arange(0, len(source_data)) np.random.shuffle(index) batch_data = [] for i in index: try: batch_data.append(source_data[i]) if len(batch_data) == batch_size: yield batch_data batch_data = [] except Exception: import traceback traceback.print_exc() continue def get_model(model_name: str): question_model = BertModel.from_pretrained(model_name) context_model = BertModel.from_pretrained(model_name) model = BiEncoder(question_model, context_model) return model model = get_model("bert-base-uncased") def get_linear_scheduler(warmup_steps, training_steps): def lr_lambda(current_step): if current_step < warmup_steps: return float(current_step) / float(max(1, warmup_steps)) return max(0.0, float(training_steps - current_step) / float(max(1, training_steps - warmup_steps))) return LambdaDecay(learning_rate=args.learning_rate, lr_lambda=lr_lambda, last_epoch=-1, verbose=False) training_steps = 58880 * args.epoches / args.batch_size scheduler = get_linear_scheduler(args.warmup_steps, training_steps) optimizer = paddle.optimizer.AdamW(learning_rate=scheduler, parameters=model.parameters()) def get_dataset(data_path: str): data = NQdataSetForDPR(data_path) dataset = data.new_data return dataset util = DataUtil() LOSS = BiEncoderNllLoss() batch_data = [] dataset = get_dataset(data_path) def train(): for epoch in range(epoches): index = np.arange(0, len(dataset)) np.random.shuffle(index) batch_data = [] for i in index: # dataLoader batch_data.append(dataset[i]) if len(batch_data) == batch_size: all_questions = [] all_contexts = [] all_batch_input = util.create_biencoder_input(batch_data, inserted_title=True) all_positions = all_batch_input.is_positive all_inputs_questions_id = all_batch_input.questions_ids all_inputs_questions_segment = all_batch_input.question_segments all_inputs_contexts_id = all_batch_input.context_ids all_inputs_contexts_segment = all_batch_input.ctx_segments sub_q_ids = paddle.split(all_inputs_questions_id, chunk_nums, axis=0) sub_c_ids = paddle.split(all_inputs_contexts_id, chunk_nums, axis=0) sub_q_segments = paddle.split(all_inputs_questions_segment, chunk_nums, axis=0) sub_c_segments = paddle.split(all_inputs_contexts_segment, chunk_nums, axis=0) all_questions = [] all_contexts = [] all_CUDA_rnd_state_question = [] all_CUDA_rnd_state_context = [] for sub_q_id, sub_q_segment in zip(sub_q_ids, sub_q_segments): with paddle.no_grad(): sub_CUDA_rnd_state = paddle.framework.random.get_cuda_rng_state() all_CUDA_rnd_state_question.append(sub_CUDA_rnd_state) sub_question_output = model.get_question_pooled_embedding(sub_q_id, sub_q_segment) all_questions.append(sub_question_output) for sub_c_id, sub_c_segment in zip(sub_c_ids, sub_c_segments): with paddle.no_grad(): sub_CUDA_rnd_state = paddle.framework.random.get_cuda_rng_state() all_CUDA_rnd_state_context.append(sub_CUDA_rnd_state) sub_context_output = model.get_context_pooled_embedding(sub_c_id, sub_c_segment) all_contexts.append(sub_context_output) model_questions = paddle.concat(all_questions, axis=0) all_questions = [] model_questions = model_questions.detach() model_questions.stop_gradient = False model_contexts = paddle.concat(all_contexts, axis=0) model_contexts = model_contexts.detach() model_contexts.stop_gradient = False all_contexts = [] model_positions = all_positions loss, _ = LOSS.calc(model_questions, model_contexts, model_positions) print("loss is:") print(loss.item()) loss.backward() grads_for_questions = paddle.split(model_questions.grad, chunk_nums, axis=0) grads_for_contexts = paddle.split(model_contexts.grad, chunk_nums, axis=0) for sub_q_id, sub_q_segment, CUDA_state, grad_for_each_question in zip( sub_q_ids, sub_q_segments, all_CUDA_rnd_state_question, grads_for_questions ): paddle.framework.random.set_cuda_rng_state(CUDA_state) sub_question_output = model.get_question_pooled_embedding(sub_q_id, sub_q_segment) finally_question_res_for_backward = paddle.dot(sub_question_output, grad_for_each_question) finally_question_res_for_backward = finally_question_res_for_backward * (1 / 8.0) finally_question_res_for_backward.backward(retain_graph=True) for sub_c_id, sub_c_segment, CUDA_state, grad_for_each_context in zip( sub_c_ids, sub_c_segments, all_CUDA_rnd_state_context, grads_for_contexts ): paddle.framework.random.set_cuda_rng_state(CUDA_state) sub_context_output = model.get_context_pooled_embedding(sub_c_id, sub_q_segment) finally_context_res_for_backward = paddle.dot(sub_question_output, grad_for_each_context) finally_context_res_for_backward = finally_context_res_for_backward * (1 / 8.0) finally_context_res_for_backward.backward(retain_graph=True) paddle.nn.ClipGradByGlobalNorm(clip_norm=args.max_grad_norm, group_name=model.parameters()) optimizer.step() scheduler.step() optimizer.clear_grad() batch_data = [] EPOCH = str(epoch) save_path_que = args.save_dir + "/question_model_" + EPOCH save_path_con = args.save_dir + "/context_model_" + EPOCH model.question_encoder.save_pretrained(save_path_que) model.context_encoder.save_pretrained(save_path_con) if __name__ == "__main__": train()