简介本资源为基于Python的实体关系抽取系统源码包面向具备一定深度学习基础、从事知识图谱构建或自然语言处理方向的开发者与研究人员尤其适用于高端装备制造领域的知识图谱自动化构建任务。项目以BERT等预训练模型为核心实现从文本中抽取结构化实体与关系信息并引入FGM对抗性训练与EMA参数平滑更新策略以增强模型鲁棒性与训练稳定性。压缩包共22个文件包含8个Python脚本、12个JSON配置与数据文件、1个txt词表及1个md说明文档整体约1.97MB涵盖数据加载、数据集生成、模型定义、训练入口与预测等完整流程模块。目前已有208人学习下载。读者可据此获得一套可直接运行的实体关系抽取工程代码理解数据处理、模型训练、对抗训练与参数平滑等关键环节的实现方式并参考其目录组织与配置结构快速迁移至自有语料或类似抽取任务中。1. 从一份能跑通的 Python 实体关系抽取源码说起高端装备制造领域的知识图谱构建最头疼的环节往往不是图谱存储或查询而是从工艺文档、设备手册、故障报告里把「谁和谁有什么关系」抽出来。这份基于 Python 的实体关系抽取系统源码解决的就是这个环节——输入一段文本输出结构化的实体对及其关系类型。它用 BERT 做底座在gpNet.py里定义了关系分类模型main.py负责训练入口predict.py负责推理还带了 FGM 对抗训练和 EMA 参数平滑两个提分策略。适合谁正在做知识图谱自动化构建、需要一套可复现的实体关系抽取 baseline 的从业者以及想搞懂 BERT 微调在关系抽取任务上怎么落地的新手。源码包里数据、配置、预训练模型一应俱全拿到手改改数据格式就能跑自己的任务。2. 拆开源码包文件结构与数据流怎么走2.1 核心文件清单与职责划分拿到一个源码包我习惯先看目录结构判断作者的设计思路。这份包的文件不算多但分工明确典型的「数据生成 → 模型定义 → 训练 → 推理」四段式。文件职责关键点main.py训练入口组装数据、模型、优化器启动训练循环gpNet.py模型定义基于 BERT 的关系抽取网络含 FGM 和 EMAdataloader.py数据加载把 JSON 数据转成模型可吃的 tensordata_gen.py数据生成从原始标注生成训练格式data_genccl.py数据生成CCL 格式适配不同来源的数据predict.py推理脚本加载模型权重对新文本做关系抽取utils.py工具函数指标计算、日志、模型保存等vote.py投票集成多模型结果投票提升稳定性config.json模型配置BERT 结构参数schemas.json关系 schema定义关系类型和实体类型train.json/train2.json训练数据标注好的实体关系样本train_ccl.json/train_ccl2.jsonCCL 训练数据另一批训练样本evalA.json/evalA2.json验证数据用于评估模型效果submit_example_A.json提交示例推理输出格式参考pretrain_model/预训练模型BERT 权重和词表这个结构里schemas.json是容易被忽略但很关键的文件——它定义了你的任务到底要抽哪些关系类型。如果关系类型和你的业务对不上后面所有代码都得改。2.2 数据格式与 schema 定义实体关系抽取的数据格式常见的有两种一种是「句子 实体对 关系标签」的扁平结构另一种是「句子 实体列表 关系列表」的嵌套结构。这份源码用的是哪种从train.json和schemas.json的配合来看大概率是前者。先看schemas.json的结构。它通常长这样{ relation_types: [位于, 属于, 包含, 连接], entity_types: [设备, 部件, 工艺, 材料] }这个文件决定了模型输出层的维度——关系类型有几个分类头就输出几维。如果你要抽的关系和这里定义的不一样改这个文件是第一步。再看train.json的样本格式。典型的实体关系抽取训练样本是这样的{ text: 液压泵安装在主机架的上方, entities: [ {id: 0, type: 设备, start: 0, end: 3, text: 液压泵}, {id: 1, type: 部件, start: 6, end: 9, text: 主机架} ], relations: [ {head: 0, tail: 1, type: 位于} ] }text是原始句子entities列出所有实体及其位置relations用实体 id 索引关系。这种格式的好处是实体和关系解耦一个实体可以参与多个关系。注意如果你的标注数据是「实体对 关系」的 CSV 格式需要先转成这种 JSON 结构。data_gen.py里应该有转换逻辑但具体字段名可能和你的原始数据对不上得手动改。2.3 从原始数据到模型输入的完整链路数据从原始 JSON 到模型可用的 tensor中间经过dataloader.py的处理。这个文件的核心任务是把文本 tokenize、把实体位置映射到 token 位置、把关系标签转成 id。常见做法是继承 PyTorch 的Dataset类实现__getitem__和__len__。__getitem__里做几件事from torch.utils.data import Dataset from transformers import BertTokenizer class REDataset(Dataset): def __init__(self, data_path, tokenizer_path, max_len128): self.tokenizer BertTokenizer.from_pretrained(tokenizer_path) self.max_len max_len self.samples self.load_data(data_path) def load_data(self, path): import json with open(path, r, encodingutf-8) as f: return json.load(f) def __getitem__(self, idx): sample self.samples[idx] text sample[text] encoding self.tokenizer( text, max_lengthself.max_len, paddingmax_length, truncationTrue, return_tensorspt ) input_ids encoding[input_ids].squeeze(0) attention_mask encoding[attention_mask].squeeze(0) # 实体位置映射字符位置 - token 位置 head_pos self.char_to_token(sample[entities][0][start]) tail_pos self.char_to_token(sample[entities][1][start]) label self.relation2id[sample[relations][0][type]] return { input_ids: input_ids, attention_mask: attention_mask, head_pos: head_pos, tail_pos: tail_pos, label: label }这段代码的关键在char_to_token——BERT 的 tokenizer 会把一个词切成多个 subword字符位置和 token 位置不是一一对应的。如果映射错了模型看到的实体位置就是错的训练 loss 会降不下去。常见做法是用tokenizer.encode_plus的offset_mapping参数拿到每个 token 对应的字符区间再反查。max_len128是个经验值。关系抽取的句子通常不会太长128 够用。如果你的文本平均长度超过 100 个 token调到 256但显存占用会翻倍。3. 模型定义与训练BERT 微调 FGM EMA 三件套3.1 gpNet.py 里的网络结构gpNet.py是整个项目的核心。从文件名猜gp 可能是「general purpose」或某个作者缩写不影响理解。这个文件定义了一个基于 BERT 的关系抽取模型结构不复杂BERT 编码器 实体位置特征 分类头。典型实现是这样的import torch import torch.nn as nn from transformers import BertModel class RelationExtractionModel(nn.Module): def __init__(self, pretrain_path, num_relations, dropout0.1): super().__init__() self.bert BertModel.from_pretrained(pretrain_path) hidden_size self.bert.config.hidden_size # 实体位置嵌入用两个可学习的向量表示 head 和 tail self.head_embed nn.Embedding(1, hidden_size) self.tail_embed nn.Embedding(1, hidden_size) self.classifier nn.Sequential( nn.Linear(hidden_size * 3, hidden_size), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_size, num_relations) ) def forward(self, input_ids, attention_mask, head_pos, tail_pos): outputs self.bert(input_idsinput_ids, attention_maskattention_mask) sequence_output outputs.last_hidden_state # [batch, seq_len, hidden] # 取实体位置的向量 batch_size input_ids.size(0) head_vec sequence_output[torch.arange(batch_size), head_pos] tail_vec sequence_output[torch.arange(batch_size), tail_pos] cls_vec sequence_output[:, 0, :] # [CLS] 向量 # 拼接 [CLS] head tail concat torch.cat([cls_vec, head_vec, tail_vec], dim-1) logits self.classifier(concat) return logits这个结构有几个设计选择值得说。第一用[CLS]向量做句子级表示这是 BERT 做分类任务的标准做法。第二head 和 tail 的向量直接从sequence_output里按位置取没有额外的注意力机制——简单但有效。第三拼接方式是[CLS] head tail维度是hidden_size * 3分类头先降维再输出。如果你的关系类型很多超过 20 种可以考虑在分类头前加一层注意力让模型自己学习 head 和 tail 的交互。但关系类型少的时候直接拼接就够了加注意力反而容易过拟合。3.2 FGM 对抗训练怎么接进去FGMFast Gradient Method是一种对抗训练策略思路是在 embedding 层加一个扰动让模型对输入的小变化更鲁棒。这份源码里用 FGM 增强鲁棒性具体实现通常在gpNet.py或utils.py里。FGM 的核心代码不长class FGM: def __init__(self, model): self.model model self.backup {} def attack(self, epsilon1.0, emb_nameword_embeddings): for name, param in self.model.named_parameters(): if param.requires_grad and emb_name in name: self.backup[name] param.data.clone() norm torch.norm(param.grad) if norm ! 0: r_at epsilon * param.grad / norm param.data.add_(r_at) def restore(self, emb_nameword_embeddings): for name, param in self.model.named_parameters(): if param.requires_grad and emb_name in name: param.data self.backup[name] self.backup {}用法是在训练循环里正常反向传播后先attack()加扰动再算一次 loss 并反向传播最后restore()恢复参数。这样模型就见到了「被扰动过的自己」泛化能力会好一些。epsilon1.0是扰动幅度常见范围是 0.5 到 1.5。太大模型学不动太小没效果。emb_nameword_embeddings指定只扰动 BERT 的词嵌入层不碰其他参数。如果你的模型里 embedding 层名字不一样得改这个参数。注意FGM 会让训练时间增加约 50%因为多了一次前向和反向。如果训练资源紧张可以只在训练后期开启或者隔几个 batch 开一次。3.3 EMA 参数平滑更新的实现细节EMAExponential Moving Average是另一种提稳策略。思路是维护一份模型参数的滑动平均训练结束后用这份平均参数做推理而不是用最后一步的参数。实现上EMA 类通常长这样class EMA: def __init__(self, model, decay0.999): self.model model self.decay decay self.shadow {} self.backup {} for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] param.data.clone() def update(self): for name, param in self.model.named_parameters(): if param.requires_grad: new_average (1.0 - self.decay) * param.data self.decay * self.shadow[name] self.shadow[name] new_average.clone() def apply_shadow(self): for name, param in self.model.named_parameters(): if param.requires_grad: self.backup[name] param.data param.data self.shadow[name] def restore(self): for name, param in self.model.named_parameters(): if param.requires_grad: param.data self.backup[name] self.backup {}decay0.999是衰减率越大平均越平滑但需要更多步才能跟上模型变化。常见范围是 0.99 到 0.9999。训练步数少的时候用 0.99步数多用 0.999。用法是在每个 batch 更新完参数后调ema.update()训练结束后调ema.apply_shadow()再推理。这样推理用的参数是「历史平均」比单步参数更稳。3.4 main.py 训练流程与关键参数main.py把上面这些串起来。典型训练流程是import torch from torch.utils.data import DataLoader from transformers import AdamW, get_linear_schedule_with_warmup from gpNet import RelationExtractionModel, FGM, EMA from dataloader import REDataset def train(): # 1. 加载数据 train_dataset REDataset(train.json, pretrain_model) train_loader DataLoader(train_dataset, batch_size16, shuffleTrue) # 2. 初始化模型 model RelationExtractionModel(pretrain_model, num_relations4) model.cuda() # 3. 优化器和调度器 optimizer AdamW(model.parameters(), lr2e-5, weight_decay0.01) total_steps len(train_loader) * 5 # 5 个 epoch scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepsint(0.1 * total_steps), num_training_stepstotal_steps ) # 4. FGM 和 EMA fgm FGM(model) ema EMA(model, decay0.999) # 5. 训练循环 for epoch in range(5): model.train() for batch in train_loader: input_ids batch[input_ids].cuda() attention_mask batch[attention_mask].cuda() head_pos batch[head_pos].cuda() tail_pos batch[tail_pos].cuda() labels batch[label].cuda() # 正常前向 logits model(input_ids, attention_mask, head_pos, tail_pos) loss torch.nn.functional.cross_entropy(logits, labels) loss.backward() # FGM 对抗 fgm.attack() logits_adv model(input_ids, attention_mask, head_pos, tail_pos) loss_adv torch.nn.functional.cross_entropy(logits_adv, labels) loss_adv.backward() fgm.restore() # 更新参数 optimizer.step() scheduler.step() optimizer.zero_grad() ema.update() # 6. 应用 EMA 参数并保存 ema.apply_shadow() torch.save(model.state_dict(), model_ema.pt)几个关键参数lr2e-5是 BERT 微调的标准学习率太大容易崩太小收敛慢。batch_size16在 8G 显存上跑 128 长度没问题显存小就降到 8。warmup比例 0.1 是经验值让学习率在前 10% 步数里线性上升避免一开始就大步更新。weight_decay0.01是 AdamW 的默认正则化对 BERT 微调通常够用。如果过拟合严重可以加到 0.05。4. 推理与集成predict.py 和 vote.py 怎么配合4.1 predict.py 的推理流程训练完模型下一步是用predict.py对新文本做关系抽取。推理流程比训练简单加载模型权重、tokenize 输入、前向传播、取 argmax。import torch import json from transformers import BertTokenizer from gpNet import RelationExtractionModel def predict(text, model, tokenizer, relation_types, devicecuda): model.eval() encoding tokenizer( text, max_length128, paddingmax_length, truncationTrue, return_tensorspt ) input_ids encoding[input_ids].to(device) attention_mask encoding[attention_mask].to(device) # 这里假设实体位置已知实际场景可能需要先做实体识别 head_pos torch.tensor([0]).to(device) tail_pos torch.tensor([1]).to(device) with torch.no_grad(): logits model(input_ids, attention_mask, head_pos, tail_pos) pred torch.argmax(logits, dim-1).item() return relation_types[pred]这段代码有个前提实体位置是已知的。实际场景里实体识别和关系抽取通常是两个模型先跑 NER 拿到实体再把实体对喂给关系模型。这份源码可能假设实体已经标注好或者predict.py里集成了 NER 逻辑。submit_example_A.json是输出格式的参考。如果你的任务需要提交到某个平台照着这个格式生成结果就行。4.2 vote.py 多模型投票的适用场景vote.py做的是多模型结果投票。思路很简单训练多个模型不同随机种子、不同超参推理时让它们各自预测取多数票作为最终结果。from collections import Counter def vote_predictions(predictions_list): # predictions_list: [[pred1, pred2, ...], [pred1, pred2, ...], ...] final [] for preds in zip(*predictions_list): counter Counter(preds) final.append(counter.most_common(1)[0][0]) return final投票能提升稳定性但代价是推理时间翻倍。如果单模型效果已经不错投票的边际收益不大。常见做法是训练 3 到 5 个模型做投票再多就划不来了。注意投票要求所有模型的输出空间一致也就是schemas.json里的关系类型必须相同。如果不同模型用了不同的 schema投票没意义。4.3 评估指标与验证集使用evalA.json和evalA2.json是验证集。评估关系抽取效果常用指标是 Precision、Recall 和 F1。utils.py里应该有计算这些指标的函数。典型评估逻辑是def evaluate(model, eval_loader, relation_types): model.eval() all_preds [] all_labels [] with torch.no_grad(): for batch in eval_loader: logits model(batch[input_ids].cuda(), batch[attention_mask].cuda(), batch[head_pos].cuda(), batch[tail_pos].cuda()) preds torch.argmax(logits, dim-1).cpu().tolist() all_preds.extend(preds) all_labels.extend(batch[label].tolist()) # 计算 P/R/F1 from sklearn.metrics import classification_report report classification_report(all_labels, all_preds, target_namesrelation_types) return reportclassification_report会输出每个关系类型的 P/R/F1以及宏平均和加权平均。重点看宏平均 F1因为它对类别不平衡更敏感。如果某个关系类型的 F1 特别低可能是样本太少需要补充数据或做数据增强。5. 避坑与排查这份源码跑不起来时先查什么5.1 数据格式对不上导致 loss 不降现象训练了几个 epochloss 一直在 2.0 附近震荡不下降。原因最常见的是数据格式问题。train.json里的字段名和dataloader.py里读的字段名不一致或者实体位置是字符位置但代码按 token 位置处理。解决先打印一条__getitem__的输出看input_ids、head_pos、tail_pos、label是否合理。head_pos应该在 0 到max_len之间label应该在 0 到num_relations-1之间。如果head_pos是负数或超过max_len说明位置映射错了。5.2 显存不够导致 batch_size 必须调小现象跑main.py时报CUDA out of memory。原因batch_size16加max_len128在 8G 显存上刚好够但如果你的 BERT 是bert-base以上或者max_len设了 256显存就不够了。解决先把batch_size降到 8 或 4再考虑用梯度累积模拟大 batch。梯度累积的做法是每跑 N 个 batch 才optimizer.step()一次等效于batch_size * N。accumulation_steps 4 for i, batch in enumerate(train_loader): loss model(batch) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()5.3 FGM 扰动后 loss 变成 NaN现象开启 FGM 后训练几个 batch 就出现lossnan。原因FGM 的epsilon太大或者梯度本身有 NaN。epsilon1.0在大多数情况下安全但如果学习率太高梯度爆炸FGM 会把 NaN 放大。解决先把epsilon降到 0.5同时加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)梯度裁剪放在loss.backward()之后、optimizer.step()之前。max_norm1.0是常见值如果还 NaN降到 0.5。5.4 EMA 应用后效果反而变差现象用 EMA 参数推理F1 比不用 EMA 还低。原因decay太大EMA 参数跟不上模型变化相当于用了一个「过时」的模型。或者训练步数太少EMA 还没收敛。解决把decay从 0.999 降到 0.99或者只在训练后期最后 1 到 2 个 epoch开启 EMA。如果训练总步数少于 1000EMA 的收益不明显可以不用。5.5 推理时实体位置未知怎么办现象predict.py需要head_pos和tail_pos但实际场景里实体位置是未知的。原因这份源码假设实体已经标注好或者只做关系分类不做实体识别。如果你的场景需要端到端抽取得先跑一个 NER 模型。解决常见做法是用transformers的pipeline先做 NER拿到实体列表后枚举实体对逐个喂给关系模型。实体对数量是 N*(N-1)/2如果实体太多推理时间会很长。可以加一个启发式规则只对距离小于 50 个 token 的实体对做关系判断。6. 进阶技巧用 schemas.json 驱动的关系类型扩展与投票集成调优这份源码的扩展性很大程度上取决于schemas.json的设计。我拿到手第一件事是看这个文件因为它决定了整个系统的「输出空间」。如果你要加一个新的关系类型比如「维护周期」需要改三个地方schemas.json里加类型、train.json里加标注样本、gpNet.py里改num_relations。少改一个模型就学不会。一个容易被忽略的点是关系类型的粒度。太粗比如把所有「空间关系」都归为「位于」模型学起来容易但业务上不够用太细比如把「上方」「下方」「左侧」「右侧」分成四类样本不够模型学不动。我一般会先跑一版粗粒度看混淆矩阵如果某两类经常混再考虑合并或加更多样本。投票集成的调优关键在「模型多样性」。如果几个模型都是同一个随机种子、同一份数据、同一套超参投票没意义。常见做法是不同随机种子42、123、2024、不同学习率1e-5、2e-5、3e-5、不同dropout0.1、0.2、0.3。这样每个模型的错误模式不一样投票才能互补。# 多模型推理 投票的完整流程 import torch from gpNet import RelationExtractionModel from transformers import BertTokenizer def ensemble_predict(text, model_paths, relation_types): tokenizer BertTokenizer.from_pretrained(pretrain_model) all_preds [] for path in model_paths: model RelationExtractionModel(pretrain_model, len(relation_types)) model.load_state_dict(torch.load(path)) model.cuda().eval() # 假设实体位置已知 pred predict(text, model, tokenizer, relation_types) all_preds.append(pred) # 投票 from collections import Counter return Counter(all_preds).most_common(1)[0][0]model_paths是多个模型权重的路径列表。每个模型独立推理最后投票。如果某个模型在验证集上 F1 明显低可以给它降权或直接剔除。还有一个技巧是「阈值调优」。关系抽取模型输出的logits经过 softmax 后是概率默认取 argmax。但如果某个关系类型的 precision 低可以设一个阈值概率低于阈值就判为「无关系」。这个阈值在验证集上调通常能提升 1 到 2 个点的 F1。def predict_with_threshold(logits, relation_types, threshold0.5): probs torch.softmax(logits, dim-1) max_prob, pred torch.max(probs, dim-1) if max_prob.item() threshold: return 无关系 return relation_types[pred.item()]threshold0.5是起点实际值在验证集上扫一遍看哪个阈值 F1 最高。如果模型对「无关系」样本的识别差这个技巧效果明显。从那以后我每次拿到新的关系抽取源码都先跑通train.json上的 baseline再改schemas.json适配自己的任务最后才动模型结构。顺序反了出了问题都不知道是数据、schema 还是模型的锅。希望帮到你。本文还有配套的精品资源点击获取