【Whole Word Mask】中文 bert wwm 预训练代码，预训练你自己的模型

基于 transformers 来进行 bert 的 wwm （全词掩码）预训练，让你的模型可以领域适配

提示：文章偏向于NLP进阶，大佬请绕路……

---

---

前言

最近看了很多关于预训练的文章和论文，最近的工作也是需要一个垂直领域的预训练模型来进行领域适配的任务，博主现在是做关于NLP医学相关领域的工作，但是现在比较适合的预训练模型又太少了，找来找去就哪些（mc bert, plc-medbert-base, med-bert， 这些更多偏向于QA，而且训练数据咱也看不到），都不怎么合适（效果不怎么好）就想着索性自己就训练一个预训练模型吧，

---

提示：能够找到的比较多的预训练代码（基于pytorch）都是 mlm 的， wwm 的太少了或者根本没找到，所有自己根据huggingface的API顺藤摸瓜，找到了DataCollatorForWholeWordMask这个类，可以进行全词掩码的操作，但是基于分词工具的，由于本人是医学领域的，所以觉得 pkuseg 这个工具不错，大家可以根据自己的需求来选择，jieba啊，哈工大的pyltp啊 什么的。。。

一、WWM（全词掩码）和 MLM 的区别是什么？

大家可以根据这个图来自行揣摩。

二、预训练代码

代码如下：

# coding=utf-8
'''
    Whole word mask for bert
'''
import pkuseg
from transformers import BertConfig, BertForMaskedLM, DataCollatorForWholeWordMask,\
    BertTokenizer, TrainingArguments, Trainer
from torch.utils.data import Dataset
from tqdm import tqdm
import torch

class My_wwm_pretrain_dataset(Dataset):

    def __init__(self, path, tokenizer, dup_factor=5,max_length=512): # dup_factor : dynamic mask for 5 times
        self.examples = []
        with open(path,'r',encoding='utf-8') as f:
            total_data = f.readlines()
            with tqdm(total_data * dup_factor) as loader:
                for data in loader:
                    # clean data
                    data = data.replace('\n', '').replace('\r', '').replace('\t','').replace(' ','').replace('　', '')
                    chinese_ref = self.get_new_segment(data)
                    input_ids = tokenizer.encode_plus(data,truncation=True,max_length=max_length).input_ids
                    dict_data = {'input_ids' : input_ids, 'chinese_ref' : chinese_ref}
                    self.examples.append(dict_data)
                    loader.set_description(f'loading data')

    def get_new_segment(self,segment):
        """
            使用分词工具获取 whole word mask
            e.g [喜,欢]-> [喜，##欢]
        """
        seq_cws = seg.cut("".join(segment))  # 利用 pkuseg 进行医学领域分词
        chinese_ref = []
        index = 1
        for seq in seq_cws:
            for i, word in enumerate(seq):
                if i>0:
                    chinese_ref.append(index)
                index +=1
        return chinese_ref

    def __getitem__(self, index):
        return self.examples[index]

    def __len__(self):
        return len(self.examples)


if __name__ == '__main__':
    # configuration
    epoch = 100
    batch_size = 1
    pretrian_model = 'mc-bert-base'
    train_file = 'data/train.txt'
    save_epoch = 10 # every 10 epoch save checkpoint
    bert_file = '../../pretrained_models/' + pretrian_model
    tokenizer_model_path='../../pretrained_models/pkuseg_medical'
    #
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    seg = pkuseg.pkuseg(model_name=tokenizer_model_path)
    config = BertConfig.from_pretrained(bert_file)
    tokenizer = BertTokenizer.from_pretrained(bert_file)
    train_dataset = My_wwm_pretrain_dataset(train_file,tokenizer)
    model = BertForMaskedLM.from_pretrained(bert_file).to(device)
    print('No of parameters: ', model.num_parameters())
    data_collator = DataCollatorForWholeWordMask(
        tokenizer=tokenizer, mlm=True, mlm_probability=0.15
    )
    print('No. of lines: ', len(train_dataset))
    save_step = len(train_dataset) * save_epoch
    tot_step = int(len(train_dataset)/batch_size *  epoch)
    print(f'\n\t***** Running training *****\n'
          f'\tNum examples = {len(train_dataset)}\n'
          f'\tNum Epochs = {epoch}\n'
          f'\tBatch size = {batch_size}\n'
          f'\tTotal optimization steps = {tot_step}\n')

    # official training
    training_args = TrainingArguments(
        output_dir='./outputs/',
        overwrite_output_dir=True,
        num_train_epochs=epoch,
        per_device_train_batch_size=batch_size,
        save_steps=save_step,
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
    )

    trainer.train()
    trainer.save_model('pretrain_outputs/wwm/')

这里我们自己写了一个Dataset对数据进行分词，并封装，调用 transformers 的 DataCollatorForWholeWordMask API 帮我们进行后续 mask 的处理，利用 BertForMaskedLM 这个API 进行模型架构的生成和预训练，代码不是很多啦，也比较简单了，其中的原理和API源码后续可能会讲，小伙伴们也可以自行多多阅读。

---

总结

例如：以上就是今天要讲的内容，本文仅仅简单介绍了如何基于transformer进行bert的全词掩码预训练，帮助小伙伴们进行领域适配的工作，后续有时间会简单讲一下 DataCollatorForWholeWordMask 这个API 的源码和原理，以及BertForMaskedLM 这个模型工作流程。

喜欢的朋友可以给我点点关注，你的支持是我持续更新的动力！感谢！！
想要数据和模型的童鞋可以私信我。