基于transformers框架实践Bert系列2--命名实体识别

本系列用于Bert模型实践实际场景，分别包括分类器、命名实体识别、选择题、文本摘要等等。（关于Bert的结构和详细这里就不做讲解，但了解Bert的基本结构是做实践的基础，因此看本系列之前，最好了解一下transformers和Bert等）
本篇主要讲解命名实体识别应用场景。本系列代码和数据集都上传到GitHub上：https://github.com/forever1986/bert_task

1 环境说明

1）本次实践的框架采用torch-2.1+transformer-4.37
2）另外还采用或依赖其它一些库，如：evaluate、pandas、datasets、accelerate、seqeval等

2 前期准备

Bert模型是一个只包含transformer的encoder部分，并采用双向上下文和预测下一句训练而成的预训练模型。可以基于该模型做很多下游任务。

2.1 了解Bert的输入输出

Bert的输入：input_ids（使用tokenizer将句子向量化），attention_mask，token_type_ids（句子序号）、labels（结果）
Bert的输出：
last_hidden_state：最后一层encoder的输出；大小是(batch_size, sequence_length, hidden_size)（注意：这是关键输出，本次任务就需要获取该值，并进行一次线性层处理）
pooler_output：这是序列的第一个token(classification token)的最后一层的隐藏状态，输出的大小是(batch_size, hidden_size)，它是由线性层和Tanh激活函数进一步处理的。（通常用于句子分类，至于是使用这个表示，还是使用整个输入序列的隐藏状态序列的平均化或池化，视情况而定）。
hidden_states： 这是输出的一个可选项，如果输出，需要指定config.output_hidden_states=True,它也是一个元组，它的第一个元素是embedding，其余元素是各层的输出，每个元素的形状是(batch_size, sequence_length, hidden_size)
attentions：这是输出的一个可选项，如果输出，需要指定config.output_attentions=True,它也是一个元组，它的元素是每一层的注意力权重，用于计算self-attention heads的加权平均值。

2.2 数据集与模型

1）数据集：weibo_senti_100k（微博公开数据集），这里只是演示，使用其中2400条数据
2）模型：bert-base-chinese
注意：本次练习都是采用本地数据集和本地权重模型，不直接从hf下载，因为速度过慢

2.3 任务说明

1）什么是命名实体识别
命名实体识别其实就是对输出结果的每个token做出标签（标注属于哪个实体），而实体类型包括地缘政治实体(GPE.NAM)、地名(LOC.NAM)、机构名(ORG.NAM)、人名(PER.NAM)及其对应的代指(以NOM为结尾）
2）NLP的序列标注
NLP的序列标注有很多种方式：IOB、IBO2、BIOES、IOE1等等，有兴趣可以自行了解。这里简单说一下IBO2，这个案例使用到的。就是通过3类标签（B-begin（实体开始），I-inside（实体中间），O-outside（实体之外）），下表就是本次weibo_ner数据集的基本标签

标签	含义	标签	含义
B-PER.NAM	名字（张三）开始标签	I-PER.NAM	名字（张三）中间标签
B-PER.NOM	代称、类别名（穷人）开始标签	I-PER.NOM	代称、类别名（穷人）中间标签
B-LOC.NAM	特指名称（紫玉山庄）开始标签	I-LOC.NAM	特指名称（紫玉山庄）中间标签
B-LOC.NOM	泛称（大峡谷、宾馆）开始标签	I-LOC.NOM	泛称（大峡谷、宾馆）中间标签
B-GPE.NAM	行政区的名称（北京）开始标签	I-GPE.NAM	行政区的名称（北京）中间标签
B-GPE.NOM	泛指行政区的名称（国家）开始标签	I-GPE.NOM	泛指行政区的名称（国家）中间标签
B-ORG.NAM	特定机构名称（通惠医院）开始标签	I-ORG.NAM	特定机构名称（通惠医院）中间标签
B-ORG.NOM	泛指名称、统称（文艺公司）开始标签	I-ORG.NOM	泛指名称、统称（文艺公司）中间标签
O	非实体标签

3）命名实体识别任务就是给每个token都标注最终属于哪个标签含义

2.4 实现关键

其实就是将bert输出的每个token预测为对应到上表中每个label的值，所以是每个token，因此需要使用last_hidden_state输出结果，将输出结果在接入到线性层（输入是hidden_size, 输出是num_labels的线性层），最终将输出做argmax可以得出每个token的标签。loss计算使用交叉熵。

3 关键代码

3.1 数据集处理

weibo_ner原始数据是如下
人 O
生 O
如 O
戏 O
， O
导 B-PER.NOM
演 I-PER.NOM
是 O
自 O
己 O
蜡 O
烛 O
经过我们处理，将其保存为如下：

{'id': '0', 'tokens': ['科', '技', '全', '方', '位', '资', '讯', '智', '能', '，', '快', '捷', '的', '汽', '车', '生', '活', '需', '要', '有', '三', '屏', '一', '云', '爱', '你'], 'ner_tags': [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]}

其中tokens是每个中文字，ner_tags为标签label。这是为了方便tokenizer进行转换。数据tokenizer后转换的函数如下：

# 关键点2处
# 1.需要增加is_split_into_words=True，因为我们的数据tokens是每个字分开的，需要将其连成一个list，比如tokenizer之后都是分开的数组
# 2.labels和input_ids匹配，因为每个中文字tokenizer之后可能对应成多个token，我们需要知道每个token对于哪个字，这时候就需要使用word_ids属性来判断，这样就能对齐label_ids与input_ids
# 3.使用word_ids就必须使用BertTokenizerFast，原先BertTokenizer返回值是没有word_ids的
def process_function(datas):
    tokenized_datas = tokenizer(datas["tokens"], max_length=256, truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(datas["ner_tags"]):
        word_ids = tokenized_datas.word_ids(batch_index=i)
        label_ids = []
        for word_id in word_ids:
            if word_id is None:
                label_ids.append(-100)
            else:
                label_ids.append(label[word_id])
        labels.append(label_ids)
    tokenized_datas["labels"] = labels
    return tokenized_datas

3.2 模型加载

model = BertForTokenClassification.from_pretrained(model_path, num_labels=len(label_list))

注意：这里使用的是BertForTokenClassification，因为我们在2.4实现关键中说了，我们可以看看BertForTokenClassification实现关键代码是否满足我们的要求：

# 在__init__方法中增加dropout和分类线性层
self.bert = BertModel(config, add_pooling_layer=False)
classifier_dropout = (
    config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
)
self.dropout = nn.Dropout(classifier_dropout)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)

# 在forward方法中，将bert输出outputs中的第一个值（也就是前面讲到的last_hidden_state），进行dropout处理并关联线性层
outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

3.3 评估函数

本次评估函数使用f1，主要考虑到命名实体识别其实就是判断其边界是否正确;实体的类型是否标注正确。对于f1能综合考虑它们的加权调和平均值。使用seqeval库来计算f1。

# 评估函数：此处的评估函数可以从https://github.com/huggingface/evaluate下载到本地
seqeval = evaluate.load("./evaluate/seqeval_metric.py")
def evaluate_function(preprediction):
    predictions, labels = preprediction
    predictions = numpy.argmax(predictions, axis=-1)

    # 将id转换为原始的字符串类型的标签
    true_predictions = [
        [label_list[p] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    true_labels = [
        [label_list[l] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    result = seqeval.compute(predictions=true_predictions, references=true_labels, mode="strict", scheme="IOB2")
    return {
        "f1": result["overall_f1"]
    }

3.4 设置TrainingArguments

这里相对于情感分类的应用场景中没有设置学习率和权重衰减，好像效果更好，大家也可以自己尝试一下，可能是数据集太少的原因，训练准确率不高。

# step 5 创建TrainingArguments
# train是1350条数据，batch_size=32，因此每个epoch的step=43，总step=129，
train_args = TrainingArguments(output_dir="./checkpoints",      # 输出文件夹
                               per_device_train_batch_size=32,  # 训练时的batch_size
                               per_device_eval_batch_size=32,    # 验证时的batch_size
                               num_train_epochs=3,              # 训练轮数
                               logging_steps=20,                # log 打印的频率
                               evaluation_strategy="epoch",     # 评估策略
                               save_strategy="epoch",           # 保存策略
                               save_total_limit=3,              # 最大保存数
                               metric_for_best_model="f1",      # 设定评估指标
                               load_best_model_at_end=True      # 训练完成后加载最优模型
                               )

3.5 设置DataCollatorForTokenClassification

注意本次实验没有使用DataCollatorWithPadding，而是使用DataCollatorForTokenClassification。因为ner任务需要对标签label进行补齐，而DataCollatorWithPadding不对标签label补齐。而在情感分类中，我们的label其实只取第一个token（pooler_output）都是一个值0或1，所以不需要补齐。

4 整体代码

"""
基于BERT做命名实体识别
1）数据集来自：weibo_ner
2）模型权重使用：bert-base-chinese
"""
# step 1 引入数据库
import numpy
import evaluate
from datasets import DatasetDict
from transformers import BertForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification, \
     pipeline, BertTokenizerFast


model_path = "./model/tiansz/bert-base-chinese"
data_path = "data/weibo_ner"


# step 2 数据集处理
datasets = DatasetDict.load_from_disk(data_path)
# 保存label真实描述，用于显示正确结果和传入模型初始化告诉模型分类数量
label_list = ['B-GPE.NAM', 'B-GPE.NOM', 'B-LOC.NAM', 'B-LOC.NOM', 'B-ORG.NAM', 'B-ORG.NOM', 'B-PER.NAM', 'B-PER.NOM',
              'I-GPE.NAM', 'I-GPE.NOM', 'I-LOC.NAM', 'I-LOC.NOM', 'I-ORG.NAM', 'I-ORG.NOM', 'I-PER.NAM', 'I-PER.NOM',
              'O']
tokenizer = BertTokenizerFast.from_pretrained(model_path)


# 借助word_ids 实现标签映射
def process_function(datas):
    tokenized_datas = tokenizer(datas["tokens"], max_length=256, truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(datas["ner_tags"]):
        word_ids = tokenized_datas.word_ids(batch_index=i)
        label_ids = []
        for word_id in word_ids:
            if word_id is None:
                label_ids.append(-100)
            else:
                label_ids.append(label[word_id])
        labels.append(label_ids)
    tokenized_datas["labels"] = labels
    return tokenized_datas


new_datasets = datasets.map(process_function, batched=True)

# step 3 加载模型
model = BertForTokenClassification.from_pretrained(model_path, num_labels=len(label_list))

# step 4 评估函数：此处的评估函数可以从https://github.com/huggingface/evaluate下载到本地
seqeval = evaluate.load("./evaluate/seqeval_metric.py")


def evaluate_function(prepredictions):
    predictions, labels = prepredictions
    predictions = numpy.argmax(predictions, axis=-1)
    # 将id转换为原始的字符串类型的标签
    true_predictions = [
        [label_list[p] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    true_labels = [
        [label_list[l] for p, l in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    result = seqeval.compute(predictions=true_predictions, references=true_labels, mode="strict", scheme="IOB2")
    return {
        "f1": result["overall_f1"]
    }


# step 5 创建TrainingArguments
# train是1350条数据，batch_size=32，因此每个epoch的step=43，总step=129
train_args = TrainingArguments(output_dir="./checkpoints",      # 输出文件夹
                               per_device_train_batch_size=32,  # 训练时的batch_size
                               # gradient_checkpointing=True,     # *** 梯度检查点 ***
                               per_device_eval_batch_size=32,    # 验证时的batch_size
                               num_train_epochs=3,              # 训练轮数
                               logging_steps=20,                # log 打印的频率
                               evaluation_strategy="epoch",     # 评估策略
                               save_strategy="epoch",           # 保存策略
                               save_total_limit=3,              # 最大保存数
                               metric_for_best_model="f1",      # 设定评估指标
                               load_best_model_at_end=True      # 训练完成后加载最优模型
                               )

# step 6 创建Trainer
trainer = Trainer(model=model,
                  args=train_args,
                  train_dataset=new_datasets["train"],
                  eval_dataset=new_datasets["validation"],
                  data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer),
                  compute_metrics=evaluate_function,
                  )

# step 7 训练
trainer.train()

# step 8 模型评估
evaluate_result = trainer.evaluate(new_datasets["test"])
print(evaluate_result)

# step 9：模型预测
ner_pipe = pipeline("token-classification", model=model, tokenizer=tokenizer, device=0, aggregation_strategy="simple")
res = ner_pipe("对，输给一个女人，的成绩。失望")
print(res)

5 运行效果

注：本文参考来自大神：https://github.com/zyds/transformers-code