W2NER 代码

本文主要讲述 W2NER 的代码，关于论文相关部分可阅读：统一NER为词词关系分类 | Yam。代码主要包括：输入、训练输出和解码部分，对于模型部分可参考前面的链接。

输入

假设使用如下量句训练数据：

1
2
3
4
5
6
7
8
inp = [
	{'ner': [{'index': [3, 4, 5], 'type': 'CONT'}, {'index': [3, 4, 7], 'type': 'CONT'}],
	 'sentence': ['i', 'am', 'having', 'aching', 'in', 'legs', 'and', 'shoulders'],
	 'word': []},
	{'ner': [{'index': [0, 1, 2], 'type': 'NAME'}],
	'sentence': ['常', '建', '良', '，', '男', '，'],
	'word': [[0], [1, 2], [3], [4], [5]]}
]

首先是 bert_inputs，它和正常使用 BERT Tokenizer 直接 encode 是一样的。注意英文会有子词，所以长度比单词数要多。

1
2
3
4
5
6
7
8
9
10
# tokens（按空格分开后再执行 tokenize）
[['i'], ['am'], ['ha', '##ving'], ['ac', '##hing'], ['in'], ['le', '##gs'], ['and'], ['sh', '##ould', '##ers']]
[['常'], ['建'], ['良'], ['，'], ['男'], ['，']]

# pieces（将上面的tokens打平）
['i', 'am', 'ha', '##ving', 'ac', '##hing', 'in', 'le', '##gs', 'and', 'sh', '##ould', '##ers']
['常', '建', '良', '，', '男', '，']

# bert_inputs
tensor([[101,151,8413,11643,10369,9226,10716,8217,8983,9726,8256,11167,11734,8755,102],[101,2382,2456,5679,8024,4511,8024,102,0,0,0,0,0,0,0]])

接下来是 grid_labels，具体结果如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# grid_labels Bx8x8
tensor([[[0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 1, 0, 0, 0],
         [0, 0, 0, 0, 0, 1, 0, 1],
         [0, 0, 0, 3, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 3, 0, 0, 0, 0]],

        [[0, 1, 0, 0, 0, 0, 0, 0],
         [0, 0, 1, 0, 0, 0, 0, 0],
         [2, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0]]])

以论文中的例子（第一句）来说明，左上三角的三个 1 就表示 NNW，右下三角的 3 则表示 THW（就是对应的类型：CONT），是实体的尾部。具体代码如下：

1
2
3
4
5
for entity in instance["ner"]:
    index = entity["index"]
    for i in range(len(index) - 1):
        grid_labels[index[i], index[i + 1]] = 1
    grid_labels[index[-1], index[0]] = type_id

需要注意的是，shape 的大小是根据句子的长度（词数）来确定的，第一句话长度为 8，第二句为 6，所以第二句会 Padding。

下一个是 grid_mask2d，这个和 grid_labels 对应，将其中 Padding 的部分 Mask 住。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Bx8x8
tensor([[[ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True, False, False],
         [False, False, False, False, False, False, False, False],
         [False, False, False, False, False, False, False, False]]])

接下来是 pieces2word，这个主要是针对子词的，中文一般一个 Token 就是一个字（词），英文可能几个 Token 合成一个词。因为我们在 NER 时并不是针对子词，而是针对一个个独立的词的，所以这块需要单独标记。

注意和前面的 tokens 结合起来看。另外需要注意，这里的首尾各增加了 BERT 的特殊标记。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 中文
array([[False,  True, False, False, False, False, False, False],
       [False, False,  True, False, False, False, False, False],
       [False, False, False,  True, False, False, False, False],
       [False, False, False, False,  True, False, False, False],
       [False, False, False, False, False,  True, False, False],
       [False, False, False, False, False, False,  True, False]])

# 英文
array([[False,  True, False, False, False, False, False, False, False, False],
       [False, False,  True, False, False, False, False, False, False, False],
       [False, False, False,  True,  True, False, False, False, False, False],
       [False, False, False, False, False,  True,  True, False, False, False],
       [False, False, False, False, False, False, False,  True, False, False],
       [False, False, False, False, False, False, False, False,  True,  True],
       [False, False, False, False, False, False, False, False, False, False],
       [False, False, False, False, False, False, False, False, False, False]])

当然，如果两句话在一个 Batch 中，Padding 也是必须的，Padding 部分会被 Mask 住。

再接下来是位置编码：dist_inputs，根据词对的相对距离设计了 20 个 Embedding，并根据前后顺序与距离的 2 进制（2**0, 2**1, 2**2, ...）进行分配。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Bx8x8
tensor([[[19, 10, 11, 11, 12, 12, 12, 12],
         [ 1, 19, 10, 11, 11, 12, 12, 12],
         [ 2,  1, 19, 10, 11, 11, 12, 12],
         [ 2,  2,  1, 19, 10, 11, 11, 12],
         [ 3,  2,  2,  1, 19, 10, 11, 11],
         [ 3,  3,  2,  2,  1, 19, 10, 11],
         [ 3,  3,  3,  2,  2,  1, 19, 10],
         [ 3,  3,  3,  3,  2,  2,  1, 19]],

        [[19, 10, 11, 11, 12, 12,  0,  0],
         [ 1, 19, 10, 11, 11, 12,  0,  0],
         [ 2,  1, 19, 10, 11, 11,  0,  0],
         [ 2,  2,  1, 19, 10, 11,  0,  0],
         [ 3,  2,  2,  1, 19, 10,  0,  0],
         [ 3,  3,  2,  2,  1, 19,  0,  0],
         [ 0,  0,  0,  0,  0,  0,  0,  0],
         [ 0,  0,  0,  0,  0,  0,  0,  0]]])

代码如下：

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
dis2idx = np.zeros((1000))
dis2idx[1] = 1
dis2idx[2:] = 2
dis2idx[4:] = 3
dis2idx[8:] = 4
dis2idx[16:] = 5
dis2idx[32:] = 6
dis2idx[64:] = 7
dis2idx[128:] = 8
dis2idx[256:] = 9

length = 10

dist_inputs = np.zeros((length, length), dtype=np.int_)

for k in range(length):
    dist_inputs[k, :] += k
    dist_inputs[:, k] -= k

for i in range(length):
    for j in range(length):
        if dist_inputs[i, j] < 0:
            dist_inputs[i, j] = dis2idx[-dist_inputs[i, j]] + 9
        else:
            dist_inputs[i, j] = dis2idx[dist_inputs[i, j]]
# 对角线
dist_inputs[dist_inputs == 0] = 19

相对位置最大相差 9，句子长度可达 2**10，不过代码中将其限制到了 1000。

下一个是句子长度，比较简单：

1
tensor([8, 6])

最后是 entity_text，其实就是实体对应的 index 和实体类型的 id，这个用于评估。

1
[{'3-4-5-#-3', '3-4-7-#-3'}, {'0-1-2-#-2'}]

其实，输入中还少了一个（在模型内部实现）用来区分上下三角的 reg_inputs，它等于对 grid_mask 下三角 Mask 与 grid_mask 之和：

1
2
tril_mask = torch.tril(grid_mask2d.clone().long())
reg_inputs = tril_mask + grid_mask2d.clone().long()

具体实例如下：

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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
# tril_mask
tensor([[[1, 0, 0, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 1, 0],
         [1, 1, 1, 1, 1, 1, 1, 1]],

        [[1, 0, 0, 0, 0, 0, 0, 0],
         [1, 1, 0, 0, 0, 0, 0, 0],
         [1, 1, 1, 0, 0, 0, 0, 0],
         [1, 1, 1, 1, 0, 0, 0, 0],
         [1, 1, 1, 1, 1, 0, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0]]])

# gird_mask2d
tensor([[[1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1]],

        [[1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [1, 1, 1, 1, 1, 1, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0]]])

# reg_inputs
tensor([[[2, 1, 1, 1, 1, 1, 1, 1],
         [2, 2, 1, 1, 1, 1, 1, 1],
         [2, 2, 2, 1, 1, 1, 1, 1],
         [2, 2, 2, 2, 1, 1, 1, 1],
         [2, 2, 2, 2, 2, 1, 1, 1],
         [2, 2, 2, 2, 2, 2, 1, 1],
         [2, 2, 2, 2, 2, 2, 2, 1],
         [2, 2, 2, 2, 2, 2, 2, 2]],

        [[2, 1, 1, 1, 1, 1, 0, 0],
         [2, 2, 1, 1, 1, 1, 0, 0],
         [2, 2, 2, 1, 1, 1, 0, 0],
         [2, 2, 2, 2, 1, 1, 0, 0],
         [2, 2, 2, 2, 2, 1, 0, 0],
         [2, 2, 2, 2, 2, 2, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0]]])

注意 reg_inputs 的值域是 {0, 1, 2}，分别代表 Padding，上三角和下三角区域。

输出

最终输出的 logits 的 shape 为 B x SeqLen x SeqLen x LabelNum，换到本例中就是 2x8x8x10，取 argmax 后则变成 2x8x8 。

如果是训练阶段，logits 会和 grid_labels 一起计算损失，损失函数为交叉熵。

1
2
3
criterion = nn.CrossEntropyLoss()
# outputs 2x8x8x10, gird_labels 2x8x8, grid_mask2d 2x8x8
loss = criterion(outputs[grid_mask2d], grid_labels[grid_mask2d])

注意这里用 grid_mask2d 将 Padding 的全部丢掉。

记录预测值时也用同样的方法操作：

1
2
3
4
# 2x8x8
pred = torch.argmax(outputs, -1)
y_true = grid_labels[grid_mask2d].contiguous().view(-1)
y_pred = pred[grid_mask2d].contiguous().view(-1)

这里得到的 y_true 和 y_pred 也是没有 Padding 的结果。

解码

最后来看一下如何进行解码，官方代码如下（略作修改）：

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
39
40
41
42
from collections import defaultdict

def decode(outputs, length):
    decode_entities = []
    for index, (instance, l) in enumerate(zip(outputs, length)):
        # 获取实体和类型的index
        forward_dict = defaultdict(list)
        head_dict = defaultdict(set)
        ht_type_dict = {}
        for i in range(l):
            for j in range(i + 1, l):
                if instance[i, j] == 1:
                    forward_dict[i].append(j)
        for i in range(l):
            for j in range(i, l):
                if instance[j, i] > 1:
                    ht_type_dict[(i, j)] = instance[j, i].numpy().tolist()
                    head_dict[i].add(j)

        # 递归执行解码
        predicts = []
        def find_entity(key, entity, tails):
            entity.append(key)
            if key not in forward_dict:
                if key in tails:
                    typ = ht_type_dict[(entity[0], entity[-1])]
                    predicts.append((entity.copy(), typ))
                entity.pop()
                return
            else:
                if key in tails:
                    typ = ht_type_dict[(entity[0], entity[-1])]
                    predicts.append((entity.copy(), typ))
            for k in forward_dict[key]:
                find_entity(k, entity, tails)
            entity.pop()

        for head in head_dict:
            find_entity(head, [], head_dict[head])
        
        decode_entities.append(predicts)
    return decode_entities

上面的例子结果为：

1
2
3
pred_labels = torch.argmax(outputs, -1)
decode(pred_labels, sent_length)
# [[], [([0, 1, 2], 2)]]

由于模型没有使用英文数据集，所以第一个句子无结果。

小结

整个模型在输入构造方面可谓设计颇多，但每一个又有其独特的意义，而且最后整体效果确实不错。值得一提的是，词信息其实并没有使用到训练过程中，与作者沟通后主要出于两个方面考虑。第一，本文主要为了证明模型在三种不同类型 NER 任务及中英文数据上的普遍性，所以提供的其实是一个兼容的方案，并没有针对这块单独设计和处理；第二，实验结果证明即使在没有分词和词典知识的情况下效果依然不错。因此，如果在中文任务上，可以结合词信息和位置编码进行更多的尝试。