本文主要讲述 W2NER 的代码，关于论文相关部分可阅读：统一NER为词词关系分类 | Yam。代码主要包括：输入、训练输出和解码部分，对于模型部分可参考前面的链接。
输入
假设使用如下量句训练数据：
1 | inp = [ |
首先是 bert_inputs，它和正常使用 BERT Tokenizer 直接 encode 是一样的。注意英文会有子词，所以长度比单词数要多。
1 | # tokens（按空格分开后再执行 tokenize） |
接下来是 grid_labels，具体结果如下：
1 | # grid_labels Bx8x8 |
以论文中的例子（第一句）来说明，左上三角的三个 1 就表示 NNW，右下三角的 3 则表示 THW（就是对应的类型：CONT），是实体的尾部。具体代码如下：
1 | for entity in instance["ner"]: |
需要注意的是，shape 的大小是根据句子的长度（词数）来确定的，第一句话长度为 8，第二句为 6，所以第二句会 Padding。
下一个是 grid_mask2d，这个和 grid_labels 对应，将其中 Padding 的部分 Mask 住。
1 | # Bx8x8 |
接下来是 pieces2word，这个主要是针对子词的，中文一般一个 Token 就是一个字（词），英文可能几个 Token 合成一个词。因为我们在 NER 时并不是针对子词，而是针对一个个独立的词的，所以这块需要单独标记。
注意和前面的 tokens 结合起来看。另外需要注意，这里的首尾各增加了 BERT 的特殊标记。
1 | # 中文 |
当然，如果两句话在一个 Batch 中，Padding 也是必须的，Padding 部分会被 Mask 住。
再接下来是位置编码：dist_inputs，根据词对的相对距离设计了 20 个 Embedding，并根据前后顺序与距离的 2 进制（2**0, 2**1, 2**2, ...）进行分配。
1 | # Bx8x8 |
代码如下：
1 | dis2idx = np.zeros((1000)) |
相对位置最大相差 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 | tril_mask = torch.tril(grid_mask2d.clone().long()) |
具体实例如下：
1 | # tril_mask |
注意 reg_inputs 的值域是 {0, 1, 2}，分别代表 Padding，上三角和下三角区域。
输出
最终输出的 logits 的 shape 为 B x SeqLen x SeqLen x LabelNum，换到本例中就是 2x8x8x10，取 argmax 后则变成 2x8x8 。
如果是训练阶段，logits 会和 grid_labels 一起计算损失，损失函数为交叉熵。
1 | criterion = nn.CrossEntropyLoss() |
注意这里用 grid_mask2d 将 Padding 的全部丢掉。
记录预测值时也用同样的方法操作：
1 | # 2x8x8 |
这里得到的 y_true 和 y_pred 也是没有 Padding 的结果。
解码
最后来看一下如何进行解码，官方代码如下（略作修改）：
1 | from collections import defaultdict |
上面的例子结果为：
1 | pred_labels = torch.argmax(outputs, -1) |
由于模型没有使用英文数据集，所以第一个句子无结果。
小结
整个模型在输入构造方面可谓设计颇多，但每一个又有其独特的意义，而且最后整体效果确实不错。值得一提的是，词信息其实并没有使用到训练过程中，与作者沟通后主要出于两个方面考虑。第一，本文主要为了证明模型在三种不同类型 NER 任务及中英文数据上的普遍性，所以提供的其实是一个兼容的方案，并没有针对这块单独设计和处理；第二，实验结果证明即使在没有分词和词典知识的情况下效果依然不错。因此，如果在中文任务上，可以结合词信息和位置编码进行更多的尝试。