HUSTAI / uie_pytorch

PaddleNLP UIE模型的PyTorch版实现
Apache License 2.0
539 stars 93 forks source link
information-extraction pytorch transformers uie

通用信息抽取 UIE(Universal Information Extraction) PyTorch版

迁移PaddleNLP中的UIE模型到PyTorch上

PyTorch版功能介绍

目录

1. 模型简介

UIE(Universal Information Extraction):Yaojie Lu等人在ACL-2022中提出了通用信息抽取统一框架UIE。该框架实现了实体抽取、关系抽取、事件抽取、情感分析等任务的统一建模,并使得不同任务间具备良好的迁移和泛化能力。为了方便大家使用UIE的强大能力,PaddleNLP借鉴该论文的方法,基于ERNIE 3.0知识增强预训练模型,训练并开源了首个中文通用信息抽取模型UIE。该模型可以支持不限定行业领域和抽取目标的关键信息抽取,实现零样本快速冷启动,并具备优秀的小样本微调能力,快速适配特定的抽取目标。

UIE的优势

2. 应用示例

UIE不限定行业领域和抽取目标,以下是一些零样本行业示例:

image

image

image

image

image

3. 开箱即用

uie_predictor提供通用信息抽取、评价观点抽取等能力,可抽取多种类型的信息,包括但不限于命名实体识别(如人名、地名、机构名等)、关系(如电影的导演、歌曲的发行时间等)、事件(如某路口发生车祸、某地发生地震等)、以及评价维度、观点词、情感倾向等信息。用户可以使用自然语言自定义抽取目标,无需训练即可统一抽取输入文本中的对应信息。实现开箱即用,并满足各类信息抽取需求

uie_predictor现在可以自动下载模型了,无需手动convert,如果想手动转换模型,可以参照以下方法。

下载并转换模型,将下载Paddle版的uie-base模型到当前目录中,并生成PyTorch版模型uie_base_pytorch

python convert.py

如果没有安装paddlenlp,则使用以下命令。这将不会导入paddlenlp,以及不会验证转换结果正确性。

python convert.py --no_validate_output

可配置参数说明:

* : 使用ernie-3.0-base-zh时不会验证模型,需要微调后才能用于预测

3.1 实体抽取

命名实体识别(Named Entity Recognition,简称NER),是指识别文本中具有特定意义的实体。在开放域信息抽取中,抽取的类别没有限制,用户可以自己定义。

3.2 关系抽取

关系抽取(Relation Extraction,简称RE),是指从文本中识别实体并抽取实体之间的语义关系,进而获取三元组信息,即<主体,谓语,客体>。

3.3 事件抽取

事件抽取 (Event Extraction, 简称EE),是指从自然语言文本中抽取预定义的事件触发词(Trigger)和事件论元(Argument),组合为相应的事件结构化信息。

3.4 评论观点抽取

评论观点抽取,是指抽取文本中包含的评价维度、观点词。

3.5 情感分类

3.6 跨任务抽取

3.7 模型选择

3.8 更多配置

>>> from uie_predictor import UIEPredictor

>>> ie = UIEPredictor('uie_nano',   
                       schema=schema)  

4. 训练定制

对于简单的抽取目标可以直接使用UIEPredictor实现零样本(zero-shot)抽取,对于细分场景我们推荐使用轻定制功能(标注少量数据进行模型微调)以进一步提升效果。下面通过报销工单信息抽取的例子展示如何通过5条训练数据进行UIE模型微调。

4.1 代码结构

.
├── utils.py          # 数据处理工具
├── model.py          # 模型组网脚本
├── doccano.py        # 数据标注脚本
├── doccano.md        # 数据标注文档
├── finetune.py       # 模型微调脚本
├── evaluate.py       # 模型评估脚本
└── README.md

4.2 数据标注

我们推荐使用数据标注平台doccano 进行数据标注,本示例也打通了从标注到训练的通道,即doccano导出数据后可通过doccano.py脚本轻松将数据转换为输入模型时需要的形式,实现无缝衔接。标注方法的详细介绍请参考doccano数据标注指南

原始数据示例:

深大到双龙28块钱4月24号交通费

抽取的目标(schema)为:

schema = ['出发地', '目的地', '费用', '时间']

标注步骤如下:

python doccano.py \
    --doccano_file ./data/doccano_ext.json \
    --task_type ext \
    --save_dir ./data \
    --splits 0.8 0.2 0

可配置参数说明:

备注:

更多不同类型任务(关系抽取、事件抽取、评价观点抽取等)的标注规则及参数说明,请参考doccano数据标注指南

此外,也可以通过数据标注平台 Label Studio 进行数据标注。本示例提供了 labelstudio2doccano.py 脚本,将 label studio 导出的 JSON 数据文件格式转换成 doccano 导出的数据文件格式,后续的数据转换与模型微调等操作不变。

python labelstudio2doccano.py --labelstudio_file label-studio.json

可配置参数说明:

4.3 模型微调

通过运行以下命令进行模型微调:

python finetune.py \
    --train_path "./data/train.txt" \
    --dev_path "./data/dev.txt" \
    --save_dir "./checkpoint" \
    --learning_rate 1e-5 \
    --batch_size 16 \
    --max_seq_len 512 \
    --num_epochs 100 \
    --model "uie_base_pytorch" \
    --seed 1000 \
    --logging_steps 10 \
    --valid_steps 100 \
    --device "gpu"

可配置参数说明:

4.4 模型评估

通过运行以下命令进行模型评估:

python evaluate.py \
    --model_path ./checkpoint/model_best \
    --test_path ./data/dev.txt \
    --batch_size 16 \
    --max_seq_len 512

评估方式说明:采用单阶段评价的方式,即关系抽取、事件抽取等需要分阶段预测的任务对每一阶段的预测结果进行分别评价。验证/测试集默认会利用同一层级的所有标签来构造出全部负例。

可开启debug模式对每个正例类别分别进行评估,该模式仅用于模型调试:

python evaluate.py \
    --model_path ./checkpoint/model_best \
    --test_path ./data/dev.txt \
    --debug

输出打印示例:

[2022-09-14 03:13:58,877] [    INFO] - -----------------------------
[2022-09-14 03:13:58,877] [    INFO] - Class Name: 疾病
[2022-09-14 03:13:58,877] [    INFO] - Evaluation Precision: 0.89744 | Recall: 0.83333 | F1: 0.86420
[2022-09-14 03:13:59,145] [    INFO] - -----------------------------
[2022-09-14 03:13:59,145] [    INFO] - Class Name: 手术治疗
[2022-09-14 03:13:59,145] [    INFO] - Evaluation Precision: 0.90000 | Recall: 0.85714 | F1: 0.87805
[2022-09-14 03:13:59,439] [    INFO] - -----------------------------
[2022-09-14 03:13:59,440] [    INFO] - Class Name: 检查
[2022-09-14 03:13:59,440] [    INFO] - Evaluation Precision: 0.77778 | Recall: 0.56757 | F1: 0.65625
[2022-09-14 03:13:59,708] [    INFO] - -----------------------------
[2022-09-14 03:13:59,709] [    INFO] - Class Name: X的手术治疗
[2022-09-14 03:13:59,709] [    INFO] - Evaluation Precision: 0.90000 | Recall: 0.85714 | F1: 0.87805
[2022-09-14 03:13:59,893] [    INFO] - -----------------------------
[2022-09-14 03:13:59,893] [    INFO] - Class Name: X的实验室检查
[2022-09-14 03:13:59,894] [    INFO] - Evaluation Precision: 0.71429 | Recall: 0.55556 | F1: 0.62500
[2022-09-14 03:14:00,057] [    INFO] - -----------------------------
[2022-09-14 03:14:00,058] [    INFO] - Class Name: X的影像学检查
[2022-09-14 03:14:00,058] [    INFO] - Evaluation Precision: 0.69231 | Recall: 0.45000 | F1: 0.54545

可配置参数说明:

4.5 定制模型一键预测

UIEPredictor装载定制模型,通过task_path指定模型权重文件的路径,路径下需要包含训练好的模型权重文件pytorch_model.bin

>>> from pprint import pprint
>>> from uie_predictor import UIEPredictor

>>> schema = ['出发地', '目的地', '费用', '时间']
# 设定抽取目标和定制化模型权重路径
>>> my_ie = UIEPredictor(model='uie-base',task_path='./checkpoint/model_best', schema=schema)
>>> pprint(my_ie("城市内交通费7月5日金额114广州至佛山"))
[{'出发地': [{'end': 17,
           'probability': 0.9975287467835301,
           'start': 15,
           'text': '广州'}],
  '时间': [{'end': 10,
          'probability': 0.9999476678061399,
          'start': 6,
          'text': '7月5日'}],
  '目的地': [{'end': 20,
           'probability': 0.9998511131226735,
           'start': 18,
           'text': '佛山'}],
  '费用': [{'end': 15,
          'probability': 0.9994474579292856,
          'start': 12,
          'text': '114'}]}]

4.6 实验指标

我们在互联网、医疗、金融三大垂类自建测试集上进行了实验:

金融医疗互联网
0-shot5-shot0-shot5-shot0-shot5-shot
uie-base (12L768H)46.4370.9271.8385.7278.3381.86
uie-medium (6L768H)41.1164.5365.4075.7278.3279.68
uie-mini (6L384H)37.0464.6560.5078.3672.0976.38
uie-micro (4L384H)37.5362.1157.0475.9266.0070.22
uie-nano (4L312H)38.9466.8348.2976.7462.8672.35
uie-m-large (24L1024H)49.3574.5570.5092.6678.4983.02
uie-m-base (12L768H)38.4674.3163.3787.3276.2780.13

0-shot表示无训练数据直接通过UIEPredictor进行预测,5-shot表示每个类别包含5条标注数据进行模型微调。实验表明UIE在垂类场景可以通过少量数据(few-shot)进一步提升效果

4.7 模型部署

以下是UIE Python端的部署流程,包括环境准备、模型导出和使用示例。