论文将命名实体识别任务转换为机器阅读理解任务MRC，也就是说，文本序列中对应的实体通过问一个问题来提取；一般来说，具体类别的问题是提取的org实体类别，query文本序列中的组织是什么？”。其使用BERT作为backbone，将文本和问句作为序列BERT，使用两个二分类器对BERT对最终数据进行分类，一个分类器判断每个数据token另一个分类器判断每个实体开始索引的可能性token实体结束索引的可能性。

• datasets --构建数据集文件
• evaluate --用于评估的文件
• inference --用于前向推理的文件
• metrics --实现metric的文件
• models --构建模型的文件
• ner2mrc --将数据转换为mrc所需格式的文件
• scripts --在每个数据集上训练的启动文件
• tests --测试所需的文件
• train --模型训练文件
• utils --其他的辅助文件
• README.md --项目详情
• requirements.txt --项目所需的python包

以下是官方代码的适当修改和详细注释。通过阅读以下代码和注释，您可以清楚地了解项目，但需要更多的时间反复思考；有效的方法是下载官方代码库并运行整个代码pipeline在主题文件中，找出每个模块过程

先针对models文件分析在路径下，了解模型构建的整个过程

使用文件MRC实体提取任务和框架Tag实体抽取任务的方式分别定义了所需的config类

#!/usr/bin/env python3 # -*- coding: utf-8 -*-  # file: model_config.py  from transformers import BertConfig   class BertQueryNerConfig(BertConfig):    # 使用MRC框架的实体提取所需config     def __init__(self, **kwargs):         super(BertQueryNerConfig, self).__init__(**kwargs)         self.mrc_dropout = kwargs.get("mrc_dropout", 0.1)         self.classifier_intermediate_hidden_size = kwargs.get("classifier_intermediate_hidden_size", 1024)         self.classifier_act_func = kwargs.get("classifier_act_func", "gelu")   class BertTaggerConfig(BertConfig
       
        )
        : 
        # 使用tag方式的实体抽取所需的config 
        def 
        __init__
        (self
        , 
        **kwargs
        )
        : 
        super
        (BertTaggerConfig
        , self
        )
        .__init__
        (
        **kwargs
        ) self
        .num_labels 
        = kwargs
        .get
        (
        "num_labels"
        , 
        6
        ) self
        .classifier_dropout 
        = kwargs
        .get
        (
        "classifier_dropout"
        , 
        0.1
        ) self
        .classifier_sign 
        = kwargs
        .get
        (
        "classifier_sign"
        , 
        "multi_nonlinear"
        ) self
        .classifier_act_func 
        = kwargs
        .get
        (
        "classifier_act_func"
        , 
        "gelu"
        ) self
        .classifier_intermediate_hidden_size 
        = kwargs
        .get
        (
        "classifier_intermediate_hidden_size"
        , 
        1024
        ) 
       

为两种实体抽取方式分别定义分类头

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# file: classifier.py

import torch.nn as nn
from torch.nn import functional as F


class SingleLinearClassifier(nn.Module):  # 通过一个全连接层直接进行类别预测
    def __init__(self, hidden_size, num_label):
        super(SingleLinearClassifier, self).__init__()
        self.num_label = num_label
        self.classifier = nn.Linear(hidden_size, num_label)

    def forward(self, input_features):
        features_output = self.classifier(input_features)
        return features_output


class MultiNonLinearClassifier(nn.Module):  # MRC框架使用的分类头，使用了两个全连接层对bert输出的特征进行处理
    def __init__(self, hidden_size, num_label, dropout_rate, act_func="gelu", intermediate_hidden_size=None):
        super(MultiNonLinearClassifier, self).__init__()
        self.num_label = num_label
        self.intermediate_hidden_size = hidden_size if intermediate_hidden_size is None else intermediate_hidden_size
        self.classifier1 = nn.Linear(hidden_size, self.intermediate_hidden_size)
        self.classifier2 = nn.Linear(self.intermediate_hidden_size, self.num_label)
        self.dropout = nn.Dropout(dropout_rate)
        self.act_func = act_func

    def forward(self, input_features):
        features_output1 = self.classifier1(input_features)
        if self.act_func == "gelu":
            features_output1 = F.gelu(features_output1)
        elif self.act_func == "relu":
            features_output1 = F.relu(features_output1)
        elif self.act_func == "tanh":
            features_output1 = F.tanh(features_output1)
        else:
            raise ValueError
        features_output1 = self.dropout(features_output1)
        features_output2 = self.classifier2(features_output1)
        return features_output2


class BERTTaggerClassifier(nn.Module):  # tag方式的分类头
    def __init__(self, hidden_size, num_label, dropout_rate, act_func="gelu", intermediate_hidden_size=None):
        super(BERTTaggerClassifier, self).__init__()
        self.num_label = num_label
        self.intermediate_hidden_size = hidden_size if intermediate_hidden_size is None else intermediate_hidden_size
        self.classifier1 = nn.Linear(hidden_size, self.intermediate_hidden_size)
        self.classifier2 = nn.Linear(self.intermediate_hidden_size, self.num_label)
        self.dropout = nn.Dropout(dropout_rate)
        self.act_func = act_func

    def forward(self, input_features):
        features_output1 = self.classifier1(input_features)
        if self.act_func == "gelu":
            features_output1 = F.gelu(features_output1)
        elif self.act_func == "relu":
            features_output1 = F.relu(features_output1)
        elif self.act_func == "tanh":
            features_output1 = F.tanh(features_output1)
        else:
            raise ValueError
        features_output1 = self.dropout(features_output1)
        features_output2 = self.classifier2(features_output1)
        return features_output2

使用bert进行tag方式的实体抽取

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# file: bert_tagger.py
#

import torch.nn as nn
from transformers import BertModel, BertPreTrainedModel
from models.classifier import BERTTaggerClassifier


# 直接把bert最后一层输出的隐含状态直接连接一个分类器进行状态分类
class BertTagger(BertPreTrainedModel):
    def __init__(self, config):
        super(BertTagger, self).__init__(config)
        self.bert = BertModel(config)  # 基于config初始化bert模型

        self.num_labels = config.num_labels
        self.hidden_size = config.hidden_size
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        if config.classifier_sign == "multi_nonlinear":  # 调用tag方式的分类头
            self.classifier = BERTTaggerClassifier(self.hidden_size, self.num_labels,
                                                   config.classifier_dropout,
                                                   act_func=config.classifier_act_func,
                                                   intermediate_hidden_size=config.classifier_intermediate_hidden_size)
        else:
            self.classifier = nn.Linear(self.hidden_size, self.num_labels)

        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None,):
        last_bert_layer, pooled_output = self.bert(input_ids, token_type_ids, attention_mask)
        last_bert_layer = last_bert_layer.view(-1, self.hidden_size)
        last_bert_layer = self.dropout(last_bert_layer)
        logits = self.classifier(last_bert_layer)
        return logits

如论文中一样，分别对文本序列中每个token进行实体开始索引预测和结束索引预测，再计算开始索引与结束索引的匹配预测结果进行返回

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# file: bert_query_ner.py

import torch
import torch.nn as nn
from transformers import BertModel, BertPreTrainedModel
from models.classifier import MultiNonLinearClassifier


class BertQueryNER(BertPreTrainedModel):
    def __init__(self, config):
        super(BertQueryNER, self).__init__(config)
        self.bert = BertModel(config)  # 初始化bert

        self.start_outputs = nn.Linear(config.hidden_size, 1)  # 用于计算实体开始索引
        self.end_outputs = nn.Linear(config.hidden_size, 1)  # 用于计算实体结束索引
        # 判断i、j是否为一个匹配
        self.span_embedding = MultiNonLinearClassifier(config.hidden_size * 2, 1, config.mrc_dropout,
                                                       intermediate_hidden_size=config.classifier_intermediate_hidden_size)

        self.hidden_size = config.hidden_size

        self.init_weights()

    def forward(self, input_ids, token_type_ids=None, attention_mask=None):
        """ Args: input_ids: bert input tokens, tensor of shape [seq_len] token_type_ids: 0 for query, 1 for context, tensor of shape [seq_len]，query在前，文本在后 attention_mask: attention mask, tensor of shape [seq_len] Returns: start_logits: start/non-start probs of shape [seq_len] end_logits: end/non-end probs of shape [seq_len] match_logits: start-end-match probs of shape [seq_len, 1]，此处的1表示匹配行的得分 """

        bert_outputs = self.bert(input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)

        sequence_heatmap = bert_outputs[0]  # [batch, seq_len, hidden]
        batch_size, seq_len, hid_size = sequence_heatmap.size()

        start_logits = self.start_outputs(sequence_heatmap).squeeze(-1)  # [batch, seq_len, 1]，开始索引预测
        end_logits = self.end_outputs(sequence_heatmap).squeeze(-1)  # [batch, seq_len, 1]，结束索引预测

        # for every position $i$ in sequence, should concate $j$ to
        # predict if $i$ and $j$ are start_pos and end_pos for an entity.
        # [batch, seq_len, hidden]->[batch, seq_len, 1, hidden]->[batch, seq_len, seq_len, hidden]
        start_extend = sequence_heatmap.unsqueeze(2).expand(-1, -1, seq_len, -1)
        # [batch, seq_len, hidden]->[batch, 1, seq_len, hidden]->[batch, seq_len, seq_len, hidden]
        end_extend = sequence_heatmap.unsqueeze(1).expand(-1, seq_len, -1, -1)
        # [batch, seq_len, seq_len, hidden]+[batch, seq_len, seq_len, hidden]->[batch, seq_len, seq_len, hidden*2]
        span_matrix = torch.cat([start_extend, end_extend], 3)
        # [batch, seq_len, seq_len, hidden*2]->[batch, seq_len, seq_len, 1]->[batch, seq_len, seq_len]
        span_logits = self.span_embedding(span_matrix).squeeze(-1)  # 开始索引和结束索引匹配情况的预测

        return start_logits, end_logits, span_logits

该项目主要使用Pytoch Lightning框架实现训练、测试等过程，使用Pytorch Lightning只用定义主要的训练/training_step()、验证/validation_step()和测试/test_step()函数，而不用写复杂的for循环，框架会自动进行训练。代码如下，该代码是在官方代码的基础上进行调整，增加了使用WandbLogger进行数据记录，可配合笔记进行阅读