AutoRE

本仓库基于LLaMA-Factory代码，实现了基于大语言模型的文档级关系抽取系统AutoRE。使用的抽取范式为RHF（论文链接）。目前基于Re-DocRED数据集进行实验，能够抽取文档级文本中的96个关系的三元组事实。

使用方法

从 huggingface上下载dante123/AutoRE。单个LoRA模块在570M左右。

0.环境准备

cd AutoRE/
pip install -r requirement.txt

因为使用了wandb，所以需要先将train_bash.py中的key进行设置

api_key = os.environ.get('WANDB_API_KEY', "your api key")

1.推理

# 根据AutoRE.sh内的提示进行修改
bash AutoRE.sh
# 输入对应文档即可自动抽取

2.模型训练

1）数据准备

cd AutoRE/utils/
python pre_process_data.py

2）微调模型

我们的代码参考自LLaMA-Factory，并进行了适当修改。

cd AutoRE/
# 选择对应的模型进行微调
# 可以指定单卡或者多卡
bash train_script/mistral_loras_D_R_H_F_desc.sh

3.模型测试

cd AutoRE/
# 选择对应的模型进行测试，数据集为Re-DocRED，将--inference去除，并且设置具体的模型和ckpt
bash AutoRE.sh

AutoRE_analysis

验证analysis过程是否对抽取是有帮助的。整个过程思路与AutoRE的框架一致，只是在每一步抽取前加入了analysis。具体可以看redocred_train_analysis.json中的例子。数据和代码已经分享，希望对大家能有些许启发～

另外，为了使AutoRE能够做更多类的关系抽取，加入其他的开源数据，包括英文的fewrel，nyt等，以及中文的hacred等。如果只关注本论文的工作，只需要将数据处理中pre_process_data.py的其他代码注释掉，只保留处理redocred的处理部分（代码中给了很多的注释，希望能帮到你们～）

引用

如果你觉得我们的工作有帮助的话，请考虑引用论文。

@article{lilong2024autore,
  title={AutoRE: Document-Level Relation Extraction with Large Language Models},
  author={Lilong, Xue and Dan, Zhang and Yuxiao, Dong and Jie, Tang},
  journal={arXiv preprint arXiv:2403.14888},
  year={2024}
}

THUDM / AutoRE