利用tensorflow2自带keras搭建BiLSTM+CRF的序列标注模型，完成中文的命名实体识别任务。这里使用数据集是提前处理过的，已经转成命名实体识别需要的“BIO”标注格式。

详细代码和数据：https://github.com/huanghao128/zh-nlp-demo

模型结构

BiLSTM+CRF模型就是在双向LSTM模型的输出位置接上一个CRF层，这样可以学习到相邻输出之间的依赖关系，从而提高输出标签的整理准确率，模型结构图如下图：

数据处理

本文使用的数据是已经预处理过的，所以直接加载数据就好了，首先我们要加载字符词典文件，还有BIO标记类别的索引化。其中BIO标记中B-PER和I-PER表示人名，B-LOC和I-LOC表示地名，B-ORG和I-ORG表示机构名。

char_vocab_path = "./data/char_vocabs.txt" # 字典文件
train_data_path = "./data/train_data" # 训练数据
test_data_path = "./data/test_data" # 测试数据

special_words = ['<PAD>', '<UNK>'] # 特殊词表示

# "BIO"标记的标签
label2idx = {"O": 0,
             "B-PER": 1, "I-PER": 2,
             "B-LOC": 3, "I-LOC": 4,
             "B-ORG": 5, "I-ORG": 6
             }
# 索引和BIO标签对应
idx2label = {idx: label for label, idx in label2idx.items()}

# 读取字符词典文件
with open(char_vocab_path, "r", encoding="utf8") as fo:
    char_vocabs = [line.strip() for line in fo]
char_vocabs = special_words + char_vocabs

# 字符和索引编号对应
idx2vocab = {idx: char for idx, char in enumerate(char_vocabs)}
vocab2idx = {char: idx for idx, char in idx2vocab.items()}

然后加载训练和测试集，并把原始数据和BIO标记转成索引和类别编号。

# 读取数据集语料
def read_corpus(corpus_path, vocab2idx, label2idx):
    datas, labels = [], []
    with open(corpus_path, encoding='utf-8') as fr:
        lines = fr.readlines()
    sent_, tag_ = [], []
    for line in lines:
        if line != '\n':
            char, label = line.strip().split()
            sent_.append(char)
            tag_.append(label)
        else:
            sent_ids = [vocab2idx[char] if char in vocab2idx else vocab2idx['<UNK>'] for char in sent_]
            tag_ids = [label2idx[label] if label in label2idx else 0 for label in tag_]
            datas.append(sent_ids)
            labels.append(tag_ids)
            sent_, tag_ = [], []
    return datas, labels
# 加载训练集
train_datas, train_labels = read_corpus(train_data_path, vocab2idx, label2idx)
# 加载测试集
test_datas, test_labels = read_corpus(test_data_path, vocab2idx, label2idx)

数据的填充，以及类别的one-hot编码。

from tensorflow import keras
from keras.preprocessing import sequence
import numpy as np

MAX_LEN = 100
VOCAB_SIZE = len(vocab2idx)
CLASS_NUMS = len(label2idx)

# padding data
train_datas = sequence.pad_sequences(train_datas, maxlen=MAX_LEN)
train_labels = sequence.pad_sequences(train_labels, maxlen=MAX_LEN)
test_datas = sequence.pad_sequences(test_datas, maxlen=MAX_LEN)
test_labels = sequence.pad_sequences(test_labels, maxlen=MAX_LEN)
print('x_train shape:', train_datas.shape)
print('x_test shape:', test_datas.shape)
# encoder one-hot
train_labels = keras.utils.to_categorical(train_labels, CLASS_NUMS)
test_labels = keras.utils.to_categorical(test_labels, CLASS_NUMS)
print('trainlabels shape:', train_labels.shape)
print('testlabels shape:', test_labels.shape)

模型构建和训练

模型构建主要使用keras自带的基础模型组装，首先是双向LSTM模型，然后输出接CRF模型，输出对每个时刻的分类。

from tensorflow import keras
from tensorflow.keras import layers, models
from tensorflow.keras import backend as K

class CRF(layers.Layer):
    def __init__(self, label_size):
        super(CRF, self).__init__()
        self.trans_params = tf.Variable(
            tf.random.uniform(shape=(label_size, label_size)), name="transition")
    
    @tf.function
    def call(self, inputs, labels, seq_lens):
        log_likelihood, self.trans_params = tfa.text.crf_log_likelihood(
                                                inputs, labels, seq_lens,
                                                transition_params=self.trans_params)
        loss = tf.reduce_sum(-log_likelihood)
        return loss

EPOCHS = 20
BATCH_SIZE = 64
EMBED_DIM = 128
HIDDEN_SIZE = 64
MAX_LEN = 100
VOCAB_SIZE = len(vocab2idx)
CLASS_NUMS = len(label2idx)

inputs = layers.Input(shape=(MAX_LEN,), name='input_ids', dtype='int32')
targets = layers.Input(shape=(MAX_LEN,), name='target_ids', dtype='int32')
seq_lens = layers.Input(shape=(), name='input_lens', dtype='int32')

x = layers.Embedding(input_dim=VOCAB_SIZE, output_dim=EMBED_DIM, mask_zero=True)(inputs)
x = layers.Bidirectional(layers.LSTM(HIDDEN_SIZE, return_sequences=True))(x)
logits = layers.Dense(CLASS_NUMS)(x)
loss = CRF(label_size=CLASS_NUMS)(logits, targets, seq_lens)

model = models.Model(inputs=[inputs, targets, seq_lens], outputs=loss)

print(model.summary())

模型训练

由于构建模型时，输出对应的是loss，所以loss定义需要进行一些修改。

class CustomLoss(keras.losses.Loss):
    def call(self, y_true, y_pred):
        loss, pred = y_pred
        return loss

# 自定义Loss
# model.compile(loss=CustomLoss(), optimizer='adam')
# 或者使用lambda表达式
model.compile(loss=lambda y_true, y_pred: y_pred, optimizer='adam')

# 训练
model.fit(x=[train_datas, train_labels, train_seq_lens], y=labels, 
        validation_split=0.1, batch_size=BATCH_SIZE, epochs=EPOCHS)

# 保存
model.save("output/bilstm_crf_ner")

结果预测

结果预测是我们训练好模型后，重新加载模型，输入新的要预测文本，然后识别出文本中的命名实体。这里首先要加载字符词典，然后加载模型，之后对输入文本预处理成字符序列，然后模型预测每个时刻的输出类别，最后把类别转成BIO标记，BIO标记组合成正确的命名实体。

# 加载模型
model = models.load_model("output/bilstm_crf_ner", compile=False)
# 如果需要继续训练，需要下面的重新compile
# model.compile(loss=lambda y_true, y_pred: y_pred, optimizer='adam')

分别提取转移矩阵参数和BiLSTM的输出

# 提取转移矩阵参数
trans_params = model.get_layer('crf').get_weights()[0]

# 获得BiLSTM的输出logits
sub_model = models.Model(inputs=model.get_layer('input_ids').input,
                        outputs=model.get_layer('dense').output)

def predict(model, inputs, input_lens):
    logits = sub_model.predict(inputs)
    # 获取CRF层的转移矩阵
    # crf_decode：viterbi解码获得结果
    pred_seq, viterbi_score = tfa.text.crf_decode(logits, trans_params, input_lens)
    return pred_seq

测试预测数据

maxlen = 100
sentence = "这里输入是需要实体识别的句子"
sent_chars = list(sentence)
sent2id = [vocab2idx[word] if word in vocab2idx else vocab2idx['<UNK>'] for word in sent_chars]
sent2id_new = np.array([[0] * (maxlen-len(sent2id)) + sent2id[:maxlen]])
test_lens = np.array([100])

pred_seq = predict(model, sent2id_new, test_lens)
print(pred_seq)

对预测结果进行命名实体解析和提取

def get_valid_nertag(input_data, result_tags):
    result_words = []
    start, end =0, 1 # 实体开始结束位置标识
    tag_label = "O" # 实体类型标识
    for i, tag in enumerate(result_tags):
        if tag.startswith("B"):
            if tag_label != "O": # 当前实体tag之前有其他实体
                result_words.append((input_data[start: end], tag_label)) # 获取实体
            tag_label = tag.split("-")[1] # 获取当前实体类型
            start, end = i, i+1 # 开始和结束位置变更
        elif tag.startswith("I"):
            temp_label = tag.split("-")[1]
            if temp_label == tag_label: # 当前实体tag是之前实体的一部分
                end += 1 # 结束位置end扩展
        elif tag == "O":
            if tag_label != "O": # 当前位置非实体 但是之前有实体
                result_words.append((input_data[start: end], tag_label)) # 获取实体
                tag_label = "O"  # 实体类型置"O"
            start, end = i, i+1 # 开始和结束位置变更
    if tag_label != "O": # 最后结尾还有实体
        result_words.append((input_data[start: end], tag_label)) # 获取结尾的实体
    return result_words

详细代码和数据：https://github.com/huanghao128/zh-nlp-demo