使用bert的预训练模型做命名实体识别NER_用大模型做实体识别

前言

在上一篇博客提到了如何使用blstm-crf模型来训练本地数据集，完成命名实体提取的任务，还详细解析了代码和对应的原理。针对特定的任务，垂直的领域，可能需要使用特定数据集去训练，从而使得模型有一个很好的效果，但是在一些非特定（垂直）领域，是完全可以使用预训练好的模型来做命名实体识别任务的。下面，介绍一下如何使用开源框架，导入预训练好的bert模型来做命名实体识别。

基于transformers完成NER任务

废话不多说，直接上代码

import os
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from seqeval.metrics.sequence_labeling import get_entities

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained("shibing624/bert4ner-base-chinese")
model = AutoModelForTokenClassification.from_pretrained("shibing624/bert4ner-base-chinese")
label_list = ['I-ORG', 'B-LOC', 'O', 'B-ORG', 'I-LOC', 'I-PER', 'B-TIME', 'I-TIME', 'B-PER']

sentence = "王宏伟来自北京，是个警察，喜欢去王府井游玩儿。"


def get_entity(sentence):
    tokens = tokenizer.tokenize(sentence)
    inputs = tokenizer.encode(sentence, return_tensors="pt")
    with torch.no_grad():
        outputs = model(inputs).logits
    predictions = torch.argmax(outputs, dim=2)
    char_tags = [(token, label_list[prediction]) for token, prediction in zip(tokens, predictions[0].numpy())][1:-1]
    print(sentence)
    print(char_tags)

    pred_labels = [i[1] for i in char_tags]
    entities = []
    line_entities = get_entities(pred_labels)
    for i in line_entities:
        word = sentence[i[1]: i[2] + 1]
        entity_type = i[0]
        entities.append((word, entity_type))

    print("Sentence entity:")
    print(entities)


get_entity(sentence)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

这段代码就是使用"shibing624/bert4ner-base-chinese"这个预训练模型，来对"王宏伟来自北京，是个警察，喜欢去王府井游玩儿。"这句话做命名实体识别任务。这里主要提取ORG， LOC， PER和TIME这四种实体。最终模型输出的结果如下：

需要注意的是，这里的预训练模型需要从HuggingFace Hub官网在线下载，如果有的小伙伴因为网络的原因在线下载模型失败，可以参考这篇博客ReadTimeoutError: HTTPSConnectionPool(host=‘cdn-lfs.huggingface.co‘, port=443)来手动下载模型并导入。

可以看到，bert4ner-base-chinese这个预训练模型很好的完成了NER任务，并且提取结果都是正确的。

基于paddlenlp的完成NER任务

当然，除了国外的HuggingFace能够提供一些开源的模型来完成NER任务，国内也有一些比较好的开源框架，比如paddlenlp。而且，在paddlenlp2.3.4的版本之后，可以实现一行代码完成一些基础的nlp任务。

首先上代码：

from paddlenlp import Taskflow

ner = Taskflow("ner")

sentence = "王宏伟来自北京，是个警察，喜欢去王府井游玩儿。"

print(ner(sentence))
1
2
3
4
5
6
7

可以看到，从paddlenlp导入Taskflow库，可以一行代码搞定NER这个任务。paddlenlp的效果如下所示：

paddlenlp的ner的结果会更加细分一些，不仅把"王宏伟"提取成’人物类_实体’、‘王府井’提取成’场所类’，还把其他的每个词或者句子分成了paddlenlp框架定义的其他实体名称。

这里需要说明的是，要完成上述简单的任务，搭建环境的时候需要注意各个依赖库的版本问题。下面给出笔者环境版本，供各位参考：
paddlepaddle-gpu 2.3.2
paddlenlp 2.3.4
python 3.7
cudatoolkit 10.1.243
cudnn 7.6.5

之前安装paddlepaddle-gpu的版本是2.4.2，按照paddlepaddle官网的命令安装，但是报错ImportError: libcudart.so.10.2: cannot open shared object file: No such file or directory 。后来一查，有人说2.4的版本不稳定，容易报各种错误，所以就改成了2.3.2的版本，果然没有报错了。

参考

PaddleNLP一键预测功能：Taskflow API
paddlepaddle安装指南