命名实体识别（Named Entity Recognition，简称 NER）是自然语言处理（NLP）中的一项重要任务，旨在从非结构化文本中识别并分类具有特定意义的实体，如人名、地名、机构名、时间表达、数量、货币和百分比等。
1. 数据处理
我们使用的是 MSRA 中文 NER 数据集，该数据集共包含三个目录：test、valid、train，分别对应了测试集、验证集、训练集。训练集有：42000 条数据，验证集有 3000 条数据，测试集有 3442 条数据。 另外也包括了一个定义了预测标签的 tags.txt 文件，数据集使用的是 BIO 的标签体系，即：O、B-PER、I-PER、B-ORG、I-ORG、B-LOC、I-LOC，目录结构为：
msra
├── tags.txt
├── test
│ ├── sentences.txt
│ └── tags.txt
├── train
│ ├── sentences.txt
│ └── tags.txt
└── valid
├── sentences.txt
└── tags.txt
数据集链接：https://www.aliyundrive.com/s/HkNk51zog6gi 提取码: 60oq
import pandas as pd import torch import pickle import os from datasets import Dataset, DatasetDict from transformers import BertTokenizer import datasets import transformers # 显示进度条 datasets.disable_progress_bar() # 设置日志级别 datasets.logging.set_verbosity(datasets.logging.CRITICAL) transformers.logging.set_verbosity(transformers.logging.CRITICAL) # 修改工作目录 os.chdir(os.path.dirname(os.path.abspath(__file__)))
我们将 valid、train 数据集合并为 train 数据集。由于我们使用的是 Bert 模型来完成 NER 任务，而 Bert 模型对输入有 max_len 限制。所以，我们得把长度大于 505 的数据过滤掉。当然这里的 505 可以修改，只要总长度不超过 512 就可以。最终，我们会生成两个 csv 文件，01-训练集.csv 和 02-测试集.csv 分别存储训练集和测试集数据。
# 加载训练集
def load_train_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]
data_inputs, data_labels = [], []
max_len = 0
for x_path, y_path in data_path:
for data_input, data_label in zip(open(x_path), open(y_path)):
data_input = data_input.split()
data_label = data_label.split()
if len(data_input) > max_len:
max_len = len(data_input)
if len(data_input) > 505:
continue
if len(data_input) != len(data_label):
continue
data_labels.append(' '.join(data_label))
data_inputs.append(' '.join(data_input))
# 数据存储
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), '最大句子长度:', max_len)
# 加载测试集
def load_valid_corpus():
x_path = 'msra/test/sentences.txt'
y_path = 'msra/test/tags.txt'
data_inputs, data_labels = [], []
max_len = 0
for data_input, data_label in zip(open(x_path), open(y_path)):
data_input = data_input.split()
data_label = data_label.split()
if len(data_input) > max_len:
max_len = len(data_input)
if len(data_input) > 505:
continue
if len(data_input) != len(data_label):
continue
data_labels.append(' '.join(data_label))
data_inputs.append(' '.join(data_input))
# 数据存储
valid_data = pd.DataFrame()
valid_data['data_inputs'] = data_inputs
valid_data['data_labels'] = data_labels
valid_data.to_csv('data/02-测试集.csv')
print('测试集数据量:', len(valid_data), '最大句子长度:', max_len)
程序的输出结果：
训练集数据量: 44968 最大句子长度: 746 测试集数据量: 3438 最大句子长度: 2427
这里要做的事情就是将 O、B-LOC、I-PER…等标签映射为数字索引表示。转换完成之后，存储到 data/03-train，具体实现代码如下：
def data_handler(data_labels_batch, data_inputs):
labels = []
for data_labels in data_labels_batch:
label = []
for data_label in data_labels.split():
label.append(label_to_index[data_label])
labels.append(label)
return {'data_label_ids': labels, 'data_inputs': data_inputs}
if __name__ == '__main__':
load_train_corpus()
load_valid_corpus()
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)
# 由于在 centos 里 load_dataset 总是报错，所以按照以下步骤加载
# 1. 先使用 pd.read_csv 从 csv 文件加载数据
# 2. 再使用 Dataset.from_pandas 从 DataFram 加载数据
# 3. 最后使用 DatasetDict 构建包含训练集和测试集的数据对象
data = DatasetDict({'train': train_data, 'valid': valid_data})
# input_columns: 指定传递到 data_handler 中的列，以位置参数的形式传递
# fn_kwargs: 指定传递到 data_handler 中的自定义关键字参数
train_data = data.map(data_handler,
input_columns=['data_labels', 'data_inputs'],
batched=True,
batch_size=1000)
train_data.save_to_disk('data/03-train')
2. 模型训练
我们使用 Bert 的 bert-base-chinese 为基础模型，在我们自己的数据集上进行微调来实现 NER 任务。我们这里使用的是 transformers 库中提供的 BertForTokenClassification 类来进行微调，该 head 非常适合字词粒度的分类任务。
import datasets from transformers import BertModel from transformers import BertConfig from transformers import BertTokenizer from datasets import load_from_disk from transformers import BertForMaskedLM from transformers import BertForTokenClassification import torch from torch.nn.utils.rnn import pad_sequence import torch.optim as optim import math import numpy as np
由于数据集有 4 万+，训练时间较长，我们这里我们使用累计多个 step 的损失之后，再更新参数，训练效率得到明显提升。训练的轮数为 40，也可以改成其他值。存储模型时，从 11 epoch 开始，每一个 epoch 结束就存储一次。
另外，由于我们希望模型对象帮我们去自动计算损失，所以，在调用 model.forward 函数时，可以传递 labels 参数，但要注意的是，我们传递的 labels 参数必须为 tensor 类型，它要求每个数据长度必须一样。所以，可以使用 pad_sequence 对 labels 填充下，模型内部会使用 attention_mask 屏蔽填充这部分标签值的。
受限于显存大小，我这里最高只能设置 batch_size 为 4，没迭代 8 次，即: 累计 8 次梯度进行一次参数更新。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def train():
# 读取训练数据
train_data = load_from_disk('data/03-train')['train']
# 初始化分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
# 初始化模型
model = BertForTokenClassification.from_pretrained('bert-base-chinese',
num_labels=7)
model.to(device)
# 优化器
optimizer = optim.AdamW(model.parameters(), lr=3e-5)
# 训练批次
batch_size = 4
# 训练轮数
epoch_num = 40
# 累计梯度
accumulate_step = 0
total_steps = int(math.ceil(len(train_data) / batch_size))
def start_train(data_inputs, data_label_ids):
data_inputs = tokenizer.batch_encode_plus(data_inputs,
add_special_tokens=False,
padding='longest',
return_tensors='pt')
data_inputs['input_ids'] = data_inputs['input_ids'].to(device)
data_inputs['attention_mask'] = data_inputs['attention_mask'].to(device)
data_inputs['token_type_ids'] = data_inputs['token_type_ids'].to(device)
# 标签转换张量
data_labels = []
for labels in data_label_ids:
data_labels.append(torch.tensor(labels, dtype=torch.int64, device=device))
labels = pad_sequence(data_labels, batch_first=True, padding_value=-1)
# 模型计算
outputs = model(**data_inputs, labels=labels)
# 反向传播
loss = outputs.loss
loss.backward()
nonlocal accumulate_step
accumulate_step += 1
if accumulate_step % 8 == 0 or accumulate_step == total_steps:
# 参数更新
optimizer.step()
# 梯度清零
optimizer.zero_grad()
nonlocal epoch_loss
epoch_loss += loss.item()
for epoch_idx in range(epoch_num):
epoch_loss = 0.0
train_data.map(start_train,
batched=True,
batch_size=batch_size,
input_columns=['data_inputs', 'data_label_ids'])
print('loss: %.5f' % epoch_loss)
if epoch_idx > 10:
model.save_pretrained('data/ner-model-%d' % epoch_idx)
训练结束后得到的模型文件如下：
ner-model-12 ner-model-17 ner-model-22 ner-model-27 ner-model-32 ner-model-37 ner-model-13 ner-model-18 ner-model-23 ner-model-28 ner-model-33 ner-model-38 ner-model-14 ner-model-19 ner-model-24 ner-model-29 ner-model-34 ner-model-39 ner-model-15 ner-model-20 ner-model-25 ner-model-30 ner-model-35 ner-model-11 ner-model-16 ner-model-21 ner-model-26 ner-model-31 ner-model-36
每个模型目录下的内容为：
ner-model-11 ├── config.json └── pytorch_model.bin
3. 模型评估
模型评估主要做的事情，提取测试集所有的实体名称，并划分为 ORG、PER、LOC 类别。分别统计每个类别的精度、召回率，以及准确率。
我编写了 extract_decode 函数，该函数接收一个句子，以及该句子中每个字对应的实体标签。从而，提取出该剧中标识的所有实体，并分类存储。
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
评估函数 evaluate 主要做三件事。首先，计算提取测试集中所有的实体名称；然后，由模型对测试集数据进行预测，从而得到模型预测出的实体标签。最后，根据模型预测的结果和真实结果计算不同类别的精度、召回率，例如：计算模型对 LOC 位置实体的预测的精度、召回率。
前面，我们训练出了多个模型，我们这里分别对每个模型都进行评估。部分评估的结果如下：
# ner-model-12 这个模型在测试集上的评估结果 ORG 查全率: 0.929 ORG 查准率: 0.844 -------------------------------------------------- PER 查全率: 0.973 PER 查准率: 0.936 -------------------------------------------------- LOC 查全率: 0.933 LOC 查准率: 0.902 -------------------------------------------------- 准确率: 0.942 # ner-model-37 这个模型在测试集上的评估结果 ORG 查全率: 0.924 ORG 查准率: 0.814 -------------------------------------------------- PER 查全率: 0.961 PER 查准率: 0.926 -------------------------------------------------- LOC 查全率: 0.924 LOC 查准率: 0.894 -------------------------------------------------- 准确率: 0.933 ...等等
从评估结果可以看到，不同的模型在 ORG、PER、LOC 上的精度、召回率都是不同的。这也可以简单理解，不同的模型在对具体的类别的命名实体识别时，性能是不同的。下面是完整的评估代码：
def evaluate():
# 读取测试数据
valid_data = load_from_disk('data/03-train')['valid_data']
# 1. 计算各个不同类别总实体数量
# 计算测试集实体数量
total_entities = {'ORG': [], 'PER': [], 'LOC': []}
def calculate_handler(data_inputs, data_label_ids):
# 将 data_inputs 转换为没有空格隔开的句子
text = ''.join(data_inputs.split())
label_list = data_label_ids
# 提取句子中的实体
extract_entities = extract_decode(data_label_ids, text)
# 统计每种实体的数量
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_label_ids'])
print(total_entities)
# 2. 计算模型预测的各个类别实体数量
# 初始化分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
for index in range(11, 38):
# 初始化模型
model = BertForTokenClassification.from_pretrained('data/ner-model-%d' % index, num_labels=7)
model.train(mode=False)
model_entities = {'ORG': [], 'PER': [], 'LOC': []}
def start_evaluate(data_inputs):
# 对输入文本进行分词
model_inputs = tokenizer(data_inputs, add_special_tokens=False, return_tensors='pt')
# 文本送入模型进行计算
with torch.no_grad():
outputs = model(**model_inputs)
# 统计预测的实体数量
label_list = torch.argmax(outputs.logits.squeeze(), dim=-1).tolist()
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. 统计每个类别的召回率
print('#%d\n' % index)
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))
4. 预测函数
预测函数就比较简单了，就是传递一个句子，抽取句子中的实体。实现代码如下：
def entity_extract(text):
# 初始化分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
# 初始化模型
model = BertForTokenClassification.from_pretrained('data/ner-model-39', num_labels=7)
model.train(False)
# 我们先按字将其分开，并在字之间添加空格，便于 Bert 分词器能够准确按字分割
input_text = ' '.join(list(text))
print(text)
inputs = tokenizer.encode_plus(input_text, add_special_tokens=False, return_tensors='pt')
with torch.no_grad():
outputs = model(**inputs)
y_pred = torch.argmax(outputs.logits, dim=-1)[0].tolist()
return extract_decode(y_pred, text)
if __name__ == '__main__':
text = '今年７月１日我国政府恢复对香港行使主权，标志着“一国两制”构想的巨大成功，标志着中国人民在祖国统一大业的道路上迈出了重要的一步。'
result = entity_extract(text)
print(result)
text = '同时，三毛集团自身也快速扩张，企业新创造了３０００多个就业岗位，安置了一大批下岗职工。'
result = entity_extract(text)
print(result)
text = '我要感谢洛杉矶市民议政论坛、亚洲协会南加中心、美中关系全国委员会、美中友协美西分会等友好团体的盛情款待。'
result = entity_extract(text)
print(result)
程序运行结果如下：
今年７月１日我国政府恢复对香港行使主权，标志着“一国两制”构想的巨大成功，标志着中国人民在祖国统一大业的道路上迈出了重要的一步。
{'ORG': [], 'PER': [], 'LOC': ['香港', '中国']}
同时，三毛集团自身也快速扩张，企业新创造了３０００多个就业岗位，安置了一大批下岗职工。
{'ORG': ['三毛集团'], 'PER': [], 'LOC': []}
我要感谢洛杉矶市民议政论坛、亚洲协会南加中心、美中关系全国委员会、美中友协美西分会等友好团体的盛情款待。
{'ORG': ['亚洲协会南加中心', '美中关系全国委员会', '美中友协美西分会'], 'PER': [], 'LOC': ['洛杉矶']}
冀公网安备13050302001966号