用于关系抽取的 BERT(S)

概述

这是一个基于 PyTorch 的实现，用于 ACL 2019 年发表的论文《匹配空缺：用于关系学习的分布相似性》（https://arxiv.org/pdf/1906.03158.pdf）中的模型。
注意：这不是该论文的官方仓库。
在此基础上，我们还实现了以下基于该论文方法的关系抽取模型：

• ALBERT (https://arxiv.org/abs/1909.11942)
• BioBERT (https://arxiv.org/abs/1901.08746)

有关实现的更多概念性细节，请参阅：https://towardsdatascience.com/bert-s-for-relation-extraction-in-nlp-2c7c3ab487c4

如果您喜欢我的工作，请考虑通过点击顶部的赞助按钮来支持我。

需求

需求：Python (3.8+)

python3 -m pip install -r requirements.txt
python3 -m spacy download en_core_web_lg

预训练的 BERT 模型（ALBERT、BERT）由 HuggingFace.co (https://huggingface.co) 提供；
预训练的 BioBERT 模型由 https://github.com/dmis-lab/biobert 提供。

要使用 BioBERT (biobert_v1.1_pubmed)，请从 这里 下载并解压模型到 ./additional_models 文件夹中。

通过匹配空缺进行预训练（BERT_EM + MTB）

运行 main_pretraining.py，并传入以下参数。预训练数据可以是任何 .txt 格式的连续文本文件。
我们使用 Spacy NLP 从文本中提取成对实体（窗口大小为 40 个标记），以形成用于预训练的关系语句。实体识别基于 NER 和主语/宾语的依存句法分析。

我使用的 CNN 数据集（cnn.txt）预训练数据可在此下载：链接。
请下载并保存为 ./data/cnn.txt。
不过请注意，论文中用于 MTB 预训练的数据是维基百科转储文件，其规模远大于 CNN 数据集。

注意：预训练可能需要较长时间，具体取决于可用的 GPU。您也可以直接在关系抽取任务上进行微调，同样可以获得合理的结果，详见下文。

main_pretraining.py [-h] 
	[--pretrain_data TRAIN_PATH] 
	[--batch_size BATCH_SIZE]
	[--freeze FREEZE]  
	[--gradient_acc_steps GRADIENT_ACC_STEPS]
	[--max_norm MAX_NORM]
	[--fp16 FP_16]  
	[--num_epochs NUM_EPOCHS]
	[--lr LR]
	[--model_no MODEL_NO (0: BERT ; 1: ALBERT ; 2: BioBERT)]  
	[--model_size MODEL_SIZE (BERT: 'bert-base-uncased', 'bert-large-uncased';   
				ALBERT: 'albert-base-v2', 'albert-large-v2';   
				BioBERT: 'bert-base-uncased' (biobert_v1.1_pubmed))]

在 SemEval2010 Task 8 上进行微调（BERT_EM/BERT_EM + MTB）

运行 main_task.py，并传入以下参数。需要 SemEval2010 Task 8 数据集，可在此下载：链接。请下载并解压到 ./data/ 文件夹中。

main_task.py [-h] 
	[--train_data TRAIN_DATA]
	[--test_data TEST_DATA]
	[--use_pretrained_blanks USE_PRETRAINED_BLANKS]
	[--num_classes NUM_CLASSES] 
	[--batch_size BATCH_SIZE]
	[--gradient_acc_steps GRADIENT_ACC_STEPS]
	[--max_norm MAX_NORM]
	[--fp16 FP_16]  
	[--num_epochs NUM_EPOCHS]
	[--lr LR]
	[--model_no MODEL_NO (0: BERT ; 1: ALBERT ; 2: BioBERT)]  
	[--model_size MODEL_SIZE (BERT: 'bert-base-uncased', 'bert-large-uncased';   
				ALBERT: 'albert-base-v2', 'albert-large-v2';   
				BioBERT: 'bert-base-uncased' (biobert_v1.1_pubmed))]    
	[--train TRAIN]
	[--infer INFER]

推理 (--infer=1)

要对一句话进行推理，您可以将句子中感兴趣的实体1和实体2分别用标签 [E1] 和 [E2] 标注。
示例：

输入句子（输入 'quit' 或 'exit' 结束）：
The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor.

句子： The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor.
预测： Cause-Effect(e1,e2) 

from src.tasks.infer import infer_from_trained

inferer = infer_from_trained(args, detect_entities=False)
test = "The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor."
inferer.infer_sentence(test, detect_entities=False)

句子： The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor.
预测： Cause-Effect(e1,e2) 

该脚本还可以自动检测输入句子中的潜在实体，在这种情况下会推断出所有可能的关系组合：

inferer = infer_from_trained(args, detect_entities=True)
test2 = "After eating the chicken, he developed a sore throat the next morning."
inferer.infer_sentence(test2, detect_entities=True)

句子： [E2]After eating the chicken[/E2] , [E1]he[/E1] developed a sore throat the next morning .
预测： Other 

句子： After eating the chicken , [E1]he[/E1] developed [E2]a sore throat[/E2] the next morning .
预测： Other 

句子： [E1]After eating the chicken[/E1] , [E2]he[/E2] developed a sore throat the next morning .
预测： Other 

句子： [E1]After eating the chicken[/E1] , he developed [E2]a sore throat[/E2] the next morning .
预测： Other 

句子： After eating the chicken , [E2]he[/E2] developed [E1]a sore throat[/E1] the next morning .
预测： Other 

句子： [E2]After eating the chicken[/E2] , he developed [E1]a sore throat[/E1] the next morning .
预测： Cause-Effect(e2,e1) 

FewRel 任务

请在此下载 FewRel 1.0 数据集：链接，并解压到 ./data/ 文件夹中。
运行 main_task.py，并将参数 'task' 设置为 'fewrel'。

python main_task.py --task fewrel

结果如下：
(5-way 1-shot)
BERT_EM 未使用 MTB，未在任何 FewRel 数据上训练

模型尺寸	准确率（41646 个样本）
bert-base-uncased	62.229 %
bert-large-uncased	72.766 %

基准测试结果

SemEval2010 Task 8

1. 基础架构：BERT base uncased（12 层，768 隐藏层，12 头注意力，1.1 亿参数）

未使用 MTB 预训练：在 100% 训练数据上训练时的 F1 分数：

2. 基础架构：ALBERT base uncased（12 重复层，128 嵌入维度，768 隐藏层，12 头注意力，1100 万参数）

未使用 MTB 预训练：在 100% 训练数据上训练时的 F1 分数：

待补充

• 使用 MTB 预训练后在基准测试（SemEval2010 Task 8）上的推理与结果
• felrel 任务

BERT-Relation-Extraction 快速上手指南

本指南基于 ACL 2019 论文《Matching the Blanks: Distributional Similarity for Relation Learning》的 PyTorch 实现，支持 BERT、ALBERT 和 BioBERT 模型进行关系抽取任务。

环境准备

系统要求：

• Python 3.8 或更高版本
• GPU 环境（推荐，用于加速预训练和微调）
• Git

前置依赖：

• PyTorch
• HuggingFace Transformers (用于加载 BERT/ALBERT)
• Spacy (用于实体识别和依存句法分析)
• 其他依赖见 requirements.txt

安装步骤

1. 克隆仓库并安装依赖

   git clone https://github.com/plkmo/BERT-Relation-Extraction.git
   cd BERT-Relation-Extraction
   python3 -m pip install -r requirements.txt
   

   提示：国内用户建议使用清华或阿里镜像源加速安装： python3 -m pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 下载 Spacy 英文模型

   python3 -m spacy download en_core_web_lg
   

   提示：若下载缓慢，可手动下载模型文件后离线安装，或使用国内镜像配置。

3. 准备预训练模型（可选）

   • BERT / ALBERT：代码运行时会自动从 HuggingFace 下载。
   • BioBERT：需手动下载。
     1. 访问 BioBERT GitHub 下载 biobert_v1.1_pubmed 模型。
     2. 解压并将文件夹放入项目根目录下的 ./additional_models 文件夹中。

4. 准备数据集

   • SemEval2010 Task 8 (微调必备)： 下载 数据集链接，解压至 ./data/ 目录。
   • 预训练数据 (可选)： 如需复现论文的 MTB 预训练，可下载 CNN 数据集保存为 ./data/cnn.txt，或使用更大的 Wiki dumps 数据。

基本使用

1. 直接推理（无需训练）

如果你已有训练好的模型，可以直接对句子进行关系推断。实体需用 [E1]...[/E1] 和 [E2]...[/E2] 标签标记。

命令行交互模式：

python main_task.py --infer 1 --model_no 0 --model_size 'bert-base-uncased'

输入示例：

Type input sentence ('quit' or 'exit' to terminate):
The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor.

输出示例：

Sentence:  The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor.
Predicted:  Cause-Effect(e1,e2) 

Python 代码调用模式：

from src.tasks.infer import infer_from_trained
import argparse

# 模拟参数配置 (实际使用时请加载你的 args)
args = argparse.Namespace()
args.model_no = 0
args.model_size = 'bert-base-uncased'
# ... 其他必要参数

# 初始化推断器 (detect_entities=False 表示实体已标记)
inferer = infer_from_trained(args, detect_entities=False)

test_sentence = "The surprise [E1]visit[/E1] caused a [E2]frenzy[/E2] on the already chaotic trading floor."
result = inferer.infer_sentence(test_sentence, detect_entities=False)
print(result)

自动实体检测模式： 若未标记实体，可设置 detect_entities=True，系统将自动尝试识别所有可能的实体组合并预测关系：

inferer = infer_from_trained(args, detect_entities=True)
test_sentence = "After eating the chicken, he developed a sore throat the next morning."
# 系统将输出多种实体组合的预测结果
inferer.infer_sentence(test_sentence, detect_entities=True)

2. 微调模型 (Fine-tuning)

使用 SemEval2010 Task 8 数据集微调模型：

python main_task.py \
    --train_data ./data/train.txt \
    --test_data ./data/test.txt \
    --model_no 0 \
    --model_size 'bert-base-uncased' \
    --num_epochs 10 \
    --lr 2e-5 \
    --batch_size 16 \
    --train 1

3. MTB 预训练 (进阶)

若需从头进行“匹配空白”(Matching the Blanks) 预训练：

python main_pretraining.py \
    --pretrain_data ./data/cnn.txt \
    --model_no 0 \
    --model_size 'bert-base-uncased' \
    --num_epochs 5 \
    --batch_size 32

注意：预训练耗时较长，建议仅在拥有充足 GPU 资源且需要特定领域适配时进行。一般情况下，直接使用预训练权重进行微调即可获得良好效果。