更新于 2026年7月25日

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

例如:

1 句子涂伊说如果有机会他想去黄州赤壁看一看
2 标签['B-PER', 'I-PER', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-LOC', 'I-LOC', 'B-LOC', 'I-LOC', 'O', 'O', 'O', 'O']
3 实体涂伊人名)、黄州地名)、赤壁地名

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

10.1 任务构造原理#

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

图 10-1. 命名体识别原理图
图 10-1. 命名体识别原理图

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

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

10.2 数据预处理#

10.2.1 语料介绍#

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

 1 涂 B-PER
 2 伊 I-PER
 3 说 O
 4  O
 5 如 O
 6 果 O
 7 有 O
 8 机 O
 9 会 O
10 他 O
11 想 O
12 去 O
13 黄 B-LOC
14 州 I-LOC
15 赤 B-LOC
16 壁 I-LOC
17 看 O
18 一 O
19 看 O
20  O

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

1 {'O': 0, 'B-ORG': 1,'B-LOC': 2,'B-PER': 3,'I-ORG': 4, 'I-LOC': 5, 'I-PER': 6}

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

10.2.2 数据集预览#

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

图 10-2. 数据集构造流程图
图 10-2. 数据集构造流程图

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

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

10.2.3 数据集构造#

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

第1步:格式化样本和 Tokenize

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

 1 class LoadChineseNERDataset(LoadSingleSentenceClassificationDataset):
 2     def __init__(self,entities=None,num_labels=None,ign_idx=-100, **kwargs):
 3         super(LoadChineseNERDataset, self).__init__(**kwargs)
 4         self.entities = entities
 5         self.num_labels = num_labels
 6         self.IGNORE_IDX = ign_idx
 7         if self.entities is None or self.num_labels is None:
 8             raise ValueError(f"类{self.__class__.__name__}中参数entities或num_labels不能为空!")
 9 
10     @cache
11     def data_process(self, filepath, postfix='cache'):
12         raw_iter = open(filepath, encoding="utf8").readlines()
13         data = []
14         max_len = 0
15         tmp_token_ids = []
16         tmp_sentence = ""
17         tmp_label = []
18         tmp_entity = []

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

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

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

 1         for raw in tqdm(raw_iter, ncols=80):
 2             line = raw.rstrip("\n").split(self.split_sep)
 3             if len(line) != 1 and len(line) != 2:
 4                 raise ValueError(f"数据标注有误{line}")
 5             if len(line) == 1: # 表示得到一个完整的token id样本
 6                 if len(tmp_token_ids) > self.max_position_embeddings - 2:
 7                     tmp_token_ids = tmp_token_ids[:self.max_position_embeddings - 2]
 8                     tmp_label = tmp_label[:self.max_position_embeddings - 2]
 9                 max_len = max(max_len, len(tmp_label) + 2)
10                 token_ids = torch.tensor([self.CLS_IDX] + tmp_token_ids +
11                                          [self.SEP_IDX], dtype=torch.long)
12                 labels = torch.tensor([self.IGNORE_IDX] + tmp_label +
13                                       [self.IGNORE_IDX], dtype=torch.long)
14                 data.append([tmp_sentence, token_ids, labels])
15                 assert len(tmp_token_ids) == len(tmp_label)
16                 tmp_token_ids = []
17                 tmp_sentence = ""
18                 tmp_label = []
19                 tmp_entity = []
20                 continue
21             tmp_sentence += line[0]
22             tmp_token_ids.append(self.vocab[line[0]])
23             tmp_label.append(self.entities[line[-1]])
24             tmp_entity.append(line[-1])
25         return data, max_len

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

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

 1  - DEBUG: ### 样本构造结果为:
 2  - DEBUG:  ## 句子: 涂伊说,如果有机会他想去黄州赤壁看一看!
 3  - DEBUG:  ## 实体: ['B-PER', 'I-PER', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-LOC', 'I-LOC', 'B-LOC', 'I-LOC', 'O', 'O', 'O', 'O']
 4  - DEBUG:  ## input_ids: [101, 3864, 823, 6432, 8024, 1963, 3362, 3300, 3322, 833, 800, 2682, 1343, 7942, 2336, 6619, 1880, 4692, 671, 4692, 8013, 102]
 5  - DEBUG:   ## label: [-100, 3, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 5, 2, 5, 0, 0, 0, 0, -100]
 6  - DEBUG:  ================================
 7 
 8  - DEBUG:  ### 样本构造结果为:
 9  - DEBUG:   ## 句子: 丽江、大理、九寨沟、黄龙等都是涂伊想去的地方!
10  - DEBUG:   ## 实体: ['B-LOC', 'I-LOC', 'O', 'B-LOC', 'I-LOC', 'O', 'B-LOC', 'I-LOC', 'I-LOC', 'O', 'B-LOC', 'I-LOC', 'O', 'O', 'O', 'B-PER', 'I-PER', 'O', 'O', 'O', 'O', 'O', 'O']
11  - DEBUG:   ## input_ids: [101, 714, 3736, 510, 1920, 4415, 510, 736, 2181, 3765, 510, 7942, 7987, 5023, 6963, 3221, 3864, 823, 2682, 1343, 4638, 1765, 3175, 8013, 102]
12  - DEBUG:   ## label: [-100, 2, 5, 0, 2, 5, 0, 2, 5, 5, 0, 2, 5, 0, 0, 0, 3, 6, 0, 0, 0, 0, 0, 0, -100]

第2步:Padding处理

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

 1     def generate_batch(self, data_batch):
 2         batch_sentence, batch_token_ids, batch_label = [], [], []
 3         for (sen, token_ids, label) in data_batch:  # 开始对一个batch中的每一个样本进行处理。
 4             batch_sentence.append(sen)
 5             batch_token_ids.append(token_ids)
 6             batch_label.append(label)
 7         batch_token_ids = pad_sequence(batch_token_ids,#[batch_size,max_len]
 8                                        padding_value=self.PAD_IDX,
 9                                        batch_first=False,
10                                        max_len=self.max_sen_len)
11         batch_label = pad_sequence(batch_label,  # [batch_size,max_len]
12                                    padding_value=self.IGNORE_IDX,
13                                    batch_first=False,
14                                    max_len=self.max_sen_len)
15         return batch_sentence, batch_token_ids, batch_label

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

第3步:使用示例

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

 1 class ModelConfig:
 2     def __init__(self):
 3         self.project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 4         self.dataset_dir=os.path.join(self.project_dir,'data','ChineseNER')
 5         self.vocab_path=os.path.join(self.pretrained_model_dir,'vocab.txt')
 6         self.train_file_path = os.path.join(self.dataset_dir, 'train.txt')
 7         self.val_file_path = os.path.join(self.dataset_dir, 'dev.txt')
 8         self.test_file_path = os.path.join(self.dataset_dir, 'test.txt')
 9         self.split_sep = ' '
10         self.entities = {'O': 0, 'B-ORG': 1, 'B-LOC': 2, 'B-PER': 3, 'I-ORG': 4, 'I-LOC': 5, 'I-PER': 6}
11         self.num_labels = len(self.entities)
12         self.ignore_idx = -100
13         # ......
14         
15 if __name__ == '__main__':
16     model_config = ModelConfig()
17     data_loader = LoadChineseNERDataset(
18         entities=model_config.entities,
19       	num_labels=model_config.num_labels,
20         ignore_idx=model_config.ignore_idx,
21       	vocab_path=model_config.vocab_path,
22         tokenizer=BertTokenizer.from_pretrained(
23           model_config.pretrained_model_dir).tokenize,
24         batch_size=model_config.batch_size,
25         max_sen_len=model_config.max_sen_len,
26         split_sep=model_config.split_sep,
27         max_position_embeddings=model_config.max_position_embeddings,
28         pad_index=model_config.pad_token_id,
29         is_sample_shuffle=model_config.is_sample_shuffle)
30     test_iter = data_loader.load_train_val_test_data(
31       	train_file_path=model_config.train_file_path,
32 		val_file_path=model_config.val_file_path,
33       	test_file_path=model_config.test_file_path, only_test=True)
34           
35     for sen, token_ids, labels in test_iter:
36         print(sen)
37         print("Input Token:", token_ids.shape)
38         print(token_ids.transpose(0, 1))
39         print("Attention Mask: \n",(token_ids == model_config.pad_token_id).transpose(0, 1))
40         print("Labels:\n", labels.shape)  # [src_len,batch_size]
41         print(labels.transpose(0, 1))

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

 1 ['涂伊说,如果有机会他想去黄州赤壁看一看!', '丽江、大理、九寨沟、黄龙等都是涂伊想去的地方!']
 2 Input Token: torch.Size([25, 2])
 3 tensor([[ 101, 3864, 823, 6432, 8024, 1963, 3362,3300,3322,833,800,2682,
 4          1343, 7942, 2336,6619, 1880, 4692,  671,4692,8013,102,  0,   0, 0],
 5         [ 101,  714,  3736, 510, 1920, 4415,  510, 736,2181,3765, 510, 7942,
 6     7987, 5023, 6963,3221, 3864, 823, 2682,1343,4638,1765, 3175,8013, 102]])
 7 Attention MaskPadding Mask): 
 8 tensor([[False, False, False, False, False, False, False, False, False, False,
 9          False, False, False, False, False, False, False, False, False, False,
10          False, False,  True,  True,  True],
11         [False, False, False, False, False, False, False, False, False, False,
12          False, False, False, False, False, False, False, False, False, False,
13          False, False, False, False, False]])
14 Labels: torch.Size([25, 2])
15 tensor([[-100,  3,  6,  0,  0,  0,  0,  0,  0,  0,  0,  0,
16             0,  2,  5,  2,  5,  0,  0,  0,  0,-100,-100,-100, -100],
17         [-100,  2,  5,  0,  2,  5,  0,  2,  5,   5,   0,   2,
18             5,  0,  0,  0,  3,  6,  0,  0,  0,   0,   0,   0,  -100]])

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

10.3 命名体识别任务#

10.3.1 前向传播#

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

 1 from ..BasicBert.Bert import BertModel
 2 import torch.nn as nn
 3 
 4 class BertForTokenClassification(nn.Module):
 5     def __init__(self, config, bert_pretrained_model_dir=None):
 6         super(BertForTokenClassification, self).__init__()
 7         self.num_labels = config.num_labels
 8         if bert_pretrained_model_dir is not None:
 9             self.bert=BertModel.from_pretrained(config,pretrained_model_dir)
10         else:
11             self.bert = BertModel(config)
12         self.dropout = nn.Dropout(config.hidden_dropout_prob)
13         self.classifier = nn.Linear(config.hidden_size, self.num_labels)
14         self.config = config

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

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

 1     def forward(self,
 2                 input_ids=None,
 3                 attention_mask=None,
 4                 token_type_ids=None,
 5                 position_ids=None,
 6                 labels=None):
 7         _, all_encoder_outputs = self.bert(input_ids=input_ids,attention_mask=attention_mask,
 8                token_type_ids=token_type_ids,position_ids=position_ids)  # [batch_size,hidden_size]
 9         sequence_output = all_encoder_outputs[-1]  # 取最后一层
10         sequence_output = self.dropout(sequence_output)
11         logits = self.classifier(sequence_output)
12         if labels is not None:  # [src_len,batch_size]
13             loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.ignore_idx)
14             loss=loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
15             return loss, logits
16         else:
17             return logits

在上述代码中,第2~6行是模型的输入部分,由于任务是单句输入所以 token_type_ids 不用输入,同时 input_ids 的形状是[src_len,batch_size],attention_mask 的形状是[batch_size, src_len],labels 的形状为[src_len,batch_size];第7~9行则是返回BERT模型的输出结果,并取最后一层的输出形状为[src_len, batch_size, hidden_size];第10~11行则是进行分类处理,输出结果形状为[src_len, batch_size, num_labels];第12~17行则是根据条件返回相应的结果,这里注意的是需要指定 ignore_index

10.3.2 模型训练#

(1)模型配置类实现

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

 1 class ModelConfig:
 2     def __init__(self):
 3         self.project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 4         self.dataset_dir=os.path.join(self.project_dir,'data','ChineseNER')
 5         self.pretrained_model_dir = os.path.join(self.project_dir, "bert_base_chinese")
 6         self.vocab_path=os.path.join(self.pretrained_model_dir,'vocab.txt')
 7         self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 8         self.train_file_path = os.path.join(self.dataset_dir, 'train.txt')
 9         self.val_file_path = os.path.join(self.dataset_dir, 'dev.txt')
10         self.test_file_path = os.path.join(self.dataset_dir, 'test.txt')
11         self.model_save_dir = os.path.join(self.project_dir, 'cache')
12         self.model_save_name = "ner_model.pt"
13         self.writer = SummaryWriter("runs")
14         self.logs_save_dir = os.path.join(self.project_dir, 'logs')
15         self.split_sep = ' '
16         self.entities = {'O': 0, 'B-ORG': 1, 'B-LOC': 2, 'B-PER': 3, 'I-ORG': 4, 'I-LOC': 5, 'I-PER': 6}
17         self.num_labels = len(self.entities)
18         self.ignore_idx = -100
19         logger_init(log_file_name='ner', log_level=logging.DEBUG,
20                     log_dir=self.logs_save_dir)

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

(2)评价指标计算实现

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

 1 def accuracy(logits, y_true, ignore_idx=-100):
 2     y_pred = logits.transpose(0, 1).argmax(axis=2).reshape(-1).tolist()
 3     # 将 [src_len,batch_size,num_labels] 转成 [batch_size, src_len,num_labels]
 4     y_true = y_true.transpose(0, 1).reshape(-1).tolist()
 5     real_pred, real_true = [], []
 6     for item in zip(y_pred, y_true):
 7         if item[1] != ignore_idx:
 8             real_pred.append(item[0])
 9             real_true.append(item[1])
10     return accuracy_score(real_true, real_pred), real_true, real_pred

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

对于accuracy函数来说,其使用示例如下:

1 y_true = torch.tensor([[-100, 0, 0, 1, -100],
2          [-100, 2, 0, -100, -100]]).transpose(0, 1)#[src_len,batch_size]
3 logits = torch.tensor([[[0.5, 0.1, 0.2],[0.5, 0.4, 0.1],[0.7, 0.2, 0.3],
4                         [0.5, 0.7,0.2],[0.1, 0.2,0.5]],
5                        [[0.3, 0.2, 0.5],[0.7, 0.2, 0.4],[0.8, 0.1, 0.3],
6                         [0.9, 0.2,0.1],[0.1, 0.5,0.2]]])
7 logits = logits.transpose(0, 1) # [src_len,batch_size,num_labels]
8 print(accuracy(logits, y_true, -100)) # (0.8, 4, 5)

(3)预测结果格式化实现

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

 1 def get_ner_tags(logits, token_ids, entities, SEP_IDX=102):
 2     label_list = list(entities.keys())
 3     logits = logits[1:].transpose(0, 1) 
 4     prob, y_pred = torch.max(logits, dim=-1) 
 5     token_ids = token_ids[1:].transpose(0, 1)#[ batch_size,src_len-1],去掉[cls]
 6     assert y_pred.shape == token_ids.shape
 7     labels = []
 8     probs = []
 9     for sample in zip(y_pred, token_ids, prob):
10         tmp_label, tmp_prob = [], []
11         for item in zip(*sample):
12             if item[1] == SEP_IDX:  # 忽略最后一个[SEP]字符
13                 break
14             tmp_label.append(label_list[item[0]])
15             tmp_prob.append(item[2].item())
16         labels.append(tmp_label)
17         probs.append(tmp_prob)
18     return labels, probs

在上述代码中,第1行 token_ids 的形状为[src_len,batch_size];第2行用来得到所有的实体,即[‘O’, ‘B-ORG’, ‘B-LOC’, ‘B-PER’, ‘I-ORG’, ‘I-LOC’, ‘I-PER’];第3行表示先忽略掉第1个Token(为[CLS]),并变形为[batch_size,src_len-1,num_samples];第4行用来得到对应的预测结果和概率值,两者的形状均为[batch_size,src_len-1];第5行用来去掉token_ids中的[CLS];第9~17行则是分别用来得到每个位置上的预测标签和对应的概率值,并且从这里可以看出token_ids的作用就是用来确定预测结果中哪位位置上的Token不是padding的部分。

实现完成后,对于get_ner_tags函数来说,使用示例如下:

 1 entities={'O:0,'B-ORG:1,'B-LOC':2,'B-PER':3,'I-ORG':4,'I-LOC':5,'I-PER':6}
 2 logits = torch.tensor([[[0.4, 0.7, 0.2],[0.5, 0.4, 0.1],
 3                         [0.1, 0.2, 0.3],[0.5, 0.7, 0.2],[0.1, 0.2, 0.5]],
 4                        [[0.3, 0.2, 0.5],[0.7, 0.8, 0.4],[0.1, 0.1, 0.3],
 5                         [0.9, 0.2, 0.1],[0.1, 0.5,0.2]]])
 6 logits = logits.transpose(0, 1)  # [src_len,batch_size,num_samples]
 7 token_ids = torch.tensor([[101, 2769, 511, 102, 0], #[src_len,batch_size]
 8                           [101, 56, 33, 22, 102]]).transpose(0, 1)  
 9 labels, probs = get_ner_tags(logits, token_ids, entities)
10 #
11 # [['O', 'B-LOC'], ['B-ORG', 'B-LOC', 'O']]
12 # [[0.5, 0.300000011], [0.80000001, 0.3000000119, 0.899999976]]

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

 1 def pretty_print(sentences, labels, entities):
 2     sep_tag = [tag for tag in list(entities.keys()) if 'I' not in tag]
 3     result = []
 4     for sen, label in zip(sentences, labels):
 5         logging.info(f"句子:{sen}")
 6         last_tag = None
 7         for item in zip(sen + "O", label + ['O']):
 8             if item[1] in sep_tag:  #
 9                 if len(result) > 0:
10                     entity = "".join(result)
11                     logging.info(f"\t{entity}:  {last_tag.split('-')[-1]}")
12                     result = []
13                 if item[1] != 'O':
14                     result.append(item[0])
15                     last_tag = item[1]
16             else:
17                 result.append(item[0])
18                 last_tag = item[1]

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

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

 1 labels = [['B-PER','I-PER', 'O','O','O','O','O','O','O','O','O','O',
 2            'B-LOC','I-LOC','B-LOC','I-LOC'],
 3           ['B-LOC','I-LOC','O','B-LOC','I-LOC','O','B-LOC','I-LOC',
 4            'I-LOC','O','B-LOC','I-LOC','O','O','O','B-PER','I-PER',
 5            'O','O','O','O','O','O']]
 6 sentences=["涂伊说,如果有机会他想去赤壁看一看!",
 7            "丽江、大理、九寨沟、黄龙等都是涂伊想去的地方!"]
 8 entities = {'O': 0, 'B-ORG': 1, 'B-LOC': 2, 'B-PER': 3, 
 9             'I-ORG': 4, 'I-LOC': 5, 'I-PER': 6}
10 pretty_print(sentences, labels, entities)

输出结果如下:

 1 句子涂伊说如果有机会他想去黄州赤壁看一看
 2 涂伊:  PER
 3 黄州:  LOC
 4 赤壁:  LOC
 5 句子丽江大理九寨沟黄龙等都是涂伊想去的地方
 6 丽江:  LOC
 7 大理:  LOC
 8 九寨沟:  LOC
 9 黄龙:  LOC
10 涂伊:  PER

当然, 在使用时需要将 get_ner_tags 函数和 pretty_print 结合起来,实现代码如下:

1 def show_result(sentences, logits, token_ids, entities):
2     labels, _ = get_ner_tags(logits, token_ids, entities)
3     pretty_print(sentences, labels, entities)

(3)模型训练实现

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

 1  def train(config):
 2     model = BertForTokenClassification(config,config.pretrained_model_dir)
 3     model_save_path = os.path.join(config.model_save_dir,config.model_save_name)
 4     global_steps = 0
 5     if os.path.exists(model_save_path):
 6         checkpoint = torch.load(model_save_path)
 7         loaded_paras = checkpoint['model_state_dict']
 8         model.load_state_dict(loaded_paras)
 9     data_loader = LoadChineseNERDataset(entities=config.entities,...)
10     train_iter,test_iter,val_iter=data_loader.load_train_val_test_data(..)
11     optimizer=torch.optim.Adam(model.parameters(),lr=config.learning_rate)
12     model.train()
13     max_acc = 0
14     for epoch in range(config.epochs):
15         losses = 0
16         start_time = time.time()
17         for idx, (sen, token_ids, labels) in enumerate(train_iter):
18             padding_mask=(token_ids==data_loader.PAD_IDX).transpose(0, 1)
19             loss, logits = model(input_ids = token_ids, attention_mask=padding_mask,
20                                  token_type_ids=None,position_ids=None,
21                                  labels=labels)  # [src_len, batch_size]
22             # ......
23             losses += loss.item()
24             global_steps += 1
25             acc, _, _ = accuracy(logits, labels, config.ignore_idx)
26             if idx % 20 == 0:
27                 config.writer.add_scalar('Training/Loss', loss.item(), global_steps)
28                 config.writer.add_scalar('Training/Acc', acc, global_steps)
29             if idx % 100 == 0:
30                 show_result(sen[:10], logits[:, :10], token_ids[:, :10], config.entities)
31         end_time = time.time()
32         train_loss = losses / len(train_iter)
33         if (epoch + 1) % config.model_val_per_epoch == 0:
34             acc = evaluate(config, val_iter, model, data_loader)
35             logging.info(f"Accuracy on val {acc:.3f}")
36             config.writer.add_scalar('Testing/Acc', acc, global_steps)

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

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

1 if __name__ == '__main__':
2     config = ModelConfig()
3     train(config)

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

 1 -INFO: Epoch: [1/10], Batch[620/1739], Train Loss: 0.115, Train acc: 0.963
 2 -INFO: Epoch: [1/10], Batch[240/1739], Train Loss: 0.098, Train acc: 0.964
 3 -INFO: Epoch: [1/10], Batch[660/1739], Train Loss: 0.087, Train acc: 0.964
 4 ...... 
 5 - INFO:句子在澳大利亚等西方国家改变反倾销政策中对中国的划分后不少欧盟人士也认识到此种划分已背离中国经济迅速发展的现实
 6 - INFO: 澳大利亚 LOC
 7 - INFO: 中国 LOC
 8 - INFO: 欧盟 LOC
 9 - INFO: 中国 LOC
10 ......
11               precision    recall  f1-score   support
12 
13            O       1.00      0.99      1.00     97640
14        B-ORG       0.86      0.93      0.89       984
15        B-LOC       0.94      0.93      0.94      1934
16        B-PER       0.97      0.97      0.97       884
17        I-ORG       0.90      0.96      0.93      3945
18        I-LOC       0.91      0.95      0.93      2556
19        I-PER       0.99      0.98      0.98      1714
20 
21     accuracy                           0.99    109657
22    macro avg       0.94      0.96      0.95    109657
23 weighted avg       0.99      0.99      0.99    109657

10.3.3 模型推理#

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

 1 def inference(config, sentences=None):
 2     model = BertForTokenClassification(config,config.pretrained_model_dir)
 3     model_save_path = os.path.join(config.model_save_dir,config.model_save_name)
 4     if os.path.exists(model_save_path):
 5         checkpoint = torch.load(model_save_path)
 6         loaded_paras = checkpoint['model_state_dict']
 7         model.load_state_dict(loaded_paras)
 8         logging.info("## 成功载入已有模型,进行追加训练......")
 9     else:
10         raise ValueError(f" 本地模型{model_save_path}不存在,请先训练模型。")
11     data_loader = LoadChineseNERDataset(entities=config.entities,......)
12     _, token_ids, _ = data_loader.make_inference_samples(sentences)
13     token_ids = token_ids.to(config.device)
14     padding_mask = (token_ids == data_loader.PAD_IDX).transpose(0, 1)
15     logits = model(input_ids=token_ids,  # [src_len, batch_size]
16                    attention_mask=padding_mask)  # [batch_size,src_len]
17     show_result(sentences, logits, token_ids, config.entities)

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

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

1 if __name__ == '__main__':
2     config = ModelConfig()
3     sens=['智光拿出石壁拓文为乔峰详述事情始末,乔峰方知自己原本姓萧,乃契丹族。',
4           '当乔峰问及带头大哥时,却发现智光大师已圆寂。',
5           '乔峰、阿朱相约找最后知情人康敏问完此事后,就到塞外骑马牧羊,再不回来。']
6     inference(config, sens)

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

1 - INFO: 句子智光拿出石壁拓文为乔峰详述事情始末乔峰方知自己原本姓萧乃契丹族
2 - INFO:     智光 PER
3 - INFO:     乔峰 PER
4 - INFO:     乔峰 PER
5 - INFO:     萧  PER
6 - INFO:     丹  PER
7 ......

到此,对于基于BERT预训练模型的命名体识别任务就介绍完了。在下一节内容中,我们将会介绍如何从零开始训练一个BERT模型,并同时通过Tensoboard来可视化网络在训练过程中各个参数的变化情况。

您当前阅读的系列内容有完整版高清 PDF ,点击右侧了解

近200页高清 PDF + 80页配套PPT、高清无水印网络结构图,打印学习更方便!

查看详情
阅读 --

第6节 基于BERT预训练模型的SWAG选择任务

本文是 BERT 精读系列第6讲,详细讲解如何基于 huggingface bert-base-uncased 预训练权重在 SWAG 常识推理四选一任务上微调。覆盖问答选择与文本蕴含的核心差异:将问题与每个候选答案分别拼接成 4 个序列, …

第7节 自定义学习率动态调整

本文是 BERT 精读系列第7讲,作为 BERT 第8讲 SQuAD 问答任务的过渡篇,系统讲清 Transformer 论文中提出的 Noam 动态学习率调整策略:从 warmup 公式 d_model^{-0.5} · …

第8节 基于BERT预训练模型的SQuAD问答任务

本文是 BERT 精读系列第8讲,详细讲解如何基于 huggingface bert-base-uncased 在 SQuAD 斯坦福问答数据集上微调阅读理解模型。覆盖问题回答与文本蕴含的核心差异:把 (问题, 描述) 拼接后让 BERT …