Transformers预训练模型使用：命名实体识别 Named Entity Recognition

命名实体识别的任务是对每一个token都进行分类。比如，识别这个token是不是一个人名、组织名或地名。命名实体识别的一个数据集是CoNLL-2003，这个数据集完全契合这个任务。

使用pipeline#

下面是一个使用pipeline实现命名实体识别的样例。首先，要定义9中标签分类：

O：不是命名实体。
B-MIS：其他类命名实体的开始标记。
I-MIS：其他类命名实体的中间标记。
B-PER：人名的开始标记。
I-PER：人名的中间标记。
B-ORG：组织名的开始标记。
I-ORG：组织名的中间标记。
B-LOC：地名的开始标记。
I-LOC：地名的中间标记。

代码示例：

1
from transformers import pipeline
2

3
nlp = pipeline("ner")
4

5
sequence = "Hugging Face Inc. is a company based in New York City. Its headquarters are in DUMBO, therefore very close to the Manhattan Bridge which is visible from the window."
6
print(nlp(sequence))

运行结果：

1
[
2
    {'word': 'Hu', 'score': 0.9995632767677307, 'entity': 'I-ORG'},
3
    {'word': '##gging', 'score': 0.9915938973426819, 'entity': 'I-ORG'},
4
    {'word': 'Face', 'score': 0.9982671737670898, 'entity': 'I-ORG'},
5
    {'word': 'Inc', 'score': 0.9994403719902039, 'entity': 'I-ORG'},
6
    {'word': 'New', 'score': 0.9994346499443054, 'entity': 'I-LOC'},
7
    {'word': 'York', 'score': 0.9993270635604858, 'entity': 'I-LOC'},
8
    {'word': 'City', 'score': 0.9993864893913269, 'entity': 'I-LOC'},
9
    {'word': 'D', 'score': 0.9825621843338013, 'entity': 'I-LOC'},
10
    {'word': '##UM', 'score': 0.936983048915863, 'entity': 'I-LOC'},
11
    {'word': '##BO', 'score': 0.8987102508544922, 'entity': 'I-LOC'},
12
    {'word': 'Manhattan', 'score': 0.9758241176605225, 'entity': 'I-LOC'},
13
    {'word': 'Bridge', 'score': 0.990249514579773, 'entity': 'I-LOC'}
14
]

使用模型和文本标记器#

过程如下：

实例化预训练模型和对应文本标记器。需要一个BERT模型。
定义模型训练时的标签列表。
定义一个有已知名命实体的序列，如“纽约市”。
分割单词用于模型预测。这里可以使用一个小技巧，先对序列编码再解码，我们就可以得到一个包含特殊标记的字符串，如BERT中的“”。
将序列编码为索引数组。（特殊标记会自动添加，不需要手动添加）
将数据输入模型获得的结果来进行标签预测。获得的结果是9个分类的概率分布。一般是使用最高概率的那个标签最为最终预测结果。
将每个标记与其预测标签一起打印出来。

示例代码

1
cache_dir="./transformersModels/ner"
2
"""
3
,cache_dir = cache_dir
4
"""
5
from transformers import AutoModelForTokenClassification, AutoTokenizer
6
import torch
7
model = AutoModelForTokenClassification.from_pretrained("dbmdz/bert-large-cased-finetuned-conll03-english",cache_dir = cache_dir, return_dict=True)
8
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased",cache_dir = cache_dir)
9

10
label_list = [
11
    "O",       # Outside of a named entity
12
    "B-MISC",  # Beginning of a miscellaneous entity right after another miscellaneous entity
13
    "I-MISC",  # Miscellaneous entity
14
    "B-PER",   # Beginning of a person's name right after another person's name
15
    "I-PER",   # Person's name
16
    "B-ORG",   # Beginning of an organisation right after another organisation
17
    "I-ORG",   # Organisation
18
    "B-LOC",   # Beginning of a location right after another location
19
    "I-LOC"    # Location
20
]
21

22
sequence = "Hugging Face Inc. is a company based in New York City. Its headquarters are in DUMBO, therefore very" \
23
           "close to the Manhattan Bridge."
24

25
# Bit of a hack to get the tokens with the special tokens
26
tokens = tokenizer.tokenize(tokenizer.decode(tokenizer.encode(sequence)))
27
inputs = tokenizer.encode(sequence, return_tensors="pt")
28

29
outputs = model(inputs).logits
30
predictions = torch.argmax(outputs, dim=2)
31

32
for token, prediction in zip(tokens, predictions[0].numpy()):
33
    print(token, label_list[prediction])

输出结果：

1
[CLS] O
2
Hu I-ORG
3
##gging I-ORG
4
Face I-ORG
5
Inc I-ORG
6
. O
7
is O
8
a O
9
company O
10
based O
11
in O
12
New I-LOC
13
York I-LOC
14
City I-LOC
15
. O
16
Its O
17
headquarters O
18
are O
19
in O
20
D I-LOC
21
##UM I-LOC
22
##BO I-LOC
23
, O
24
therefore O
25
very O
26
##c O
27
##lose O
28
to O
29
the O
30
Manhattan I-LOC
31
Bridge I-LOC
32
. O
33
[SEP] O

整理结果可以得到最终预测实体：

1
ORG:Hugging Face Inc
2
LOC:New York City
3
LOC:DUMBO
4
LOC:Manhattan Bridge

与pipline不同，这里并不会自动整合实体并删除“O”分类，需要自己写代码实现。

原文链接：CSDN