IDCNN模型实现NER任务

一、IDCNN模型介绍

IDCNN模型的结构和原理 IDCNN模型源于论文<< Fast and Accurate Entity Recognition with Iterated Dilated Convolutions >>, 核心操作是膨胀卷积. 膨胀卷积可以跳过⼀部分神经元, 同时扩⼤感受野的范围:

下图左侧展示膨胀率等于1的标准卷积操作, 也是经典CNN的实现⽅式; 下图右侧展示膨胀率等于2的膨胀卷积, 也是论⽂中要采用的方式:

当我们采⽤传统卷积操作, 神经⽹络的更⾼层可以如下图所示的获取⼀定范围内的⽂本语义信息:

当我们采⽤膨胀卷积操作, 神经⽹络的更⾼层可以如下图所示的获取更⼤范围内的⽂本语义信息：

论文中对几个经典NER模型做了对⽐测试, 显示IDCNN模型的效果可以媲美BiLSTM + CRF, 并显著的超越 BiLSTM模型本身:

IDCNN在⼯业界最⼤的亮点在于处理速度上, 论⽂中也进⾏了对⽐测试, 相⽐于效果相同的BiLSTM + CRF, 可 以达到8倍的推理加速:

二、序列标注

标注模式: NER任务本质上都是给任意token打上标签, 但是表层上又分为⼏种不同的标注模式.

• 模式⼀: BIO模式
• 模式⼆: BIEO模式
• 模式三: BIEOS模式

BIO模式: 包含3⼤类标签

• B: 任意命名实体的起始token
• I: 任意命名实体从中间到结尾的token
• O: 任意非命名实体的token

BIEO模式: 包含4⼤类标签

• B: 任意命名实体的起始token
• I: 任意命名实体中间的token
• E: 任意命名实体结尾的token
• O: 任意⾮命名实体的token

BIEOS模式: 包含5⼤类标签

• B: 任意命名实体的起始token
• I: 任意命名实体中间的token
• E: 任意命名实体结尾的token
• O: 任意⾮命名实体的token
• S: 任意只有单⼀字符实体的token

示例：

NER任务中两⼤重要矩阵：

发射矩阵: 可以理解为底层base model神经⽹络的输出张量.

转移矩阵: 可以理解为CRF模型.

三、实现代码

1.编写配置函数---config.py

# 设定label_to_id的映射字典, 采用BIEO标注模式, 外加3个特殊字符<pad>, <start>, <eos>
l2i_dic = {'O': 0, u'B-sym': 1, u'B-dis': 2, u'I-sym': 3, u'I-dis': 4, u'E-sym': 5, u'E-dis': 6, '<pad>': 7, '<start>': 8, '<eos>': 9}

# 上面的逆向字典, id_to_label
i2l_dic = {0: 'O', 1: u'B-sym', 2: u'B-dis', 3: u'I-sym', 4: u'I-dis', 5: u'E-sym', 6: u'E-dis', 7: '<pad>', 8: '<start>', 9: '<eos>'}


# 训练集, 测试集, 词表
train_file = './data/train.txt'
dev_file = './data/test.txt'
vocab_file = './data/vocab.txt'

save_model_path =  './saved_model/idcnn_crf.pt'

model_path = './saved_model/idcnn_crf_1.pt'

# 设置关键的超参数
max_length = 256
batch_size = 32
epochs = 30
lr = 1e-3
tagset_size = len(l2i_dic)
dropout = 0.4
use_cuda = True

2.编写工具类函数---utils.py

import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append('.')
from config import *


class InputFeatures(object):
    def __init__(self, text, label, input_id, label_id, input_mask):
        self.text = text
        self.label = label
        self.input_id = input_id
        self.label_id = label_id
        self.input_mask = input_mask


# 读取字典文件, 并构造字典
def load_vocab(vocab_file):
    vocab = {}
    index = 0
    with open(vocab_file, 'r', encoding='utf-8') as reader:
        while True:
            token = reader.readline()
            if not token:
                break
            token = token.strip()
            vocab[token] = index
            index += 1
    return vocab


# 加载训练集, 测试集的原始数据
def load_file(file_path):
    contents = open(file_path, encoding='utf-8').readlines()
    text =[]
    label = []
    texts = []
    labels = []
    for line in contents:
        if line != '\n':
            line = line.strip().split(' ')
            text.append(line[0])
            label.append(line[-1])
        else:
            texts.append(text)
            labels.append(label)
            text = []
            label = []
    return texts, labels


def load_data(file_path, max_length, label_dic, vocab):
    texts, labels = load_file(file_path)
    assert len(texts) == len(labels)
    result = []
    for i in range(len(texts)):
        assert len(texts[i]) == len(labels[i])
        token = texts[i]
        label = labels[i]
        if len(token) > max_length - 2:
            token = token[0: (max_length - 2)]
            label = label[0: (max_length - 2)]
        
        # 调用BERT的时候需要添加[CLS]和[SEP], 此处仿照同样的规则进行
        tokens_f =['[CLS]'] + token + ['[SEP]']

        # 调用IDCNN的时候添加首尾字符, 以对应上面的CLS和SEP
        label_f = ['<start>'] + label + ['<eos>']

        input_ids = [int(vocab[i]) if i in vocab else int(vocab['[UNK]']) for i in tokens_f]
        label_ids = [label_dic[i] for i in label_f]
        
        # 调用BERT的时候需要mask, 调用IDCNN的时候不需要mask
        input_mask = [1] * len(input_ids)
        while len(input_ids) < max_length:
            input_ids.append(0)
            input_mask.append(0)
            label_ids.append(label_dic['<pad>'])
        
        assert len(input_ids) == max_length
        assert len(input_mask) == max_length
        assert len(label_ids) == max_length
        
        feature = InputFeatures(text=tokens_f, label=label_f, input_id=input_ids,
                                input_mask=input_mask, label_id=label_ids)
        
        result.append(feature)
    
    return result


# 利用标签字典恢复真实的预测标签
def recover_label(pred_var, gold_var, l2i_dic, i2l_dic):
    assert len(pred_var) == len(gold_var)
    pred_variable = []
    gold_variable = []

    for i in range(len(gold_var)):
        start_index = gold_var[i].index(l2i_dic['<start>'])
        end_index = gold_var[i].index(l2i_dic['<eos>'])
        pred_variable.append(pred_var[i][s