【信息抽取】UIE——基于prompt的信息抽取模型（附源码）-技术圈

作者简介

作者：何枝

原文：

https://zhuanlan.zhihu.com/p/589054073

转载者：杨夕

推荐系统百面百搭地址：

https://github.com/km1994/RES-Interview-Notes

NLP 百面百搭地址：

https://github.com/km1994/NLP-Interview-Notes

个人笔记：

https://github.com/km1994/nlp_paper_study

1. 什么是信息抽取（Information Extraction）

信息抽取是NLP任务中非常常见的一种任务，其目的在于从一段自然文本中提取出我们想要的关键信息结构。

举例来讲，现在有下面这样一个句子：

新东方烹饪学校在成都。

我们想要提取这句话中所有有意义的词语，例如：

机构	新东方烹饪学校
城市	成都

这个关键词提取任务就叫做命名实体识别（Named Entity Recognition, NER）任务，文中的「新东方烹饪学校」和「成都」就被称为实体（Entity）。

如果我们还想进一步的知道这些词语之间的关系，例如：

实体1	关系名	实体2
新东方烹饪学校	所在地	成都

这种提取实体之间关系的任务就叫做关系抽取（Relation Extraction, RE）任务。

2. 信息抽取的几种方法

2.1 序列标注（Sequence Labeling）

序列标注通常是指对文中的每一个字（以下简称token）进行分类，即本质是 token classification任务。

我们对第一小节中的例子做序列标注任务，得到的结果如下：

新	东	...	学	校	在	成	都
B-机构	I-机构	I-机构*N	I-机构	I-机构	O	B-城市	I-城市

可以看到，我们对句子中的每一个字（token）都打上了一个类别标签，我们期望模型要做的事就是去学会每一个字所属的类别是什么。

Note: 这里用的标注方法是「BIO 标记法」，其中「B-」代表该位置 token 是某一个实体词语（span）的起始 token；「I-」代表该位置 token 处于某一个词语的中间（或结尾），「O」则代表该位置 token 不在任何一个实体词语中。除了「BIO 标记法」外，还有许多其他的标注方式（如 BIOES 等），其本质思路都很类似。

2.2 指针网络（Pointer Network）

序列标注模型有一个天然的缺陷，无法解决解决实体重叠（overlap）的问题。

举例来讲，如果今天我们不仅要提取「机构」，还同时要提取「机构类型」，那么我们期望的提取结果应该为：

机构	新东方烹饪学校
机构类型	学校
城市	成都

可以看到，对于「学校」这两个字，即属于「新东方烹饪学校」（机构）这个词，也存在于「学校」（机构类型）这个词，那我们在给这两个字打标签的时候，究竟应该打成哪个类别呢？

新	东	…	学	校	在	成	都
B-机构	I-机构	I-机构	？	？	O	B-城市	I-城市

由此我们可以看到，因为在进行分类时我们通常对一个字（token）只赋予一个标签，这就导致了 token classification 不能很好的解决实体重叠（一字多标签）的复杂情况。

Note: 存在一些技巧可以解决该问题，例如可以从单字单分类（CE）衍生到单字多分类（BCE），这里不展开讨论。

指针网络（Pointer Network）通过分别对每一个实体单独做预测来解决了实体之前的重叠冲突问题。

例如，我们现在要同时预测「机构」和「机构类型」这两个实体，那么我们就可以设计一个多头网络（Multi-Head）来分别预测这两个实体的实体词。

其中，

「机构」实体头中「起始」向量代表这一句话中是「机构」词语的首字（例子中为「新」）；

「机构」实体中「结束」向量代表这一句话中时「机构」词语的尾字（例子中为「校」）。

通过「起始」和「结束」向量中的首尾字索引就能找到对应实体的词语。

可以看到，通过构建多头的任务，指针网络能够分别预测「机构」和「机构类型」中的实体词起始/终止位置，即「学校」这个词语在两个任务层中都能被抽取出来。

3. UIE —— 基于 prompt 的指针网络

3.1 UIE中的 prompt 是什么？

多头指针网络能够很好的解决实体重叠问题，但缺点在于：不够灵活。

假定今天我们已经通过指针网络训练好了一个提取「机构」、「机构类型」的模型，即将交付时甲方突然提出一个新需求：我们想再多提取一个「机构简称」的属性。

草（一种植物）。

从 2.2 节中的示意图中我们可以看到，每一个实体类型会对应一个单独的网络头。

这就意味着我们不仅需要重标数据，还需要为新属性添加一个新的网络头，即模型结构会随着实体类型个数改变而发生变化。

那，能不能有一种办法去固定住模型的结构，不管今天来多少种类型要识别都能使用同样的模型结构完成呢？

我们思考一下，模型结构变化的部分是和实体类型强绑定的「头」部分。

而不同「头」之间结构其实是完全一样的：一个「起始」向量 + 一个「终止」向量。

既然「头」结构完全一样，我们能不能干脆直接使用一个「头」去提取不同实体类型的信息呢？

不同「头」之间的区别在于它们关注的信息不同：「机构头」只关注「机构」相关的实体词，「城市头」只关注「城市」相关的实体词。

那么我们是不是可以直接在模型输入的时候就告诉模型：我现在需要提取「某个头」的信息。

这个用来告诉模型做具体任务的参数就叫 prompt，我们把它拼在输入中一并喂给模型即可。

通过上图可以看到，我们将不同的「实体类型」作为 prompt 参数喂给模型，用于「激活」模型参数跟当前「实体类型」相关的参数，从而输出不同的抽取结果。

Note: 「通过一个输入参数去激活一个大模型中的不同参数，从而完成不同任务的思路」并不是首次出现，在 meta-learning 中也存在相关的研究，这里的 prompt 参数和 meta-parameter 有着非常类似的思路。

通过引入 prompt，UIE 也能很方便的解决实体之间的关系抽取（Relation Extraction）任务，例如：

3.2 UIE 的实现

看完了基本思路，我们来一起看看 UIE 是怎么实现的吧。

模型部分

UIE 的模型代码比较简单，只需要在 encoder 后构建一个起始层和一个结束层即可：

class UIE(nn.Module):
    def __init__(self, encoder):        """        init func.        Args:            encoder (transformers.AutoModel): backbone, 默认使用 ernie 3.0                Reference:            https://github.com/PaddlePaddle/PaddleNLP/blob/a12481fc3039fb45ea2dfac3ea43365a07fc4921/model_zoo/uie/model.py        """        super().__init__()        self.encoder = encoder        hidden_size = 768        self.linear_start = nn.Linear(hidden_size, 1)        self.linear_end = nn.Linear(hidden_size, 1)        self.sigmoid = nn.Sigmoid()
    def forward(        self,        input_ids: torch.tensor,        token_type_ids: torch.tensor,        attention_mask=None,        pos_ids=None,    ) -> tuple:        """        forward 函数，返回开始/结束概率向量。        Args:            input_ids (torch.tensor): (batch, seq_len)            token_type_ids (torch.tensor): (batch, seq_len)            attention_mask (torch.tensor): (batch, seq_len)            pos_ids (torch.tensor): (batch, seq_len)        Returns:            tuple:  start_prob -> (batch, seq_len)                    end_prob -> (batch, seq_len)        """        sequence_output = self.encoder(            input_ids=input_ids,            token_type_ids=token_type_ids,            position_ids=pos_ids,            attention_mask=attention_mask,        )["last_hidden_state"]        start_logits = self.linear_start(sequence_output)       # (batch, seq_len, 1)        start_logits = torch.squeeze(start_logits, -1)          # (batch, seq_len)        start_prob = self.sigmoid(start_logits)                 # (batch, seq_len)        end_logits = self.linear_end(sequence_output)           # (batch, seq_len, 1)        end_logits = torch.squeeze(end_logits, -1)              # (batch, seq_len)        end_prob = self.sigmoid(end_logits)                     # (batch, seq_len)        return start_prob, end_prob

2. 训练部分

训练部分主要关注一下 loss 的计算即可。

由于每一个 token 都是一个二分类任务，因此选用 BCE Loss 作为损失函数。

分别计算起始/结束向量的 BCE Loss 再取平均值即可，如下所示：

criterion = torch.nn.BCELoss()...
start_prob, end_prob = model(input_ids=batch['input_ids'].to(args.device),                                        token_type_ids=batch['token_type_ids'].to(args.device),                                        attention_mask=batch['attention_mask'].to(args.device))start_ids = batch['start_ids'].to(torch.float32).to(args.device)    # (batch, seq_len)end_ids = batch['end_ids'].to(torch.float32).to(args.device)        # (batch, seq_len)loss_start = criterion(start_prob, start_ids)                       # 起止向量loss -> (1,)loss_end = criterion(end_prob, end_ids)                             # 结束向量loss -> (1,)loss = (loss_start + loss_end) / 2.0                                # 求平均 -> (1,)loss.backward()...

好啦，以上就是 UIE 的全部内容，感谢观看。

完整源码在这里：

UIE 复现源码
https://github.com/HarderThenHarder/transformers_tasks/tree/main/UIE

4. UIE 实战

4.1 UIE 环境安装

本项目基于 pytorch + transformers 实现，运行前请安装相关依赖包：

pip install -r ../requirements.txt

4.2 UIE 数据集准备

项目中提供了一部分示例数据，我们使用一个简单的ner任务（关系抽取同理）来进行信息抽取任务，数据在 data/simple_ner 。

若想使用自定义数据训练，只需要仿照示例数据构建数据集构建prompt和content即可：

{"content": "6月1日交通费68元", "result_list": [], "prompt": "出发地"}
{"content": "9月3日2点18分，加班打车回家，25元", "result_list": [{"text": "家", "start": 15, "end": 16}], "prompt": "目的地"}
{"content": "5月31号晚上10点54分打车回家49元", "result_list": [{"text": "5月31号晚上10点54分", "start": 0, "end": 13}], "prompt": "时间"}
...

Notes: 数据标注建议使用 doccano 完成，标注方法和标注转换可以参考 UIE 官方的详细介绍：

UIE 数据标注
https://github.com/PaddlePaddle/PaddleNLP/tree/develop/model_zoo/uie#数据标注

4.3 UIE 模型训练

修改训练脚本 train.sh 里的对应参数, 开启模型训练：

python train.py \    --pretrained_model "uie-base-zh" \    --save_dir "checkpoints/simple_ner" \    --train_path "data/simple_ner/train.txt" \    --dev_path "data/simple_ner/dev.txt" \    --img_log_dir "logs/simple_ner" \    --img_log_name "ERNIE-3.0" \    --batch_size 8 \    --max_seq_len 128 \    --num_train_epochs 100 \    --logging_steps 10 \    --valid_steps 100 \    --device cuda:0

正确开启训练后，终端会打印以下信息：

...global step 1880, epoch: 94, loss: 0.01507, speed: 10.06 step/sglobal step 1890, epoch: 95, loss: 0.01499, speed: 10.09 step/sglobal step 1900, epoch: 95, loss: 0.01492, speed: 10.05 step/sEvaluation precision: 0.94444, recall: 1.00000, F1: 0.97143best F1 performence has been updated: 0.94118 --> 0.97143global step 1910, epoch: 96, loss: 0.01484, speed: 10.19 step/s...

4.4 UIE 模型预测

完成模型训练后，运行 inference.py 以加载训练好的模型并应用：

 if __name__ == "__main__":    from rich import print    sentence = '5月17号晚上10点35分从公司加班打车回家，36块五。'        # NER 示例    ner_example(        sentence=sentence,         schema=['出发地', '目的地', '时间']    )
    # 事件抽取示例    event_extract_example(        sentence=sentence,         schema={                '加班触发词': ['时间','地点'],                '出行触发词': ['时间', '出发地', '目的地', '花费']            }    )

NER和事件抽取在schema的定义上存在一些区别：

NER的schema结构为 List 类型，列表中包含所有要提取的 实体类型。
事件的schema结构为 Dict 类型，其中 Key 的值是所有 事件触发词，Value 对应每一个触发词下的所有 事件属性。

python inference.py

得到以下推理结果：

[+] NER Results: {    '出发地': ['公司'],     '目的地': ['家'],     '时间': ['5月17号晚上10点35分']}
[+] Event-Extraction Results: {    '加班触发词': {},    '出行触发词': {        '时间': ['5月17号晚上10点35分', '公司'],        '出发地': ['公司'],        '目的地': ['公司', '家'],        '花费': ['36块五']    }}