# coding=utf-8 # Copyright 2021 The IDEA 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. from logging import basicConfig import torch from torch import nn import json from tqdm import tqdm import os import numpy as np from transformers import BertTokenizer import pytorch_lightning as pl from pytorch_lightning.callbacks import ModelCheckpoint from pytorch_lightning import trainer, loggers from torch.utils.data import Dataset, DataLoader from transformers.optimization import get_linear_schedule_with_warmup from transformers import BertForMaskedLM, AlbertTokenizer from transformers import AutoConfig from transformers.pipelines.base import Pipeline from transformers import MegatronBertForMaskedLM from fengshen.models.deberta_v2.modeling_deberta_v2 import DebertaV2ForMaskedLM from fengshen.models.albert.modeling_albert import AlbertForMaskedLM import argparse import copy from fengshen.utils.universal_checkpoint import UniversalCheckpoint import warnings from transformers import TextClassificationPipeline as HuggingfacePipe class UniMCDataset(Dataset): def __init__(self, data, yes_token, no_token, tokenizer, args, used_mask=True): super().__init__() self.tokenizer = tokenizer self.max_length = args.max_length self.num_labels = args.num_labels self.used_mask = used_mask self.data = data self.args = args self.yes_token = yes_token self.no_token = no_token def __len__(self): return len(self.data) def __getitem__(self, index): return self.encode(self.data[index], self.used_mask) def get_token_type(self, sep_idx, max_length): token_type_ids = np.zeros(shape=(max_length,)) for i in range(len(sep_idx)-1): if i % 2 == 0: ty = np.ones(shape=(sep_idx[i+1]-sep_idx[i],)) else: ty = np.zeros(shape=(sep_idx[i+1]-sep_idx[i],)) token_type_ids[sep_idx[i]:sep_idx[i+1]] = ty return token_type_ids def get_position_ids(self, label_idx, max_length, question_len): question_position_ids = np.arange(question_len) label_position_ids = np.arange(question_len, label_idx[-1]) for i in range(len(label_idx)-1): label_position_ids[label_idx[i]-question_len:label_idx[i+1]-question_len] = np.arange( question_len, question_len+label_idx[i+1]-label_idx[i]) max_len_label = max(label_position_ids) text_position_ids = np.arange( max_len_label+1, max_length+max_len_label+1-label_idx[-1]) position_ids = list(question_position_ids) + \ list(label_position_ids)+list(text_position_ids) if max_length <= 512: return position_ids[:max_length] else: for i in range(512, max_length): if position_ids[i] > 511: position_ids[i] = 511 return position_ids[:max_length] def get_att_mask(self, attention_mask, label_idx, question_len): max_length = len(attention_mask) attention_mask = np.array(attention_mask) attention_mask = np.tile(attention_mask[None, :], (max_length, 1)) zeros = np.zeros( shape=(label_idx[-1]-question_len, label_idx[-1]-question_len)) attention_mask[question_len:label_idx[-1], question_len:label_idx[-1]] = zeros for i in range(len(label_idx)-1): label_token_length = label_idx[i+1]-label_idx[i] if label_token_length <= 0: print('label_idx', label_idx) print('question_len', question_len) continue ones = np.ones(shape=(label_token_length, label_token_length)) attention_mask[label_idx[i]:label_idx[i+1], label_idx[i]:label_idx[i+1]] = ones return attention_mask def random_masking(self, token_ids, maks_rate, mask_start_idx, max_length, mask_id, tokenizer): rands = np.random.random(len(token_ids)) source, target = [], [] for i, (r, t) in enumerate(zip(rands, token_ids)): if i < mask_start_idx: source.append(t) target.append(-100) continue if r < maks_rate * 0.8: source.append(mask_id) target.append(t) elif r < maks_rate * 0.9: source.append(t) target.append(t) elif r < maks_rate: source.append(np.random.choice(tokenizer.vocab_size - 1) + 1) target.append(t) else: source.append(t) target.append(-100) while len(source) < max_length: source.append(0) target.append(-100) return source[:max_length], target[:max_length] def encode(self, item, used_mask=False): while len(self.tokenizer.encode('[MASK]'.join(item['choice']))) > self.max_length-32: item['choice'] = [c[:int(len(c)/2)] for c in item['choice']] if 'textb' in item.keys() and item['textb'] != '': if 'question' in item.keys() and item['question'] != '': texta = '[MASK]' + '[MASK]'.join(item['choice']) + '[SEP]' + \ item['question'] + '[SEP]' + \ item['texta']+'[SEP]'+item['textb'] else: texta = '[MASK]' + '[MASK]'.join(item['choice']) + '[SEP]' + \ item['texta']+'[SEP]'+item['textb'] else: if 'question' in item.keys() and item['question'] != '': texta = '[MASK]' + '[MASK]'.join(item['choice']) + '[SEP]' + \ item['question'] + '[SEP]' + item['texta'] else: texta = '[MASK]' + '[MASK]'.join(item['choice']) + \ '[SEP]' + item['texta'] encode_dict = self.tokenizer.encode_plus(texta, max_length=self.max_length, padding='max_length', truncation='longest_first') encode_sent = encode_dict['input_ids'] token_type_ids = encode_dict['token_type_ids'] attention_mask = encode_dict['attention_mask'] sample_max_length = sum(encode_dict['attention_mask']) if 'label' not in item.keys(): item['label'] = 0 item['answer'] = '' question_len = 1 label_idx = [question_len] for choice in item['choice']: cur_mask_idx = label_idx[-1] + \ len(self.tokenizer.encode(choice, add_special_tokens=False))+1 label_idx.append(cur_mask_idx) token_type_ids = [0]*question_len+[1] * \ (label_idx[-1]-label_idx[0]+1)+[0]*self.max_length token_type_ids = token_type_ids[:self.max_length] attention_mask = self.get_att_mask( attention_mask, label_idx, question_len) position_ids = self.get_position_ids( label_idx, self.max_length, question_len) clslabels_mask = np.zeros(shape=(len(encode_sent),)) clslabels_mask[label_idx[:-1]] = 10000 clslabels_mask = clslabels_mask-10000 mlmlabels_mask = np.zeros(shape=(len(encode_sent),)) mlmlabels_mask[label_idx[0]] = 1 # used_mask=False if used_mask: mask_rate = 0.1*np.random.choice(4, p=[0.3, 0.3, 0.25, 0.15]) source, target = self.random_masking(token_ids=encode_sent, maks_rate=mask_rate, mask_start_idx=label_idx[-1], max_length=self.max_length, mask_id=self.tokenizer.mask_token_id, tokenizer=self.tokenizer) else: source, target = encode_sent[:], encode_sent[:] source = np.array(source) target = np.array(target) source[label_idx[:-1]] = self.tokenizer.mask_token_id target[label_idx[:-1]] = self.no_token target[label_idx[item['label']]] = self.yes_token input_ids = source[:sample_max_length] token_type_ids = token_type_ids[:sample_max_length] attention_mask = attention_mask[:sample_max_length, :sample_max_length] position_ids = position_ids[:sample_max_length] mlmlabels = target[:sample_max_length] clslabels = label_idx[item['label']] clslabels_mask = clslabels_mask[:sample_max_length] mlmlabels_mask = mlmlabels_mask[:sample_max_length] return { "input_ids": torch.tensor(input_ids).long(), "token_type_ids": torch.tensor(token_type_ids).long(), "attention_mask": torch.tensor(attention_mask).float(), "position_ids": torch.tensor(position_ids).long(), "mlmlabels": torch.tensor(mlmlabels).long(), "clslabels": torch.tensor(clslabels).long(), "clslabels_mask": torch.tensor(clslabels_mask).float(), "mlmlabels_mask": torch.tensor(mlmlabels_mask).float(), } class UniMCDataModel(pl.LightningDataModule): @staticmethod def add_data_specific_args(parent_args): parser = parent_args.add_argument_group('TASK NAME DataModel') parser.add_argument('--num_workers', default=8, type=int) parser.add_argument('--batchsize', default=16, type=int) parser.add_argument('--max_length', default=512, type=int) return parent_args def __init__(self, train_data, val_data, yes_token, no_token, tokenizer, args): super().__init__() self.batchsize = args.batchsize self.train_data = UniMCDataset( train_data, yes_token, no_token, tokenizer, args, True) self.valid_data = UniMCDataset( val_data, yes_token, no_token, tokenizer, args, False) def train_dataloader(self): return DataLoader(self.train_data, shuffle=True, collate_fn=self.collate_fn, batch_size=self.batchsize, pin_memory=False) def val_dataloader(self): return DataLoader(self.valid_data, shuffle=False, collate_fn=self.collate_fn, batch_size=self.batchsize, pin_memory=False) def collate_fn(self, batch): ''' Aggregate a batch data. batch = [ins1_dict, ins2_dict, ..., insN_dict] batch_data = {'sentence':[ins1_sentence, ins2_sentence...], 'input_ids':[ins1_input_ids, ins2_input_ids...], ...} ''' batch_data = {} for key in batch[0]: batch_data[key] = [example[key] for example in batch] batch_data['input_ids'] = nn.utils.rnn.pad_sequence(batch_data['input_ids'], batch_first=True, padding_value=0) batch_data['clslabels_mask'] = nn.utils.rnn.pad_sequence(batch_data['clslabels_mask'], batch_first=True, padding_value=-10000) batch_size, batch_max_length = batch_data['input_ids'].shape for k, v in batch_data.items(): if k == 'input_ids' or k == 'clslabels_mask': continue if k == 'clslabels': batch_data[k] = torch.tensor(v).long() continue if k != 'attention_mask': batch_data[k] = nn.utils.rnn.pad_sequence(v, batch_first=True, padding_value=0) else: attention_mask = torch.zeros( (batch_size, batch_max_length, batch_max_length)) for i, att in enumerate(v): sample_length, _ = att.shape attention_mask[i, :sample_length, :sample_length] = att batch_data[k] = attention_mask return batch_data class UniMCModel(nn.Module): def __init__(self, pre_train_dir, yes_token): super().__init__() self.config = AutoConfig.from_pretrained(pre_train_dir) if self.config.model_type == 'megatron-bert': self.bert = MegatronBertForMaskedLM.from_pretrained(pre_train_dir) elif self.config.model_type == 'deberta-v2': self.bert = DebertaV2ForMaskedLM.from_pretrained(pre_train_dir) elif self.config.model_type == 'albert': self.bert = AlbertForMaskedLM.from_pretrained(pre_train_dir) else: self.bert = BertForMaskedLM.from_pretrained(pre_train_dir) self.loss_func = torch.nn.CrossEntropyLoss() self.yes_token = yes_token def forward(self, input_ids, attention_mask, token_type_ids, position_ids=None, mlmlabels=None, clslabels=None, clslabels_mask=None, mlmlabels_mask=None): batch_size, seq_len = input_ids.shape outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, token_type_ids=token_type_ids, labels=mlmlabels) # (bsz, seq, dim) mask_loss = outputs.loss mlm_logits = outputs.logits cls_logits = mlm_logits[:, :, self.yes_token].view(-1, seq_len)+clslabels_mask if mlmlabels == None: return 0, mlm_logits, cls_logits else: cls_loss = self.loss_func(cls_logits, clslabels) all_loss = mask_loss+cls_loss return all_loss, mlm_logits, cls_logits class UniMCLitModel(pl.LightningModule): @staticmethod def add_model_specific_args(parent_args): parser = parent_args.add_argument_group('BaseModel') parser.add_argument('--learning_rate', default=1e-5, type=float) parser.add_argument('--weight_decay', default=0.1, type=float) parser.add_argument('--warmup', default=0.01, type=float) parser.add_argument('--num_labels', default=2, type=int) return parent_args def __init__(self, args, yes_token, model_path, num_data=100): super().__init__() self.args = args self.num_data = num_data self.model = UniMCModel(model_path, yes_token) def setup(self, stage) -> None: if stage == 'fit': num_gpus = self.trainer.gpus if self.trainer.gpus is not None else 0 self.total_step = int(self.trainer.max_epochs * self.num_data / (max(1, num_gpus) * self.trainer.accumulate_grad_batches)) print('Total training step:', self.total_step) def training_step(self, batch, batch_idx): loss, logits, cls_logits = self.model(**batch) cls_acc = self.comput_metrix( cls_logits, batch['clslabels'], batch['mlmlabels_mask']) self.log('train_loss', loss) self.log('train_acc', cls_acc) return loss def validation_step(self, batch, batch_idx): loss, logits, cls_logits = self.model(**batch) cls_acc = self.comput_metrix( cls_logits, batch['clslabels'], batch['mlmlabels_mask']) self.log('val_loss', loss) self.log('val_acc', cls_acc) def configure_optimizers(self): no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] paras = list( filter(lambda p: p[1].requires_grad, self.named_parameters())) paras = [{ 'params': [p for n, p in paras if not any(nd in n for nd in no_decay)], 'weight_decay': self.args.weight_decay }, { 'params': [p for n, p in paras if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] optimizer = torch.optim.AdamW(paras, lr=self.args.learning_rate) scheduler = get_linear_schedule_with_warmup( optimizer, int(self.total_step * self.args.warmup), self.total_step) return [{ 'optimizer': optimizer, 'lr_scheduler': { 'scheduler': scheduler, 'interval': 'step', 'frequency': 1 } }] def comput_metrix(self, logits, labels, mlmlabels_mask): logits = torch.nn.functional.softmax(logits, dim=-1) logits = torch.argmax(logits, dim=-1) y_pred = logits.view(size=(-1,)) y_true = labels.view(size=(-1,)) corr = torch.eq(y_pred, y_true).float() return torch.sum(corr.float())/labels.size(0) class UniMCPredict: def __init__(self, yes_token, no_token, model, tokenizer, args): self.tokenizer = tokenizer self.args = args self.data_model = UniMCDataModel( [], [], yes_token, no_token, tokenizer, args) self.model = model def predict(self, batch_data): batch = [self.data_model.train_data.encode( sample) for sample in batch_data] batch = self.data_model.collate_fn(batch) batch = {k: v.cuda() for k, v in batch.items()} _, _, logits = self.model.model(**batch) soft_logits = torch.nn.functional.softmax(logits, dim=-1) logits = torch.argmax(soft_logits, dim=-1).detach().cpu().numpy() soft_logits = soft_logits.detach().cpu().numpy() clslabels_mask = batch['clslabels_mask'].detach( ).cpu().numpy().tolist() clslabels = batch['clslabels'].detach().cpu().numpy().tolist() for i, v in enumerate(batch_data): label_idx = [idx for idx, v in enumerate( clslabels_mask[i]) if v == 0.] label = label_idx.index(logits[i]) answer = batch_data[i]['choice'][label] score = {} for c in range(len(batch_data[i]['choice'])): score[batch_data[i]['choice'][c]] = float( soft_logits[i][label_idx[c]]) batch_data[i]['label_ori'] = copy.deepcopy(batch_data[i]['label']) batch_data[i]['label'] = label batch_data[i]['answer'] = answer batch_data[i]['score'] = score return batch_data class UniMCPipelines(Pipeline): @staticmethod def piplines_args(parent_args): total_parser = parent_args.add_argument_group("piplines args") total_parser.add_argument( '--pretrained_model_path', default='', type=str) total_parser.add_argument('--load_checkpoints_path', default='', type=str) total_parser.add_argument('--train', action='store_true') total_parser.add_argument('--language', default='chinese', type=str) total_parser = UniMCDataModel.add_data_specific_args(total_parser) total_parser = UniversalCheckpoint.add_argparse_args(total_parser) total_parser = UniMCLitModel.add_model_specific_args(total_parser) total_parser = pl.Trainer.add_argparse_args(parent_args) return parent_args def __init__(self, args, model_path): self.args = args self.checkpoint_callback = UniversalCheckpoint(args).callbacks self.logger = loggers.TensorBoardLogger(save_dir=args.default_root_dir) self.trainer = pl.Trainer.from_argparse_args(args, logger=self.logger, callbacks=[self.checkpoint_callback]) self.config = AutoConfig.from_pretrained(model_path) if self.config.model_type == 'albert': self.tokenizer = AlbertTokenizer.from_pretrained( model_path) else: self.tokenizer = BertTokenizer.from_pretrained( model_path) if args.language == 'chinese': self.yes_token = self.tokenizer.encode('是')[1] self.no_token = self.tokenizer.encode('非')[1] else: self.yes_token = self.tokenizer.encode('yes')[1] self.no_token = self.tokenizer.encode('no')[1] if args.load_checkpoints_path != '': self.model = UniMCLitModel.load_from_checkpoint( args.load_checkpoints_path, args=args, yes_token=self.yes_token, model_path=model_path) print('load model from: ', args.load_checkpoints_path) else: self.model = UniMCLitModel( args, yes_token=self.yes_token, model_path=model_path) def train(self, train_data, dev_data, process=True): if process: train_data = self.preprocess(train_data) dev_data = self.preprocess(dev_data) data_model = UniMCDataModel( train_data, dev_data, self.yes_token, self.no_token, self.tokenizer, self.args) self.model.num_data = len(train_data) self.trainer.fit(self.model, data_model) def predict(self, test_data, cuda=True, process=True): if process: test_data = self.preprocess(test_data) result = [] start = 0 if cuda: self.model = self.model.cuda() self.model.model.eval() predict_model = UniMCPredict( self.yes_token, self.no_token, self.model, self.tokenizer, self.args) while start < len(test_data): batch_data = test_data[start:start+self.args.batchsize] start += self.args.batchsize batch_result = predict_model.predict(batch_data) result.extend(batch_result) if process: result = self.postprocess(result) return result def preprocess(self, data): for i, line in enumerate(data): if 'task_type' in line.keys() and line['task_type'] == '语义匹配': data[i]['choice'] = ['不能理解为：'+data[i] ['textb'], '可以理解为：'+data[i]['textb']] # data[i]['question']='怎么理解这段话？' data[i]['textb'] = '' if 'task_type' in line.keys() and line['task_type'] == '自然语言推理': data[i]['choice'] = ['不能推断出：'+data[i]['textb'], '很难推断出：'+data[i]['textb'], '可以推断出：'+data[i]['textb']] # data[i]['question']='根据这段话' data[i]['textb'] = '' return data def postprocess(self, data): for i, line in enumerate(data): if 'task_type' in line.keys() and line['task_type'] == '语义匹配': data[i]['textb'] = data[i]['choice'][0].replace('不能理解为：', '') data[i]['choice'] = ['不相似', '相似'] ns = {} for k, v in data[i]['score'].items(): if '不能' in k: k = '不相似' if '可以' in k: k = '相似' ns[k] = v data[i]['score'] = ns data[i]['answer'] = data[i]['choice'][data[i]['label']] if 'task_type' in line.keys() and line['task_type'] == '自然语言推理': data[i]['textb'] = data[i]['choice'][0].replace('不能推断出：', '') data[i]['choice'] = ['矛盾', '自然', '蕴含'] ns = {} for k, v in data[i]['score'].items(): if '不能' in k: k = '矛盾' if '很难' in k: k = '自然' if '可以' in k: k = '蕴含' ns[k] = v data[i]['score'] = ns data[i]['answer'] = data[i]['choice'][data[i]['label']] return data def _forward(self, model_inputs): return self.model(**model_inputs) def _sanitize_parameters(self, return_all_scores=None, function_to_apply=None, top_k="", **tokenizer_kwargs): # Using "" as default argument because we're going to use `top_k=None` in user code to declare # "No top_k" preprocess_params = tokenizer_kwargs postprocess_params = {} if hasattr(self.model.config, "return_all_scores") and return_all_scores is None: return_all_scores = self.model.config.return_all_scores if isinstance(top_k, int) or top_k is None: postprocess_params["top_k"] = top_k postprocess_params["_legacy"] = False elif return_all_scores is not None: warnings.warn( "`return_all_scores` is now deprecated, if want a similar funcionality use `top_k=None` instead of" " `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`.", UserWarning, ) if return_all_scores: postprocess_params["top_k"] = None else: postprocess_params["top_k"] = 1 if function_to_apply is not None: postprocess_params["function_to_apply"] = function_to_apply return preprocess_params, {}, postprocess_params def load_data(data_path): with open(data_path, 'r', encoding='utf8') as f: lines = f.readlines() samples = [json.loads(line) for line in tqdm(lines)] return samples def comp_acc(pred_data, test_data): corr = 0 for i in range(len(pred_data)): if pred_data[i]['label'] == test_data[i]['label']: corr += 1 return corr/len(pred_data) def main(): total_parser = argparse.ArgumentParser("TASK NAME") total_parser.add_argument('--data_dir', default='./data', type=str) total_parser.add_argument('--train_data', default='train.json', type=str) total_parser.add_argument('--valid_data', default='dev.json', type=str) total_parser.add_argument('--test_data', default='test.json', type=str) total_parser.add_argument('--output_path', default='', type=str) total_parser = UniMCPipelines.piplines_args(total_parser) args = total_parser.parse_args() train_data = load_data(os.path.join(args.data_dir, args.train_data)) dev_data = load_data(os.path.join(args.data_dir, args.valid_data)) test_data = load_data(os.path.join(args.data_dir, args.test_data)) dev_data_ori = copy.deepcopy(dev_data) model = UniMCPipelines(args) print(args.data_dir) if args.train: model.train(train_data, dev_data) result = model.predict(dev_data) for line in result[:20]: print(line) acc = comp_acc(result, dev_data_ori) print('acc:', acc) if args.output_path != '': test_result = model.predict(test_data) with open(args.output_path, 'w', encoding='utf8') as f: for line in test_result: json_data = json.dumps(line, ensure_ascii=False) f.write(json_data+'\n') if __name__ == "__main__": main()