fengshen/examples/pretrain_bert/pretrain_bert.py

from data.bert_dataloader.load import BertDataModule
from transformers import (
    BertTokenizer,
    BertConfig,
    BertForPreTraining,
    BertModel,
    BertForMaskedLM
)
from pytorch_lightning import (
    LightningDataModule,
    LightningModule,
    loggers,
    Trainer,
)
from pytorch_lightning.callbacks import (
    ModelCheckpoint,
    LearningRateMonitor,
)
from typing import Optional
from torch.utils.data import DataLoader
from transformers.optimization import get_linear_schedule_with_warmup
import argparse
import sys
import torch
import os
import re
import jieba
import numpy as np

# 如果没有安装fengshen模块，请把Fengshenbang-LM/fengshen加入到系统环境变量
sys.path.insert(0, '../../../fengshen')

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'


class DataCollate(object):

    def __init__(self, tokenizer, max_length, mask_rate=0.15, max_ngram=3, if_padding=True) -> None:
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.word_cuter = jieba.cut
        self.vocab_length = len(tokenizer)
        self.mask_rate = mask_rate
        self.ignore_labels = -100
        self.ngrams = np.arange(1, max_ngram + 1, dtype=np.int64)
        pvals = 1. / np.arange(1, max_ngram + 1)
        pvals /= pvals.sum(keepdims=True)  # p(n) = 1/n / sigma(1/k)
        self.pvals = pvals
        self.padding = if_padding

    def token_process(self, token_id):
        rand = np.random.random()
        if rand <= 0.8:
            return self.tokenizer.mask_token_id
        elif rand <= 0.9:
            return token_id
        else:
            return np.random.randint(1, self.vocab_length)

    def __call__(self, samples):
        input_ids = []
        attention_mask = []
        token_type_ids = []
        batch_labels = []
        # print('^-^ batch size :',len(samples))
        for sample in samples:
            word_list = list(self.word_cuter(sample['text']))
            mask_ids, labels = [], []

            record = []
            for i in range(len(word_list)):
                rands = np.random.random()
                if i in record:
                    continue
                word = word_list[i]
                if rands > self.mask_rate and len(word) < 4:
                    word = word_list[i]
                    word_encode = tokenizer.encode(word, add_special_tokens=False)
                    for token in word_encode:
                        mask_ids.append(token)
                        labels.append(self.ignore_labels)
                    record.append(i)
                else:
                    n = np.random.choice(self.ngrams, p=self.pvals)
                    for index in range(n):
                        ind = index + i
                        if ind in record or ind >= len(word_list):
                            continue
                        record.append(ind)
                        word = word_list[ind]
                        word_encode = tokenizer.encode(word, add_special_tokens=False)
                        for token in word_encode:
                            mask_ids.append(self.token_process(token))
                            labels.append(token)
            if self.padding:
                if len(mask_ids) > self.max_length:
                    input_ids.append(mask_ids[:self.max_length])
                    batch_labels.append(labels[:self.max_length])
                else:
                    lenght = len(mask_ids)
                    mask_ids.extend([0]*(self.max_length-lenght))
                    labels.extend([-100]*(self.max_length-lenght))
                    input_ids.append(mask_ids)
                    batch_labels.append(labels)
            attention_mask.append([1]*self.max_length)
            token_type_ids.append([0]*self.max_length)

        #     print('sentence:',sample['text'])
        #     print('input_ids:',mask_ids)
        #     print('decode inputids:',self.tokenizer.decode(mask_ids))
        #     print('labels',labels)
        #     print('decode labels:',self.tokenizer.decode(labels))
        #     print('*'*20)
        return {
            'input_ids': torch.tensor(input_ids),
            'labels': torch.tensor(batch_labels),
            'attention_mask': torch.tensor(attention_mask),
            'token_type_ids': torch.tensor(token_type_ids)
        }


class Bert(LightningModule):
    @staticmethod
    def add_module_specific_args(args_parser):
        parser = args_parser.add_argument_group('Bert')
        parser.add_argument('--model_path', type=str, default='')
        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)
        return args_parser

    def __init__(self, args):
        super().__init__()
        self.save_hyperparameters(args)
        self.bertconfig = BertConfig.from_pretrained(args.model_path)
        # self.model = BertForPreTraining(self.bertconfig)
        self.model = BertForMaskedLM(self.bertconfig)

    def setup(self, stage) -> None:
        if stage == 'fit':
            train_loader = self.trainer._data_connector._train_dataloader_source.dataloader()

            # Calculate total steps
            tb_size = self.hparams.train_batchsize * max(1, self.trainer.gpus)
            ab_size = self.trainer.accumulate_grad_batches * float(self.trainer.max_epochs)
            self.total_steps = (len(train_loader.dataset) // tb_size) // ab_size

    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.hparams.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.hparams.learning_rate)
        scheduler = get_linear_schedule_with_warmup(
            optimizer, int(self.total_steps * self.hparams.warmup),
            self.total_steps)

        return [{
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'step',
                'frequency': 1
            }
        }]

    def training_step(self, batch, batch_idx):
        output = self.model(**batch)
        # print(output)
        self.log('train_loss', output.loss)
        return output.loss

    def comput_metrix(self, logits, labels):
        ones = torch.ones_like(labels)
        zero = torch.zeros_like(labels)
        mask = torch.where(labels < 0, zero, ones)
        mask = mask.view(size=(-1,)).float()
        # y_true=labels.view(size=(-1,)).float()

        y_pred = torch.argmax(logits, dim=-1)
        y_pred = y_pred.view(size=(-1,))
        y_true = labels.view(size=(-1,)).float()
        corr = torch.eq(y_pred, y_true)
        corr = torch.multiply(corr.float(), mask)
        acc = torch.sum(corr.float()) / torch.sum(mask)
        return acc

    def validation_step(self, batch, batch_idx):
        output = self.model(**batch)
        # print(output)
        acc = self.comput_metrix(output.logits, batch['labels'])
        print('val_loss ', output.loss)
        self.log('val_loss', output.loss)
        self.log('val_acc', acc)
        # pass

    def predict_step(self, batch, batch_idx):
        output = self.model(**batch)
        return output.prediction_logits


class CustomCKPT:
    @staticmethod
    def add_argparse_args(parent_args):
        parser = parent_args.add_argument_group('ckpt call back')

        parser.add_argument('--monitor', default='train_loss', type=str)
        parser.add_argument('--mode', default='min', type=str)
        parser.add_argument('--dirpath', default='./ckpt/', type=str)
        parser.add_argument(
            '--filename', default='model-{epoch:02d}-{train_loss:.4f}', type=str)
        parser.add_argument('--save_last', action='store_true', default=True)
        parser.add_argument('--save_top_k', default=3, type=float)
        parser.add_argument('--every_n_train_steps', default=100, type=float)
        parser.add_argument('--save_weights_only', action='store_true', default=False)

        return parent_args

    def __init__(self, args):
        self.callbacks = ModelCheckpoint(monitor=args.monitor,
                                         save_top_k=args.save_top_k,
                                         mode=args.mode,
                                         every_n_train_steps=args.every_n_train_steps,
                                         save_weights_only=args.save_weights_only,
                                         dirpath=args.dirpath,
                                         filename=args.filename,
                                         save_last=args.save_last)


if __name__ == '__main__':
    args_parser = argparse.ArgumentParser()
    args_parser = BertDataModule.add_data_specific_args(args_parser)
    args_parser = Trainer.add_argparse_args(args_parser)
    args_parser = Bert.add_module_specific_args(args_parser)
    args_parser = CustomCKPT.add_argparse_args(args_parser)
    args_parser.add_argument('--deepspeed')
    args_parser.add_argument('--seq_max_length')

    args = args_parser.parse_args()

    tokenizer = BertTokenizer.from_pretrained(args.model_path)
    collate_fn = DataCollate(tokenizer, 512)
    data_module = BertDataModule(tokenizer=tokenizer, args=args, collate_fn=collate_fn)

    print('data load complete')

    model = Bert(args)
    print('model load complete')

    lr_monitor = LearningRateMonitor(logging_interval='step')
    logger = loggers.TensorBoardLogger(save_dir=os.path.join(
        args.default_root_dir, 'logs/'),
        name=os.path.basename(os.path.dirname(args.model_path)))
    checkpoint_callback = CustomCKPT(args).callbacks

    if args.resume_from_checkpoint is not None and \
            not os.path.exists(args.resume_from_checkpoint):
        print('--------warning no checkpoint found--------, remove args')
        del args.resume_from_checkpoint

    # autotuning
    if args.deepspeed is not None:
        os.environ['PL_DEEPSPEED_CONFIG_PATH'] = args.deepspeed

    trainer = Trainer.from_argparse_args(args, logger=logger,
                                         callbacks=[
                                             lr_monitor,
                                             checkpoint_callback])

    trainer.fit(model, data_module)