fengshen/examples/pretrain_erlangshen_bert/pretrain_erlangshen.py

from dataclasses import dataclass
from transformers import (
    MegatronBertConfig,
    MegatronBertForPreTraining,
    AutoTokenizer,
)
from pytorch_lightning import (
    LightningModule,
    Trainer,
)
from pytorch_lightning.callbacks import (
    LearningRateMonitor,
)
import argparse
import torch
import os
import numpy as np
import time
from fengshen.data.universal_datamodule import UniversalDataModule
from fengshen.data.data_utils.sop_utils import get_a_and_b_segments
from fengshen.data.data_utils.truncate_utils import truncate_segments
from fengshen.data.data_utils.token_type_utils import create_tokens_and_tokentypes
from fengshen.data.data_utils.mask_utils import create_masked_lm_predictions
from fengshen.models.model_utils import (
    add_module_args,
    configure_optimizers,
    get_total_steps,
)
from fengshen.utils.universal_checkpoint import UniversalCheckpoint
from torch.utils.data._utils.collate import default_collate

SHOW_DATA = False


@dataclass
class ErLangShenCollator:
    '''
    由input处理成samples，也就是最终模型的输入
    其中主要处理逻辑在__call__里
    包含Mask和Sop任务
    '''
    tokenizer: None  # 分词
    max_seq_length: 512
    masked_lm_prob: 0.15
    content_key: str = 'text'
    # 一些预处理操作

    def setup(self):
        from fengshen.data.data_utils.sentence_split import ChineseSentenceSplitter
        self.sentence_split = ChineseSentenceSplitter()
        self.np_rng = np.random.RandomState(seed=((int(time.time()) % 2**32)))
        inv_vocab = {v: k for k, v in self.tokenizer.vocab.items()}
        self.vocab_id_list = list(inv_vocab.keys())
        self.vocab_id_to_token_dict = inv_vocab

    def __call__(self, samples):
        '''
        samples: 一个sample长这样{"text": "hello world"}
        '''
        model_inputs = []
        for s in samples:
            sentences = self.sentence_split.tokenize(s[self.content_key])
            # Divide sample into two segments (A and B).
            tokenized_sentences = [self.tokenizer.convert_tokens_to_ids(
                self.tokenizer.tokenize(sent)) for sent in sentences]
            if len(tokenized_sentences) == 0:
                print('find empty sentence')
                continue
            if len(tokenized_sentences) > 1:
                tokens_a, tokens_b, is_next_random = get_a_and_b_segments(tokenized_sentences,
                                                                          self.np_rng)
            else:
                tokens_a = tokenized_sentences[0]
                tokens_b = []
                is_next_random = False
            # max_seq_length - 3因为还需要拼上[CLS] [SEP] [SEP]
            if len(tokens_a) == 0:
                continue
            _ = truncate_segments(tokens_a, tokens_b, len(tokens_a),
                                  len(tokens_b), self.max_seq_length-3, self.np_rng)
            # Build tokens and toketypes.
            tokens, tokentypes = create_tokens_and_tokentypes(tokens_a, tokens_b,
                                                              self.tokenizer.cls_token_id, self.tokenizer.sep_token_id)
            # Masking.
            max_predictions_per_seq = self.masked_lm_prob * len(tokens)
            (tokens, masked_positions, masked_labels, _, _) = create_masked_lm_predictions(
                tokens, self.vocab_id_list, self.vocab_id_to_token_dict, self.masked_lm_prob,
                self.tokenizer.cls_token_id, self.tokenizer.sep_token_id, self.tokenizer.mask_token_id,
                max_predictions_per_seq, self.np_rng,
                masking_style='bert')

            # Some checks.
            num_tokens = len(tokens)
            padding_length = self.max_seq_length - num_tokens
            assert padding_length >= 0
            assert len(tokentypes) == num_tokens
            assert len(masked_positions) == len(masked_labels)

            # Tokens and token types.
            filler = [self.tokenizer.pad_token_id] * padding_length
            tokens_np = np.array(tokens + filler, dtype=np.int64)
            tokentypes_np = np.array(tokentypes + filler, dtype=np.int64)

            # Padding mask.
            padding_mask_np = np.array([1] * num_tokens + [0] * padding_length,
                                       dtype=np.int64)

            # Lables and loss mask.
            labels = [-100] * self.max_seq_length
            for i in range(len(masked_positions)):
                assert masked_positions[i] < num_tokens
                labels[masked_positions[i]] = masked_labels[i]
            labels_np = np.array(labels, dtype=np.int64)
            model_inputs.append(
                {
                    'input_ids': tokens_np,
                    'attention_mask': padding_mask_np,
                    'token_type_ids': tokentypes_np,
                    'labels': labels_np,
                    'next_sentence_label': int(is_next_random)
                }
            )
        return default_collate(model_inputs)


class ErLangShenBert(LightningModule):
    @staticmethod
    def add_module_specific_args(parent_parser):
        parser = parent_parser.add_argument_group('Erlangshen Bert')
        parser.add_argument('--masked_lm_prob', type=float, default=0.15)
        parser.add_argument('--max_seq_length', type=int, default=512)
        parser.add_argument('--sample_content_key', type=str, default='text')
        return parent_parser

    def __init__(self, args, tokenizer, **kwargs) -> None:
        super().__init__()
        self.save_hyperparameters(args)
        config = MegatronBertConfig.from_pretrained(args.model_path)
        self.config = config
        self.tokenizer = tokenizer
        self.model = MegatronBertForPreTraining(config)

    def setup(self, stage) -> None:
        if stage == 'fit':
            self.total_steps = get_total_steps(self.trainer, self.hparams)
            print('Total steps: {}' .format(self.total_steps))

    def configure_optimizers(self):
        return configure_optimizers(self)

    def forward(self, **batch):
        return self.model(**batch)

    def detokenize(self, token_ids):
        toks = self.tokenizer.convert_ids_to_tokens(token_ids)
        return self.tokenizer.convert_tokens_to_string(toks)

    def comput_metrix(self, logits, labels):
        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)
        acc = torch.sum(corr.float())/labels.shape[0]
        return acc

    def training_step(self, batch, batch_idx):
        if self.trainer.global_rank == 0:
            global SHOW_DATA
            if not SHOW_DATA:
                print(self.config)
                print(self.model)
                SHOW_DATA = True
                print('source: {}'.format(batch['input_ids'][0]))
                print('target: {}'.format(batch['labels'][0]))
                print('source: {}'.format(self.detokenize(batch['input_ids'][0])))
                label_idx = batch['labels'][0] != -100
                print('target: {}'.format(self.detokenize(
                    batch['labels'][0][label_idx])))
        output = self(**batch)
        self.log('train_loss', output.loss, sync_dist=True)
        label_idx = batch['labels'] != -100
        acc = self.comput_metrix(
            output.prediction_logits[label_idx].view(-1, output.prediction_logits.size(-1)), batch['labels'][label_idx])
        self.log('train_acc', acc, sync_dist=True)
        return output.loss

    def validation_step(self, batch, batch_idx):
        output = self(**batch)
        self.log('val_loss', output.loss, sync_dist=True)
        return output.loss

    def on_load_checkpoint(self, checkpoint) -> None:
        # 兼容低版本lightning，低版本lightning从ckpt起来时steps数会被重置为0
        global_step_offset = checkpoint["global_step"]
        if 'global_samples' in checkpoint:
            self.consumed_samples = checkpoint['global_samples']
        self.trainer.fit_loop.epoch_loop._batches_that_stepped = global_step_offset


if __name__ == '__main__':
    args_parser = argparse.ArgumentParser()
    args_parser = add_module_args(args_parser)
    args_parser = UniversalDataModule.add_data_specific_args(args_parser)
    args_parser = Trainer.add_argparse_args(args_parser)
    args_parser = ErLangShenBert.add_module_specific_args(args_parser)
    args_parser = UniversalCheckpoint.add_argparse_args(args_parser)
    args = args_parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    collate_fn = ErLangShenCollator(
        tokenizer=tokenizer,
        max_seq_length=args.max_seq_length,
        masked_lm_prob=args.masked_lm_prob,
        content_key=args.sample_content_key,
    )
    collate_fn.setup()
    data_module = UniversalDataModule(tokenizer=tokenizer, args=args, collate_fn=collate_fn)
    print('data load complete')

    model = ErLangShenBert(args, tokenizer=tokenizer)
    print('model load complete')

    lr_monitor = LearningRateMonitor(logging_interval='step')
    checkpoint_callback = UniversalCheckpoint(args)

    # 做兼容，如果目录不存在的话把这个参数去掉，不然会报错
    if args.load_ckpt_path is not None and \
            not os.path.exists(args.load_ckpt_path):
        print('--------warning no checkpoint found--------, remove args')
        args.load_ckpt_path = None

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

    trainer.fit(model, data_module, ckpt_path=args.load_ckpt_path)