这是一篇完整的 BiLSTM + CRF 中文命名实体识别（NER）实战教程，把从数据处理、模型构建到训练预测的整个闭环放在一篇文章里。数据准备篇负责把原始语料整理成模型能吃的 id 序列，模型构建篇实现 BiLSTM 编码层和 CRF 解码层，训练预测篇完成训练、评估与实体抽取。三部分代码全部基于 PyTorch，按顺序读完即可复现整套流程。
本教程使用 MSRA 中文实体语料，数据集地址 https://www.aliyundrive.com/s/HkNk5zog6gi ，提取码 60oq，训练产出的模型文件链接见训练篇内。读之前只要知道神经网络的输入是整数 id、会手写一版神经网络前向即可。
1. 数据处理
import pandas as pd import torch import pickle from datasets import Dataset from datasets import DatasetDict
1.1 读取并合并
原始语料是按字对齐的两个文件，我们先把它们对齐拼成一张表。MSRA 语料里，sentences.txt 每行一句、字之间用空格隔开，tags.txt 对应的每一行就是这一句的标签序列。两个文件按行 zip 起来，一行字配一行标签，就是一条对齐样本。
合并时要顺手做两件过滤。句子长度超过 505 的直接丢掉，太长的句子会撑爆显存，也拖慢训练。还有一个兜底，某行字数和标签数对不上就跳过，防止错位把整批数据带歪。处理完分别存成 01-训练集.csv 和 02-测试集.csv 两个文件。
def load_corpus():
train_path = ['msra/train/sentences.txt', 'msra/train/tags.txt']
valid_path = ['msra/valid/sentences.txt', 'msra/valid/tags.txt']
data_path = [train_path, valid_path]
# 1. 读取训练集数据
data_inputs, data_labels = [], []
for x_path, y_path in data_path:
for inputs, labels in zip(open(x_path), open(y_path)):
inputs = inputs.split()
labels = labels.split()
if len(inputs) > 505:
continue
if len(inputs) != len(labels):
continue
data_labels.append(' '.join(labels))
data_inputs.append(' '.join(inputs))
# 存储训练集数据
train_data = pd.DataFrame()
train_data['data_inputs'] = data_inputs
train_data['data_labels'] = data_labels
train_data.to_csv('data/01-训练集.csv')
print('训练集数据量:', len(train_data))
# 2. 读取测试集数据
test_input_path = 'msra/test/sentences.txt'
test_label_path = 'msra/test/tags.txt'
data_inputs, data_labels = [], []
for inputs, labels in zip(open(test_input_path), open(test_label_path)):
inputs = inputs.split()
labels = labels.split()
if len(inputs) > 505:
continue
if len(inputs) != len(labels):
continue
data_labels.append(' '.join(labels))
data_inputs.append(' '.join(inputs))
# 存储测试集数据
test_data = pd.DataFrame()
test_data['data_inputs'] = data_inputs
test_data['data_labels'] = data_labels
test_data.to_csv('data/02-测试集.csv')
print('测试集数据量:', len(test_data))
1.2 构建字表
语料整理成 csv 之后，下一步要给每个字编一个号，这就是字表。字表只从 01-训练集.csv 里统计，测试集一个字都不提前看。我们把训练集里所有字收集起来去重，排成一张表。表最前面手动留两个特殊符号，[PAD] 用来补空位，[UNK] 用来收纳没见过的字。
为什么字表只认训练集？因为测试集必须模拟真实上线时的陌生字。训练集里没出现过的字，统一落到 [UNK]，不能图省事把测试集的字也统计进字表，那等于提前偷看答案，评估出来的效果是假的。字表文件每行一个字，开头两行就是 [PAD] 和 [UNK]。
def build_vocab():
data_inputs = pd.read_csv('data/01-训练集.csv', usecols=['data_inputs',]).values
words = []
for data_input in data_inputs:
data_input = data_input[0].split()
words.extend(data_input)
unique_words = list(set(words))
unique_words.insert(0, '[UNK]')
unique_words.insert(0, '[PAD]')
# 将字写入到 data/vocab.txt 词典文件中
with open('data/bilstm_crf_vocab.txt', 'w') as file:
for word in unique_words:
file.write(word + '\n')
1.3 标签数字化
字要变 id，标签同样要变 id，而且标签到 id 的映射是写死的。七类标签按固定顺序编号，O 是 0，B-ORG 是 1，I-ORG 是 2，B-PER 是 3，I-PER 是 4，B-LOC 是 5，I-LOC 是 6。这个顺序一旦定下来，训练和预测全程不能改。
我们用 datasets 库把两个 csv 包成 DatasetDict，再 map 一个函数，把每个标签字符串查成对应的数字 id。字那边则按字表查表，没见过的字落成 [UNK]。整包处理完存到磁盘，后面训练直接从这里读。
def encode_label():
labels = ['O', 'B-ORG', 'I-ORG', 'B-PER', 'I-PER', 'B-LOC', 'I-LOC']
label_to_index = {label: index for index, label in enumerate(labels)}
# 将 csv 数据转换成 Dataset 类型
train_data = pd.read_csv('data/01-训练集.csv')
valid_data = pd.read_csv('data/02-测试集.csv')
train_data = Dataset.from_pandas(train_data)
valid_data = Dataset.from_pandas(valid_data)
corpus_data = DatasetDict({'train': train_data, 'valid': valid_data})
# 将标签数据转换为索引表示
def data_handler(data_labels, data_inputs):
data_label_ids = []
for labels in data_labels:
label_ids = []
for label in labels.split():
label_ids.append(label_to_index[label])
data_label_ids.append(label_ids)
return {'data_labels': data_label_ids, 'data_inputs': data_inputs}
corpus_data = corpus_data.map(data_handler, input_columns=['data_labels', 'data_inputs'], batched=True)
# 数据存储
corpus_data.save_to_disk('data/bilstm_crf_data')
if __name__ == '__main__':
load_corpus()
build_vocab()
encode_label()
整个程序跑完，会打印出两份数据量，实际输出是这样：
训练集数据量: 44968 测试集数据量: 3438
1.4 一条样本流程
我们把前面三步串起来，看一个字到底怎么变成 id 的。拿”王晓明读书”这句走一遍。假设字表里 [PAD] 是 0，[UNK] 是 1，王、晓、明、读、书分别是 2、3、4、5、6。这句的标签是 B-PER、I-PER、I-PER、O、O，按上一节的映射查成数字。字 id 和标签 id 并排看就是这样：
| 字 | 王 | 晓 | 明 | 读 | 书 |
|---|---|---|---|---|---|
| 字 id | 2 | 3 | 4 | 5 | 6 |
| 标签 | B-PER | I-PER | I-PER | O | O |
| 标签 id | 3 | 4 | 4 | 0 | 0 |
于是这一条样本最终变成一对等长的 id 序列，输入是 [2, 3, 4, 5, 6]，标签是 [3, 4, 4, 0, 0]。字 id 喂给模型，标签 id 用来算损失。如果测试集冒出一个训练集没见过的字，比如”曦”，字表里查不到，就落成 [UNK] 也就是 1，不报错。同一个 batch 里句子长短不一，短的用 [PAD] 也就是 0 补到一样长，这些补出来的位置在算损失时会被屏蔽掉。
1.5 几个注意点
实际跑的时候有几个坑要提前知道。第一，字表只认训练集，这是铁律，别图省事把测试集也统计进去。第二，超长句子 505 这个阈值是和显存、模型长度一起定的，换了硬件配置要重新调。第三，每一行字数和标签数必须对得上，错位一行，后面所有 id 全跟着错。第四，没见过的字统一走 [UNK] 兜底，别让它在查表时崩出来。
2. 模型构建
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.nn.utils.rnn import pad_sequence
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence
from datasets import load_from_disk
from transformers import BertTokenizer
# 定义计算设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
2.1 CRF 的作用
先看不用 CRF 会出什么问题，它的价值就清楚了。BiLSTM 对每个字输出一个 7 维向量，表示这个字属于 7 个标签各自的分数，这组分数通常叫发射矩阵。如果每个字独立选最大，就是逐字 argmax。问题在于标签之间有合法顺序，B-PER 后面可以接 I-PER，不能接 I-LOC，一个 I-ORG 也不能孤零零出现在句首。逐字选优不管这些，拼出来的序列可能自相矛盾。
CRF 的做法是额外学一张转移矩阵，记的是从标签 i 跳到标签 j 有多合理。解码时把发射分数和相邻标签的转移分数加在一起，在所有可能的标签序列里挑总分最高的一整条，而不是每个字各挑各的。这样选出来的序列，相邻标签衔接自然，顺序就合法了。
2.2 损失计算方式
CRF 训练时不直接说哪条序列对，而是用打分把所有序列放一起比。这里有两条核心量。一条是真实路径分数，就是真实标签序列上每个字的发射分，加上相邻标签之间的转移分，两者相加。另一条是所有路径的总分，句子长度为 n、标签数为 7，合法序列有指数级多条，不能真一条条枚举，而是用前向算法配合 log_sum_exp 递推，把所有路径的分数汇总成一个归一化因子。
损失就是总分减真实路径分。真实路径分越接近总分，损失越小，反向传播就朝这个方向调参数。举个递推的小例子。某一步要从两个前标签扩展到当前标签，两个分支的分数是 1.0 和 2.0，log_sum_exp 不是粗暴取最大的 2.0，而是 log(e 的 1.0 次方加 e 的 2.0 次方)，算出来约 2.31。比最大多出来的那一小截，正是把另一条路径也按权重考虑了进去。一层层这样累加，就把所有路径都算进了分母。
class CRF(nn.Module):
def __init__(self, label_num):
super(CRF, self).__init__()
# 转移矩阵的标签数量
self.label_num = label_num
# [TAG1, TAG2, TAG3...STAR, END]
params = torch.randn(self.label_num + 2, self.label_num + 2)
self.transition_scores = nn.Parameter(params)
# 开始和结束标签
START_TAG, ENG_TAG = self.label_num, self.label_num + 1
self.transition_scores.data[:, START_TAG] = -1000
self.transition_scores.data[ENG_TAG, :] = -1000
# 定义一个较小值用于扩展发射和转移矩阵时填充
self.fill_value = -1000.0
def _log_sum_exp(self, score):
# 计算 e 的指数时，每个元素都减去最大值，避免数值溢出
max_score, _ = torch.max(score, dim=0)
max_score_expand = max_score.expand(score.shape)
return max_score + torch.log(torch.sum(torch.exp(score - max_score_expand), dim=0))
def _get_real_path_score(self, emission_score, sequence_label):
# 计算标签的数量
seq_length = len(sequence_label)
# 计算真实路径发射分数
real_emission_score = torch.sum(emission_score[list(range(seq_length)), sequence_label])
# 在真实标签序列前后增加一个 start 和 end
b_id = torch.tensor([self.label_num], dtype=torch.int32, device=device)
e_id = torch.tensor([self.label_num + 1], dtype=torch.int32, device=device)
sequence_label_expand = torch.cat([b_id, sequence_label, e_id])
# 计算真实路径转移分数
pre_tag = sequence_label_expand[list(range(seq_length + 1))]
now_tag = sequence_label_expand[list(range(1, seq_length + 2))]
real_transition_score = torch.sum(self.transition_scores[pre_tag, now_tag])
# 计算真实路径分数
real_path_score = real_emission_score + real_transition_score
return real_path_score
def _expand_emission_matrix(self, emission_score):
# 计算标签的数量
sequence_length = emission_score.shape[0]
# 扩展时会增加 START 和 END 标签，定义该标签的值
b_s = torch.tensor([[self.fill_value] * self.label_num + [0, self.fill_value]], device=device)
e_s = torch.tensor([[self.fill_value] * self.label_num + [self.fill_value, 0]], device=device)
# 扩展发射矩阵为 (self.label_num + 2, self.label_num + 2)
expand_matrix = self.fill_value * torch.ones([sequence_length, 2], dtype=torch.float32, device=device)
emission_score_expand = torch.cat([emission_score, expand_matrix], dim=1)
emission_score_expand = torch.cat([b_s, emission_score_expand, e_s], dim=0)
return emission_score_expand
def _get_total_path_score(self, emission_score):
# 扩展发射分数矩阵
emission_score_expand = self._expand_emission_matrix(emission_score)
# 计算所有路径分数
pre = emission_score_expand[0]
for obs in emission_score_expand[1:]:
# 扩展 pre 维度
pre_expand = pre.reshape(-1, 1).expand([self.label_num + 2, self.label_num + 2])
# 扩展 obs 维度
obs_expand = obs.expand([self.label_num + 2, self.label_num + 2])
# 扩展之后 obs pre 和 self.transition_scores 维度相同
score = obs_expand + pre_expand + self.transition_scores
# 计算对数分数
pre = self._log_sum_exp(score)
return self._log_sum_exp(pre)
def forward(self, emission_scores, sequence_labels):
total_loss = 0.0
for emission_score, sequence_label in zip(emission_scores, sequence_labels):
# 计算真实路径得分
real_path_score = self._get_real_path_score(emission_score, sequence_label)
# 计算所有路径分数
total_path_score = self._get_total_path_score(emission_score)
# 最终损失
finish_loss = total_path_score - real_path_score
# 累加不同句子的损失
total_loss += finish_loss
return total_loss
def predict(self, emission_score):
"""使用维特比算法，结合发射矩阵+转移矩阵计算最优路径"""
# 扩展发射分数矩阵
emission_score_expand = self._expand_emission_matrix(emission_score)
# 计算分数
ids = torch.zeros(1, self.label_num + 2, dtype=torch.long, device=device)
val = torch.zeros(1, self.label_num + 2, device=device)
pre = emission_score_expand[0]
for obs in emission_score_expand[1:]:
# 扩展 pre 维度
pre_expand = pre.reshape(-1, 1).expand([self.label_num + 2, self.label_num + 2])
# 扩展 obs 维度
obs_expand = obs.expand([self.label_num + 2, self.label_num + 2])
# 扩展之后 obs pre 和 self.transition_scores 维度相同
score = obs_expand + pre_expand + self.transition_scores
# 获得当前多分支中最大值的分支索引
value, index = score.max(dim=0)
# 拼接每一个时间步的结果
ids = torch.cat([ids, index.unsqueeze(0)], dim=0)
val = torch.cat([val, value.unsqueeze(0)], dim=0)
# 计算分数
pre = value
# 先取出最后一个的最大值
index = torch.argmax(val[-1])
best_path = [index]
# 再回溯前一个最大值
# 由于为了方便拼接，我们在第一个位置默认填充了0
for i in reversed(ids[1:]):
# 获得分数最大的索引
# index = torch.argmax(v)
# 获得索引对应的标签ID
index = i[index].item()
best_path.append(index)
best_path = best_path[::-1][1:-1]
return best_path
我们点一下这段代码里的关键数字。标签数是 7，转移矩阵的尺寸就是 (7 加 2) 乘 (7 加 2)，也就是 9 乘 9，多出的两行两列专门留给开始和结束标签。代码里把指向开始标签的那一列全设成负的 1000，把从结束标签出发的那一行也全设成负的 1000，等于硬规定序列必须从开始标签出发、到结束标签收尾，中途不能乱穿。fill_value 取负的 1000 是同一道理，把非法衔接压到极小。
2.3 维特比找路径
训练时有真实标签当对照，预测时没有，就得自己把最优那条路径找出来，这就是维特比解码。它和前面算总分用的是同一套递推，区别只在每一步不取 log_sum_exp，而是直接取最大值，并且记下这个最大值是从哪个前标签跳过来的。递推到最后一个字，选总分最大的位置当终点，再顺着记下的索引一层层往回跳，就还原出整条最优标签序列。
对应代码里的 predict 方法。ids 这个矩阵存的是每一步的来源索引，递推结束后从最后一个位置的最大值出发，倒着查 ids 把前因后果串出来，最后切掉开头为拼接方便补的占位，以及开始和结束标签，剩下的就是每个字的标签 id。
2.4 双向编码层
CRF 负责顺序，那每个字的发射分数从哪来？靠 BiLSTM。它由三件套拼成：Embedding 把字 id 变成 256 维向量，双向 LSTM 从左到右和从右到左各看一遍、把两边表示拼成 1024 维，全连接层再把它压成 7 维，这就是每个字的发射分数。双向的好处是每个字的表示同时带了前后文，判断一个字是不是实体开头时，左右两边都能参考。
批量训练时句子长短不一，不能直接喂给 LSTM。代码里先用 pack_padded_sequence 把补出来的填充位压缩掉再算，算完用 pad_packed_sequence 还原成补零的形状，免得填充位干扰语义提取。维度上要对得上，embedding 是 256 维，单向 LSTM 隐藏层 512 维，双向拼起来正好 1024 维，全连接层就按 1024 进、7 出来接。
class BiLSTM(nn.Module):
def __init__(self, vocab_size, label_num):
super(BiLSTM, self).__init__()
# 用于将输入转换为词向量
self.embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=256)
# 用于提取输入的双向语义表示向量
self.blstm = nn.LSTM(input_size=256,
hidden_size=512,
bidirectional=True,
num_layers=1)
# 用于将 self.blstm 的输出向量映射为标签 logits
self.liner = nn.Linear(in_features=1024, out_features=label_num)
def forward(self, inputs, length):
# 将输入的 token 索引转换为词向量
outputs_embed = self.embed(inputs)
# 由于填充了很多0，此处将0进行压缩
outputs_packd = pack_padded_sequence(outputs_embed, length)
# BiLSTM 用于提取双向语义, 提取每个句子中的 token 表示
outputs_blstm, (hn, cn) = self.blstm(outputs_packd)
# outputs_paded 表示填充后的 BiLSTM 对每个 token 的输出
# outputs_length 表示每个句子实际的长度
outputs_paded, output_lengths = pad_packed_sequence(outputs_blstm)
outputs_paded = outputs_paded.transpose(0, 1)
# 线性层计算，计算出发射矩阵，形状: (16, 57, 7)
output_logits = self.liner(outputs_paded)
outputs = []
for output_logit, outputs_length in zip(output_logits, output_lengths):
outputs.append(output_logit[:outputs_length])
return outputs
def predict(self, inputs):
# 将输入的 token 索引转换为词向量
outputs_embed = self.embed(inputs)
# 增加一个 batch 维度在 1 位置
outputs_embed = outputs_embed.unsqueeze(1)
# 对每个 Token 进行语义表示
outputs_blstm, (hn, cn) = self.blstm(outputs_embed)
# 把 1 位置的 batch 值去掉
outputs_blstm = outputs_blstm.squeeze(1)
# 计算每个 Token 的发射分数
output_liner = self.liner(outputs_blstm)
return output_liner
代码注释里写的形状 (16, 57, 7)，就是一个 batch 装 16 句、每句最长 57 字、每个字输出 7 个标签分。predict 方法是给单句用的，多塞一个 batch 维进去，算完再 squeeze 掉。
2.5 拼完整模型
两层都有了，拼到一起就是一个能训能预测的 NER 模型。NER 类把 BiLSTM 和 CRF 抱在里面。forward 阶段，BiLSTM 吐出一批发射矩阵，喂给 CRF 算这批句子的总损失，拿这个损失反向传播。predict 阶段，BiLSTM 逐字打分，CRF 用维特比解出最优标签序列。save_model 把结构参数和权重一起存盘，方便下次加载接着用。
class NER(nn.Module):
def __init__(self, vocab_size, label_num):
super(NER, self).__init__()
self.vocab_size = vocab_size
self.label_num = label_num
# 双向长短记忆网络
self.bilstm = BiLSTM(vocab_size=self.vocab_size, label_num=self.label_num)
# 条件随机场网络层
self.crf = CRF(label_num=self.label_num)
def forward(self, inputs, labels, length):
# 计算输入批次样本的每个 Token 的分数，即: 每个句子的发射矩阵
emission_scores = self.bilstm(inputs, length)
# 计算批次样本的总损失
batch_loss = self.crf(emission_scores, labels)
# 返回总损失
return batch_loss
def save_model(self, save_apth):
save_info = {
'init': {'vocab_size': self.vocab_size, 'label_num': self.label_num},
'state': self.state_dict()
}
torch.save(save_info, save_apth)
def predict(self, inputs):
# 计算输入批次样本的每个 Token 的分数，即: 每个句子的发射矩阵
emission_scores = self.bilstm.predict(inputs)
# viterbi_decode 函数接收的发射矩阵为二维的 (seq_len, scores)
logits = self.crf.predict(emission_scores)
return logits
2.6 常见误区
几个容易踩的认识偏差，我们逐个澄清。以为 CRF 是分类器，它不是逐字分类，它学的是标签之间的转移规则，输出的是一整条序列的联合分数。以为负的 1000 是学出来的参数，那两行两列是写死的边界约束，真正参与学习的是其余转移分数。以为 BiLSTM 输出直接就是答案，它吐的只是发射分数，必须再过一层 CRF 才能得到合法标签序列，单独取 argmax 会拼出乱序。
3. 训练预测
3.1 训练函数
由于我们希望批次输入训练数据，在使用 RNN、GRU、LSTM 时，可以使用 pad_sequence、packed_pad_sequence、pad_packed_sequence 等函数组合来实现批次数据的输入。在将输入进行编码时，使用 BertTokenizer，为了能够准确地按字分开，输入给 BertTokenizer 的句子最好用空格隔开，便于它正确地切分出字粒度的 Token。
def pad_batch_inputs(data, labels, tokenizer):
# 函数需要返回一个按照内容长度从大到小排序过的，sentence 和 label, 还要返回 sentence 长度
# 将批次数据的输入和标签值分开，并计算批次的输入长度
data_inputs, data_length, data_labels = [], [], []
for data_input, data_label in zip(data, labels):
# 对输入句子进行编码
data_input_encode = tokenizer.encode(data_input,
return_tensors='pt',
add_special_tokens=False)
data_input_encode = data_input_encode.to(device)
data_inputs.append(data_input_encode.squeeze())
# 去除多余空格，计算句子长度
data_input = ''.join(data_input.split())
data_length.append(len(data_input))
# 将标签转换为张量
data_labels.append(torch.tensor(data_label, device=device))
# 对一个批次的内容按照长度从大到小排序, 符号表示降序
sorted_index = np.argsort(-np.asarray(data_length))
# 根据长度的索引进行排序
sorted_inputs, sorted_labels, sorted_length = [], [], []
for index in sorted_index:
sorted_inputs.append(data_inputs[index])
sorted_labels.append(data_labels[index])
sorted_length.append(data_length[index])
# 对张量进行填充，使其变成长度一样的张量
pad_inputs = pad_sequence(sorted_inputs)
return pad_inputs, sorted_labels, sorted_length
def train():
# 读取数据集
train_data = load_from_disk('data/bilstm_crf_data')['train']
# 构建分词器
tokenizer = BertTokenizer(vocab_file='data/bilstm_crf_vocab.txt')
# 构建模型
model = NER(vocab_size=tokenizer.vocab_size, label_num=7).cuda(device)
# 批次大小
batch_size = 16
# 优化器
optimizer = optim.AdamW(model.parameters(), lr=3e-5)
# 训练轮数
num_epoch = 50
# 开始训练
def start_train(data_inputs, data_labels, tokenizer):
# 对批量数据进行填充对齐
pad_inputs, sorted_labels, sorted_length = \
pad_batch_inputs(data_inputs, data_labels, tokenizer)
# 计算损失
loss = model(pad_inputs, sorted_labels, sorted_length)
# 梯度清零
optimizer.zero_grad()
# 反向传播
loss.backward()
# 参数更新
optimizer.step()
# 统计损失
nonlocal total_loss
total_loss += loss.item()
for epoch in range(num_epoch):
# 统计损失
total_loss = 0.0
# 开始训练
train_data.map(start_train,
input_columns=['data_inputs', 'data_labels'],
batched=True,
batch_size=batch_size,
fn_kwargs={'tokenizer': tokenizer},
desc='epoch: %d' % (epoch + 1))
# 打印损失
print('epoch: %d loss: %.3f' % (epoch + 1, total_loss))
# 存储模型
model.save_model('data/BiLSTM-CRF-%d.bin' % (epoch + 1))
if __name__ == '__main__':
train()
训练过程输入损失如下：
epoch: 1 loss: 646530.528 epoch: 2 loss: 322091.170 epoch: 3 loss: 231602.128 epoch: 4 loss: 182503.388 epoch: 5 loss: 149661.446 epoch: 6 loss: 125265.489 epoch: 7 loss: 106000.216 epoch: 8 loss: 90156.823 epoch: 9 loss: 76753.042 epoch: 10 loss: 65195.731 epoch: 11 loss: 55057.508 epoch: 12 loss: 46142.997 epoch: 13 loss: 38262.882 epoch: 14 loss: 31367.848 epoch: 15 loss: 25804.164 epoch: 16 loss: 23088.218 epoch: 17 loss: 18572.793 epoch: 18 loss: 15087.369 epoch: 19 loss: 13349.071 epoch: 20 loss: 11036.982 epoch: 21 loss: 9877.440 epoch: 22 loss: 8125.003 epoch: 23 loss: 6185.219 epoch: 24 loss: 5038.318 epoch: 25 loss: 4788.188 epoch: 26 loss: 3986.904 epoch: 27 loss: 3301.328 epoch: 28 loss: 2872.219 epoch: 29 loss: 2691.445 epoch: 30 loss: 3137.029 epoch: 31 loss: 1619.925 epoch: 32 loss: 2354.775 epoch: 33 loss: 2398.700 epoch: 34 loss: 1662.177 epoch: 35 loss: 1534.003 epoch: 36 loss: 1720.913 epoch: 37 loss: 1415.495 epoch: 38 loss: 1402.394 epoch: 39 loss: 1028.855 epoch: 40 loss: 1089.962 epoch: 41 loss: 1016.576 epoch: 42 loss: 1140.854 epoch: 43 loss: 1103.749 epoch: 44 loss: 735.662 epoch: 45 loss: 926.108 epoch: 46 loss: 1021.045 epoch: 47 loss: 1039.586 epoch: 48 loss: 608.818 epoch: 49 loss: 846.596 epoch: 50 loss: 685.285
最终得到的所有模型文件链接为：https://www.aliyundrive.com/s/HV1EoeYvKQ3 提取码: su42
3.2 评估函数
我们这里就对 data/BiLSTM-CRF-50.bin 模型进行评估，当然你也可以使用其他模型来评估。评估的主要步骤如下：
- 统计测试集中各个不同类别的实体名称，并将这些名称分别存储到不同列表中。
- 将测试集送入模型得到预测结果，并存储不同类别的实体名称。
- 计算不同类别的实体的精度、召回率，以及整个测试集的准确率。
def evaluate():
# 读取测试数据
valid_data = load_from_disk('data/bilstm_crf_data')['valid']
# 1. 计算各个不同类别总实体数量
# 计算测试集实体数量
total_entities = {'ORG': [], 'PER': [], 'LOC': []}
def calculate_handler(data_inputs, data_labels):
# 将 data_inputs 转换为没有空格隔开的句子
data_inputs = ''.join(data_inputs.split())
# 提取句子中的实体
extract_entities = extract_decode(data_labels, data_inputs)
# 统计每种实体的数量
nonlocal total_entities
for key, value in extract_entities.items():
total_entities[key].extend(value)
# 统计不同实体的数量
valid_data.map(calculate_handler, input_columns=['data_inputs', 'data_labels'])
print(total_entities)
# 2. 计算模型预测的各个类别实体数量
model_param = torch.load('data/BiLSTM-CRF-50.bin')
model = NER(**model_param['init']).cuda(device)
model.load_state_dict(model_param['state'])
# 构建分词器
tokenizer = BertTokenizer(vocab_file='data/bilstm_crf_vocab.txt')
model_entities = {'ORG': [], 'PER': [], 'LOC': []}
def start_evaluate(data_inputs):
# 对输入文本进行分词
model_inputs = tokenizer.encode(data_inputs, add_special_tokens=False, return_tensors='pt')[0]
model_inputs = model_inputs.to(device)
# 文本送入模型进行计算
with torch.no_grad():
label_list = model.predict(model_inputs)
# 统计预测的实体数量
text = ''.join(data_inputs.split())
# 从预测结果提取实体名字
extract_entities = extract_decode(label_list, text)
nonlocal model_entities
for key, value in extract_entities.items():
model_entities[key].extend(value)
# 统计预测不同实体的数量
valid_data.map(start_evaluate, input_columns=['data_inputs'], batched=False)
print(model_entities)
# 3. 统计每个类别的召回率
total_pred_correct = 0
total_true_correct = 0
for key in total_entities.keys():
# 获得当前 key 类别真实和模型预测实体列表
true_entities = total_entities[key]
true_entities_num = len(true_entities)
pred_entities = model_entities[key]
# 分解预测实体中，pred_correct 表示预测正确，pred_incorrect 表示预测错误
pred_correct, pred_incorrect = 0, 0
for pred_entity in pred_entities:
if pred_entity in true_entities:
pred_correct += 1
continue
pred_incorrect += 1
# 模型预测的 key 类别的实体数量
model_pred_key_num = true_entities_num + pred_incorrect
# 计算共预测正确多少个实体
total_pred_correct += pred_correct
# 计算共有多少个真实的实体
total_true_correct += true_entities_num
# 计算精度
print(key, '查全率: %.3f' % (pred_correct / true_entities_num))
print(key, '查准率: %.3f' % (pred_correct / model_pred_key_num))
print('-' * 50)
print('准确率: %.3f' % (total_pred_correct / total_true_correct))
def extract_decode(label_list, text):
"""
:param label_list: 模型输出的包含标签序列的一维列表
:param text: 模型输入的句子
:return: 提取到的实体名字
"""
labels = ['O', 'B-ORG', 'I-ORG', 'B-PER', 'I-PER', 'B-LOC', 'I-LOC']
label_to_index = {label: index for index, label in enumerate(labels)}
B_ORG, I_ORG = label_to_index['B-ORG'], label_to_index['I-ORG']
B_PER, I_PER = label_to_index['B-PER'], label_to_index['I-PER']
B_LOC, I_LOC = label_to_index['B-LOC'], label_to_index['I-LOC']
# 提取连续的标签代表的实体
def extract_word(start_index, next_label):
# index 表示最后索引的位置
index, entity = start_index + 1, [text[start_index]]
for index in range(start_index + 1, len(label_list)):
if label_list[index] != next_label:
break
entity.append(text[index])
return index, ''.join(entity)
# 存储提取的命名实体
extract_entites, index = {'ORG': [], 'PER': [], 'LOC': []}, 0
# 映射下一个持续的标签
next_label = {B_ORG: I_ORG, B_PER: I_PER, B_LOC: I_LOC}
# 映射词的所属类别
word_class = {B_ORG: 'ORG', B_PER: 'PER', B_LOC: 'LOC'}
while index < len(label_list):
# 获得当前位置的标签
label = label_list[index]
if label in next_label.keys():
# 将当前位置和对应的下一个持续标签传递到 extract_word 函数
index, word = extract_word(index, next_label[label])
extract_entites[word_class[label]].append(word)
continue
index += 1
return extract_entites
if __name__ == '__main__':
evaluate()
程序输出结果：
ORG 查全率: 0.762 ORG 查准率: 0.655 -------------------------------------------------- PER 查全率: 0.824 PER 查准率: 0.716 -------------------------------------------------- LOC 查全率: 0.846 LOC 查准率: 0.778 -------------------------------------------------- 准确率: 0.821
3.3 预测函数
预测函数就是输入一个句子，提取句子中的实体。步骤如下：
- 输入的句子先添加上空格隔开。
- 用 BertTokenizer 进行编码。
- 使用 model 进行预测。
- 使用维特比解码解析出实体名称。
def entity_extract(text):
# 构建分词器
tokenizer = BertTokenizer(vocab_file='data/bilstm_crf_vocab.txt')
# 初始化模型
model_param = torch.load('data/BiLSTM-CRF-48.bin')
model = NER(**model_param['init']).cuda(device)
model.load_state_dict(model_param['state'])
# 我们先按字将其分开，并在字之间添加空格，便于 Bert 分词器能够准确按字分割
input_text = ' '.join(list(text))
model_inputs = tokenizer.encode(input_text, add_special_tokens=False, return_tensors='pt')[0]
model_inputs = model_inputs.to(device)
with torch.no_grad():
outputs = model.predict(model_inputs)
return extract_decode(outputs, ''.join(input_text.split()))
if __name__ == '__main__':
text = '我要感谢洛杉矶市民议政论坛、亚洲协会南加中心、美中关系全国委员会、美中友协美西分会等友好团体的盛情款待。'
result = entity_extract(text)
print(result)
程序输出结果：
{'ORG': ['亚洲协会南加中心', '美中关系全国委员会', '美中友协美西分会'], 'PER': [], 'LOC': ['洛杉矶']}
这篇实战把中文 NER 从数据到部署的完整链路走了一遍。数据处理环节，把 MSRA 语料按字对齐合并成 csv，只从训练集构建含 [PAD]、[UNK] 的字表，七类标签映射为 0 到 6 的数字，一条样本最终变成一对等长的 id 序列。模型构建环节，BiLSTM 双向编码给每个字打出发射分数，CRF 用转移矩阵约束标签顺序，训练时以全路径归一化分数减真实路径分数为损失，预测时用维特比算法回溯最优标签序列。训练预测环节，把两层拼成可训练可预测的完整模型，在测试集上 ORG 查全率 0.762、查准率 0.655，PER 查全率 0.824、查准率 0.716，LOC 查全率 0.846、查准率 0.778，整体准确率 0.821。数据、模型、训练三个环节环环相扣，这就是一套可以实际跑起来的中文 NER 实现。
冀公网安备13050302001966号