基于BERT预训练模型的命名体识别任务

1 引言

各位朋友大家好，欢迎来到月来客栈，我是掌柜空字符。

经过前面一系列文章的介绍相信大家对于BERT模型已经有了非常清晰的认识，不过为了满足不同人群的学习需求，在这篇文章中掌柜将会介绍基于BERT预训练模型的第五个下游任务场景，即如何完成命名体识别（Named Entity Recognition, NER）任务。所谓命名体指的是给模型输入一句文本，最后需要模型将其中的实体（例如人名、地名、组织等等）标记出来。

例如：

对于这样一个形式的任务我们应该怎么来构建模型呢？

通常来讲，对于任意一个NLP任务来说模型最后所要完成的基本上都是一个分类任务，尽管表面上看起来可能不太像。根据给出的标签来看，对于原始句子中的每个字符来说其都有一个对应的类别标签，因此对于NER任务来说只需要对原始句子里的每个字符进行分类即可，然后再将预测后的结果进行后处理便能够得到句子从存在的相应实体。

以下所有完整示例代码均可从仓库 https://github.com/moon-hotel/BertWithPretrained 中获取！

2 模型构建

2.1 构建原理

正如上面所说，对于命名体识别这个任务场景来说其本质上依旧可以归结为分类任务，只是关键在于如何构建这个任务以及整个数据集。对于这个任务场景来说，其整体原理如图1所示。

图 1. 命名体识别原理图

如图1所示便是一个基于BERT预训练模型的NER任务原理图。从图中可以看出原始数据输入为一个句子，我们只需要在句子的首尾分别加上和，然后输入到模型当中进行特征提取并最终通过一个分类层对输出的每个Token进行分类即可，最后只需要对各个Token的预测结果进行后处理便能够实现整个NER任务。

到此，对于问答选择整个模型的原理我们算是清楚了，下面首先来看如何构造数据集。

2.2 语料介绍

在这里，我们使用到的是一个中文命名体识别数据集[1]，如下所示便是原始数据的存储形式：

其中每一行包含一个字符和其对应的所属类别，表示该类实体的开始标志，表示该类实体的延续标志。例如对于13-16行来说其对应了“黄州”和“赤壁”这两个实体。同时，对于这个数据集来说，其一共包含有3类实体（人名、地名和组织），因此其对应的分类总数便为7，如下所示：

接下里，便可以根据需要来构造模型训练时所需要的数据集了。

2.3 数据集预览

同样，在正式介绍如何构建数据集之前我们先通过一张图来了解一下整个大致构建的流程。假如我们现在有两个样本构成了一个batch，那么其整个数据的处理过程则如图2所示。

图 2. 数据集构造流程图

如图2所示，首先我们需要将原始的语料格式化成一个一个的句子；然后再将其转换成对应的Token ID，并在首尾分别加上和,同时对于标签来说因为首尾加入的这两个Token不需要进行预测所以在对应位置上需要标识出来以便在计算损失的时候对其忽略；最后，在输入模型之前在对其进行Padding处理和构造得到相应的PaddingMask向量。

下面开始分别对各部分的处理进行详细介绍。

2.4 数据集构建

在说完数据集构造的整理思路后，下面我们就来正式编码实现整个数据集的构造过程。同样，对于数据预处理部分我们可以继续继承之前文本分类处理中的类，然后再稍微修改其中的部分方法即可。同时，由于在前两个示例[2] [3]中已经就tokenize和词表构建等内容做了详细的介绍，所以这部分内容就不再赘述。

第1步：格式化样本和Tokenize

如图2过程所示，需要对原始样本进行格式化以及转换得到每个序列对应的Token id，下面首先是在函数中来定义如何完成上述步骤，实现代码如下：

在上述代码中，第1-8行是类的初始化方法，其中第4行是指定对应的实体字典，即第2.2节中末尾的内容，第5-6行是分别指定分类数和需要被忽略的Token ID。

这里值得一提的是，因为分类的类别数中有0这个类别，因此就不能直接拿（对应的Token ID为0）来作为Ignore Id了，而对于前面介绍的在MLM任务中我们便可以用0来同时作为Padding Id和Ignore Id。

第10行是对方法处理后的结果进行缓存，只要不改变相关参数那么每次运行模型时都会先找对应的缓存文件从而减少等待时间，具体讲解可参见文章[4]；第12行是一次性读取原始数据中的所有行；第13行用于保存预处理结束后的数据；第14行记录所有样本中的最大长度；第15行用于记录每个样本点的Token Id；第16行用于保存每个原始样本；第17行用于保存每个样本对应的预测标签；第18行用于记录每个样本的原始标签，主要用于观察处理时的中间结果。

在上述代码中，第1行开始遍历原始数据中的每一行；第5行表示已经将上一个样本处理完毕；第6-8行是判断长度是否超过最大长度；第9行用来得到所有句子的最大长度；第10-13行是用来构造输入模型的输入和正确标签；第21-24行分别是对当前样本中的每个Token进行相应的处理。

以下所有完整示例代码均可从仓库 https://github.com/moon-hotel/BertWithPretrained 中获取！

最后，对于前面两个样本来说，经过方法处理后便会得到如下所示结果：

第2步：Padding处理

在处理得到每个样本对应的Token Id和标签后，我们再来定义一个方法对每个batch中的数据集进行padding处理，代码如下：

在上述代码中，第3-6行用来得到每个batch中的各个部分；第7-10行是对模型的输入进行padding处理；第11-14行是对标签值进行padding处理；这里需要注意的是两处padding的值并不一样，如果非要保持一样的话可以都用来作为padding值；第15行则是返回对应处理完成的结果。

2.5 使用示例

在完成数据集构造部分的相关代码实现之后，便可以通过如下所示的方式进行使用，代码如下：

在上述代码运行结束之后便可以得到类似如下所示的输出结果：

到此，对于整个数据集的构建流程及使用方式就介绍完了，下面再来看模型的实现部分。

3 命名体识别模型

3.1 前向传播

正如第1节内容所介绍，我们只需要在原始BERT模型的基础上再加一个对所有Token进行分类的分类层即可，因此这部分代码相对来说也比较容易理解。首先需要在目录下新建一个模块，并完成整个模型的初始化和前向传播过程，代码如下：

在上述代码中，第1行用于导入对应的BERT模型；第8-11行用于根据参数返回预训练模型或者是初始化模型；第12-13行则是完成命名体识别的分类任务。

接着是模型的前向传播过程，实现代码如下：

在上述代码中，第2-6行是模型的输入部分，由于任务是单句输入所以不用输入，同时的形状是，的形状是，的形状为；第7-9行则是返回BERT模型的输出结果，并取最后一层的输出形状为；第10-11行则是进行分类处理，输出结果形状为；第12-17行则是根据条件返回相应的结果，这里注意的是需要指定。

3.2 模型训练

模型配置类实现

对于模型训练这部分内容来说，首先我们需要在目录下新建一个模块，并新建一个配置类来管理整个模型需要用到的参数，代码实现如下：

在上述代码中，第3-11行用来获取工程的目录以及预训练模型、数据集等的绝对路径；第12-14行分别用来指定保存模型的名称、tensorboard可视化文件和日志目录；第15行指定原始数据中的分割符；第16行指定实体的类别映射关系；第17-18行分别用来指定分类类别数和忽略Token的id；第19行用来初始化日志打印方法，详见文章训练模型时如何保存训练日志[5]？

所有完整示例代码均可从仓库 https://github.com/moon-hotel/BertWithPretrained 中获取！

评价指标计算实现

因为在模型训练过程中需要计算相关的评价指标，如准确率、精确率和召回率等[6]，因此需要对这部分进行实现，代码如下：

在上述代码中，第1行的形状为，的形状为；第2-4行是用来将原始的预测值和正确值转化为一个一维列表；第6-9行用来去掉需要进行忽略的预测结果，即真实标签中padding部分、和对应的部分；第10行是返回准确率和对应的预测-真实结果，用于后续通过来计算模型的准确率和精确率等。

对于函数来说，其使用示例如下：

预测结果格式化实现

为了能够在模型训练或推理过程中输入模型的预测结果，因此我们需要实现3个辅助函数来完成。首先需要实现根据和来得到每个预测值对应的实体标签，代码如下：

在上述代码中，第1行的形状为；第2行用来得到所有的实体，即；第3行表示先忽略掉第1个Token（为），并变形为；第4行用来得到对应的预测结果和概率值，两者的形状均为；第5行用来去掉中的；第9-17行则是分别用来得到每个位置上的预测标签和对应的概率值，并且从这里可以看出的作用就是用来确定预测结果中哪位位置上的Token不是padding的部分。

实现完成后，对于函数来说，使用示例如下：

进一步，在得到每个输入句子的预测结果后，还需要将其进行格式化处理得到最终的预测结果，实现代码如下：

在上述代码中，第2行用来得到预测结果分割的标签，即；第4-18行则是依次遍历每个样本及样本中的每个预测结果，并得到一个连续的实体序列及其对应的实体类别，其中第7行中加上是为了处理当输入句子最后一个预测结果仍为时的情况。

最后，在实现完上述代码后，便可以通过如下方式进行使用：

输出结果如下：

当然， 在使用时需要将函数和结合起来，实现代码如下：

模型训练实现

在完成上述所有铺垫之后，便可以来实现模型的训练部分，代码如下（下面只摘录核心部分进行介绍）：

在上述代码中，第2行用来实例化一个模型对象；第3-8行用来判断本地是否已有本地模型存在进行追加训练；第9-10行是返回得到对应的训练集、验证集和测试集；第11行是实例化一个优化器；第14行是开始对模型进行训练；第18-21行是完成模型的前向传播及损失计算过程，其中的形状为，的形状为，的形状为，同时输出结果的形状为；第26-28行是将训练时的巡视和准确率通过Tensorboard可视化；第29-30行是输出训练过程中的预测结果；第33-36行则是把测试集上的准去率进行可视化。

上述代码完成后，便可以通过如下示例进行使用：

训练时的输出结果如下所示：

3.3 模型推理

在完成模型训练部分的代码实现后最后再来看模型的推理实现。通常，模型训练完成后都会得到一个持久化的模型参数，通过载入参数模型便可以对新样本进行推理预测，实现代码如下：

在上述代码中，第2-8行用来先随机实例化一个模型，然后再通过本地的模型参数来重新赋值模型中的参数；第11-12行则是实例化一个数据集载入类，并将新输入样本处理成模型所接受的输入形式；第15-17行则是完成相应的前向传播过程，并对输出结果进行格式化。

上述代码完成后便可以通过如下示例进行使用：

推理时的输出结果如下所示：

4 总结

在这篇文章中，掌柜首先简单介绍了什么是命名体识别任务；然后详细介绍了基于BERT预训练模型的NER模型构建原理，并同时一步一步地介绍了如何构造模型训练时的数据集；最后介绍了整个模型的实现部分，包括模型的前向传播、训练和推理过程等。

引用

[1] https://github.com/zjy-ucas/ChineseNER

[7]https://github.com/huggingface/transformers/tree/main/examples/pytorch/token-classification