NLP实践——利用自己的语料进行Mask Language Model预训练

1.1 什么是MLM

作为 Bert预训练的两大任务之一， MLM和 NSP大家应该并不陌生，其中 NSP任务在后续的一些预训练任务中经常被嫌弃，例如 Roberta中将 NSP任务直接放弃， Albert中将NSP替换成了句子顺序预测，这是由于 NSP作为一个分类任务过于简单，对模型的学习并没有太大的帮助，而 MLM则被多数预训练模型保留下来。由 Roberta的实验结果也可以证明， Bert的主要能力应该是来自于 MLM任务的训练。

Bert为代表的预训练语言模型是在大规模预料的基础上训练以获得的基础的学习能力，而实际应用时，我们所面临的预料或许具有某些特殊性，这就使得重新进行 MLM训练具有了必要性。

1.2 怎样进行MLM训练

MLM的训练，在不同的预训练模型中其实是有所不同的，今天介绍的内容以最基础的Bert为例。Bert的MLM是静态mask，而在后续的其他预训练模型中，这一策略通常被替换成了动态mask。除此之外还有whole word mask的模型，这些都不在今天的讨论范围内。

所谓mask language model的任务，通俗来讲，就是将句子中的一部分token替换掉，然后根据句子的剩余部分，试图去还原这部分被mask的token。

mask的比列一般是 15%，这一比例也被后续的多数模型所继承，而在最初BERT的论文中，没有对这一比例的界定给出具体的说明。在我印象中，似乎是知道后来同样是Google提出的T5模型的论文中，对此进行了解释，对mask的比例进行了实验，最终得出结论，15%的比例是最合理的（如果我记错了，还请指正）。

将 15%的token选出之后，并不是所有的都替换成[mask]标记符，而是从这 15%选出的部分中，将其中的 80%替换成[mask]， 10%替换成一个随机的token，剩下的 10%保留原来的token。这样做可以提高模型的鲁棒性。这个比例也可以自己控制。

到这里可能有同学要问了，既然有10%保留不变的话，为什么不干脆只选择15%*90% = 13.5%的token呢？如果看完后面的代码，就会很清楚的理解这个问题了。

因为MLM的任务是将选出的这15%的token全部进行预测，不管这个token是否被替换成了[mask]，也就是说，即使它被保留了原样，也还是需要被预测的。

介绍完了基础内容之后，接下来的内容，我将基于transformers模块，介绍如何进行mask language model的训练。

其实 transformers模块中，本身是提供了MLM训练任务的，模型都写好了，只需要调用它内置的trainer和 datasets模块即可。感兴趣的同学可以去huggingface的官网搜索相关教程。

然而我觉得 datasets每次调用的时候都要去写数据集的py文件，对 arrow的数据格式不熟悉的话还很容易出错，而且trainer我觉得也不是很好用，任何一点小小的修改都挺费劲（就是它以为它写的很完备，考虑了用户的所有需求，但是实际上有一些冗余的部分）。

所以我就参考它的实现方式，把它的代码拆解，又按照自己的方式重新组织了一下。

首先在写核心代码之前，先做好准备工作。
import 所有需要的模块：

import os
import json
import copy
from tqdm.notebook import tqdm

import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from transformers import BertForMaskedLM, BertTokenizerFast

写一个config类，将所有参数集中起来：

class Config:
    def __init__(self):
        pass

    def mlm_config(
        self,
        mlm_probability=0.15,
        special_tokens_mask=None,
        prob_replace_mask=0.8,
        prob_replace_rand=0.1,
        prob_keep_ori=0.1,
    ):
"""
        :param mlm_probability: 被mask的token总数
        :param special_token_mask: 特殊token
        :param prob_replace_mask: 被替换成[MASK]的token比率
        :param prob_replace_rand: 被随机替换成其他token比率
        :param prob_keep_ori: 保留原token的比率
"""
        assert sum([prob_replace_mask, prob_replace_rand, prob_keep_ori]) == 1,                 ValueError("Sum of the probs must equal to 1.")
        self.mlm_probability = mlm_probability
        self.special_tokens_mask = special_tokens_mask
        self.prob_replace_mask = prob_replace_mask
        self.prob_replace_rand = prob_replace_rand
        self.prob_keep_ori = prob_keep_ori

    def training_config(
        self,
        batch_size,
        epochs,
        learning_rate,
        weight_decay,
        device,
    ):
        self.batch_size = batch_size
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.device = device

    def io_config(
        self,
        from_path,
        save_path,
    ):
        self.from_path = from_path
        self.save_path = save_path

然后设置各种配置：

config = Config()
config.mlm_config()
config.training_config(batch_size=4, epochs=10, learning_rate=1e-5, weight_decay=0, device='cuda:0')
config.io_config(from_path='/data/BERTmodels/huggingface/chinese_wwm/',
                 save_path='./finetune_embedding_model/mlm/')

然后创建BERT模型，注意，这里的tokenizer就是一个普通的tokenizer，而BERT模型则是带了下游任务的BertForMaskedLM，它是transformers中写好的一个类，

bert_tokenizer = BertTokenizerFast.from_pretrained(config.from_path)
bert_mlm_model = BertForMaskedLM.from_pretrained(config.from_path)

由于舍弃了 datasets这个包，现在我们需要自己实现数据的输入了。方案就是使用torch的 Dataset类。这个类一般在构建 DataLoader的时候，会与一个聚合函数一起使用，以实现对batch的组织。而我这里偷个懒，就没有写聚合函数，batch的组织方法放在dataset中进行。

在这个类中，有一个mask tokens的方法，是从数据中选择出所有需要mask的token，并且采用三种mask方式中的一个，这个方法是从transformers中拿出来的，将其从类方法转为静态方法测试之后，再将其放在自己的这个类中为我们所用。仔细阅读这一段代码，也就可以回答1.2中提出的那个问题了。

取batch的原理很简单，一开始我们将原始数据deepcopy备份一下，然后每次从中截取一个batch的大小，这个时候的当前数据就少了一个batch，我们定义这个类的长度为当前长度除以batch size向下取整，所以当类的长度变为0的时候，就说明这一个epoch的所有step都已经执行结束，要进行下一个epoch的训练，此时，再将当前数据变为原始数据，就可以实现对epoch的循环了。

class TrainDataset(Dataset):
"""
    注意：由于没有使用data_collator，batch放在dataset里边做，
    因而在dataloader出来的结果会多套一层batch维度，传入模型时注意squeeze掉
"""
    def __init__(self, input_texts, tokenizer, config):
        self.input_texts = input_texts
        self.tokenizer = tokenizer
        self.config = config
        self.ori_inputs = copy.deepcopy(input_texts)

    def __len__(self):
        return len(self.input_texts) // self.config.batch_size

    def __getitem__(self, idx):
        batch_text = self.input_texts[: self.config.batch_size]
        features = self.tokenizer(batch_text, max_length=512, truncation=True, padding=True, return_tensors='pt')
        inputs, labels = self.mask_tokens(features['input_ids'])
        batch = {"inputs": inputs, "labels": labels}
        self.input_texts = self.input_texts[self.config.batch_size: ]
        if not len(self):
            self.input_texts = self.ori_inputs

        return batch

    def mask_tokens(self, inputs):
"""
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.

"""
        labels = inputs.clone()

        probability_matrix = torch.full(labels.shape, self.config.mlm_probability)
        if self.config.special_tokens_mask is None:
            special_tokens_mask = [
                self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
            ]
            special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
        else:
            special_tokens_mask = self.config.special_tokens_mask.bool()

        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100

        indices_replaced = torch.bernoulli(torch.full(labels.shape, self.config.prob_replace_mask)).bool() & masked_indices
        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)

        current_prob = self.config.prob_replace_rand / (1 - self.config.prob_replace_mask)
        indices_random = torch.bernoulli(torch.full(labels.shape, current_prob)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)
        inputs[indices_random] = random_words[indices_random]

        return inputs, labels

然后取一些用于训练的语料，格式很简单，就是把所有文本放在一个list里边，注意长度不要超过512个token，不然多出来的部分就浪费掉了。可以做适当的预处理。

[
    "这是一条文本",
    "这是另一条文本",
    ...,
]

然后构建dataloader：

train_dataset = TrainDataset(training_texts, bert_tokenizer, config)
train_dataloader = DataLoader(train_dataset)

构建一个训练方法，输入参数分别是我们实例化好的待训练模型，数据集，还有config：

def train(model, train_dataloader, config):
"""
    训练
    :param model: nn.Module
    :param train_dataloader: DataLoader
    :param config: Config
    ---------------
    ver: 2021-11-08
    by: changhongyu
"""
    assert config.device.startswith('cuda') or config.device == 'cpu', ValueError("Invalid device.")
    device = torch.device(config.device)

    model.to(device)

    if not len(train_dataloader):
        raise EOFError("Empty train_dataloader.")

    param_optimizer = list(model.named_parameters())
    no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {"params": [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], "weight_decay": 0.01},
        {"params": [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], "weight_decay": 0.0}]

    optimizer = AdamW(params=optimizer_grouped_parameters, lr=config.learning_rate, weight_decay=config.weight_decay)

    for cur_epc in tqdm(range(int(config.epochs)), desc="Epoch"):
        training_loss = 0
        print("Epoch: {}".format(cur_epc+1))
        model.train()
        for step, batch in enumerate(tqdm(train_dataloader, desc='Step')):
            input_ids = batch['inputs'].squeeze(0).to(device)
            labels = batch['labels'].squeeze(0).to(device)
            loss = model(input_ids=input_ids, labels=labels).loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            model.zero_grad()
            training_loss += loss.item()
        print("Training loss: ", training_loss)

调用它训练几轮：

train(model=bert_mlm_model, train_dataloader=train_dataloader, config=config)

使用过预训练模型的同学应该都了解，普通的bert有两项输出，分别是每一个token对应的768维编码结果，以及用于表征整个句子的句子特征。

而这个句子特征，是由模型中的一个Pooler模块对原句池化得来的。可是，这个Pooler的训练，并不是由MLM任务来的，而是由NSP任务中来的。

由于没有NSP任务，所以无法对Pooler进行训练，故而没有必要在模型中加入Pooler。

所以在保存的时候需要分别保存embedding和encoder，
加载的时候也需要分别读取embedding和encoder，这样训练出来的模型拿不到CLS层的句子表征，如果需要的话，可以手动pooling。

torch.save(bert_mlm_model.bert.embeddings.state_dict(), os.path.join(config.save_path, 'bert_mlm_ep_{}_eb.bin'.format(config.epochs)))
torch.save(bert_mlm_model.bert.encoder.state_dict(), os.path.join(config.save_path, 'bert_mlm_ep_{}_ec.bin'.format(config.epochs)))

加载的话，也是实例化完bert模型之后，用bert的embedding组件和encoder组件分别读取这两个权重文件即可。

到这里，本期内容就全部结束了，希望看完这篇博客的同学，能够对Bert的基础原理有更深入的了解。

Original: https://blog.csdn.net/weixin_44826203/article/details/121439850
Author: 常鸿宇
Title: NLP实践——利用自己的语料进行Mask Language Model预训练