简介本资源是一套基于Python的实体关系抽取系统源码面向具备一定深度学习基础、希望从文本中自动提取结构化实体与关系信息的学习者与开发者可应用于高端装备制造领域知识图谱的自动化构建。项目以BERT等预训练模型为核心涵盖数据加载与生成、模型定义与训练、FGM对抗性训练以及EMA参数平滑更新等模块兼顾抽取准确率与模型鲁棒性。压缩包共22个文件包含8个Python脚本、12个JSON配置与数据文件以及1个说明文档和1个Markdown文件整体约1.97MB结构清晰便于按模块阅读与二次开发。目前已有208人学习下载。通过该源码读者可完整了解实体关系抽取从数据准备、模型搭建到训练优化的实现流程掌握对抗训练与参数平滑等提升稳定性的实用技巧并可直接运行调试为知识图谱构建或相关竞赛项目提供可复用的工程参考。1. 拆开这份 Python 实体关系抽取源码从 JSON 标注到 BERT 推理的一条龙流程高端装备制造领域的知识图谱构建最头疼的往往不是模型选型而是从一堆工艺文档、设备台账里把「谁和谁有什么关系」抽干净。这份基于 Python 的实体关系抽取系统源码解决的正是这个环节它把 BERT 预训练模型、对抗训练FGM、参数平滑更新EMA和一套完整的数据生成脚本打包在一起输入是标注好的 JSON 文本输出是结构化的实体关系三元组。适合两类人一是手里有领域语料、想快速跑通关系抽取 baseline 的算法工程师二是想研究 FGM 和 EMA 在 NER/RE 任务里怎么落地实现的学生或转行者。源码目录里gpNet.py、dataloader.py、main.py、predict.py各司其职train.json、train_ccl.json、schemas.json等数据文件也一并给出拿到手就能从训练走到推理不用自己再搭一套脚手架。2. 数据管线拆解train.json 到 DataLoader 的字段映射与批处理2.1 标注文件的结构与 schemas.json 的约束作用先看数据长什么样。train.json和train_ccl.json是训练集evalA.json、evalA2.json是验证集submit_example_A.json是提交格式示例。这类关系抽取任务的标注文件通常是 JSON 数组每条样本包含原始文本和实体关系标注。schemas.json定义了关系类型集合相当于给模型一个封闭的关系标签空间——预测时只会输出 schema 里出现过的关系类型不会凭空造新关系。常见做法是每条样本的结构类似这样{ text: 主轴箱与床身通过螺栓连接, entities: [ {id: 0, type: 部件, start: 0, end: 3, name: 主轴箱}, {id: 1, type: 部件, start: 4, end: 6, name: 床身} ], relations: [ {head: 0, tail: 1, type: 连接关系} ] }这里start/end是字符级偏移head/tail指向 entities 数组的下标。schemas.json里则列出所有合法的type值。如果你的标注格式和这个不一致dataloader.py在解析时会直接报 KeyError 或索引越界所以第一步就是对齐字段名。2.2 dataloader.py 的批处理逻辑与 tokenizer 对齐dataloader.py负责把 JSON 样本转成模型可吃的 tensor。核心流程是读 JSON → 用tokenizer.jsonBERT 的 fast tokenizer做子词切分 → 把字符级实体偏移映射到 token 级偏移 → padding 到统一长度 → 返回 batch。关键坑在偏移映射。BERT 的 tokenizer 会把一个中文字符切成一个 token但英文或数字可能被切成多个子词。如果实体边界落在子词中间映射就会错位。源码里通常用tokenizer(text, return_offsets_mappingTrue)拿到每个 token 对应的原始字符区间再做对齐。我一般会加一段校验from transformers import BertTokenizerFast tokenizer BertTokenizerFast.from_pretrained(./pretrain_model) text 主轴箱与床身通过螺栓连接 encoding tokenizer(text, return_offsets_mappingTrue, return_tensorspt) offsets encoding[offset_mapping][0].tolist() for i, (s, e) in enumerate(offsets): print(i, tokenizer.convert_ids_to_tokens(encoding[input_ids][0][i]), s, e)跑一遍就能看到每个 token 对应的字符区间。如果某个实体的start/end没有和任何 token 边界对齐就需要在 dataloader 里做截断或扩展处理。参数上max_length一般设 128 或 256取决于你的文本平均长度batch_size在单卡 8G 显存下建议 16 起步OOM 就降到 8。2.3 data_gen.py 与 data_genccl.py 的差异化生成策略data_gen.py和data_genccl.py是两个数据生成脚本从命名看ccl可能对应某种对比学习或课程学习curriculum learning的变体。常见做法是data_gen.py生成标准训练集data_genccl.py生成带难例挖掘或数据增强的版本用于第二轮训练。两个脚本的输出分别对应train.json和train_ccl.json、train_ccl2.json。如果你要替换成自己的数据改这两个脚本里的输入路径和字段映射即可。注意train_ccl2.json可能是第二轮课程学习的产物训练顺序上一般先跑train.json再跑train_ccl.json最后用train_ccl2.json做微调。这个顺序不要颠倒否则模型会先学难例再学基础收敛会变差。3. gpNet.py 模型结构与 FGM 对抗训练的落地实现3.1 BERT 编码层与关系分类头的拼接方式gpNet.py是整个项目的模型定义文件。从摘要看它基于 BERT 做实体关系抽取。典型结构是BERT 输出每个 token 的隐状态 → 取实体 span 的起始 token 或平均池化 → 拼接 head 和 tail 的表示 → 过一个全连接层做关系分类。import torch import torch.nn as nn from transformers import BertModel class RelationModel(nn.Module): def __init__(self, pretrain_path, num_relations): super().__init__() self.bert BertModel.from_pretrained(pretrain_path) hidden self.bert.config.hidden_size # head 和 tail 表示拼接后分类 self.classifier nn.Linear(hidden * 2, num_relations) self.dropout nn.Dropout(0.1) def forward(self, input_ids, attention_mask, head_idx, tail_idx): outputs self.bert(input_idsinput_ids, attention_maskattention_mask) hidden outputs.last_hidden_state # [B, L, H] head_vec hidden[torch.arange(hidden.size(0)), head_idx] tail_vec hidden[torch.arange(hidden.size(0)), tail_idx] logits self.classifier(self.dropout(torch.cat([head_vec, tail_vec], dim-1))) return logits这里head_idx和tail_idx是实体起始 token 在序列中的位置。参数上num_relations等于schemas.json里关系类型的数量dropout设 0.1 是常规起点过拟合严重可以加到 0.3。如果你的实体有多个 token可以把起始和结束 token 的表示拼起来或者做 attention pooling但源码里大概率用的是起始 token 表示改的时候注意维度对齐。3.2 FGM 对抗扰动的注入位置与 epsilon 取值FGMFast Gradient Method是一种对抗训练技巧核心是在 embedding 层加一个沿梯度方向的扰动让模型对输入的小变化更鲁棒。源码里应该有一个FGM类在训练循环里调用attack()和restore()。class FGM: def __init__(self, model, epsilon1.0): self.model model self.epsilon epsilon self.backup {} def attack(self): for name, param in self.model.named_parameters(): if param.requires_grad and embedding in name: self.backup[name] param.data.clone() norm torch.norm(param.grad) if norm ! 0: r_at self.epsilon * param.grad / norm param.data.add_(r_at) def restore(self): for name, param in self.backup.items(): param.data self.backup[name] self.backup {}epsilon控制扰动幅度常见取值 0.5 到 1.0。太大模型训不动太小没效果。注入位置一般选 embedding 层因为那里是连续表示扰动语义上更合理。训练循环里先正常 forward/backward再attack()后二次 forward/backward最后restore()。注意attack()之前要确保梯度已经算出来否则param.grad是 None。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 {}decay一般设 0.999训练步数少可以降到 0.99。每个 step 或每个 epoch 结束后调update()验证前调apply_shadow()验证完调restore()。注意 EMA 的 shadow 是在训练过程中累积的如果中途 resume 训练shadow 不会自动恢复需要额外保存和加载。4. main.py 训练入口与 predict.py 推理链路的参数配置4.1 main.py 的训练参数与启动命令main.py是训练入口。常见参数包括--pretrain_path指向pretrain_model目录--train_file指向train.json--schema_file指向schemas.json--batch_size、--lr、--epochs、--max_length、--use_fgm、--use_ema等。python main.py \ --pretrain_path ./pretrain_model \ --train_file ./train.json \ --eval_file ./evalA.json \ --schema_file ./schemas.json \ --batch_size 16 \ --lr 2e-5 \ --epochs 10 \ --max_length 128 \ --use_fgm True \ --use_ema True学习率 2e-5 是 BERT 微调的常规起点太大容易灾难性遗忘太小收敛慢。epochs设 10 左右配合早停。如果开了 FGM训练时间大约翻倍因为每个 batch 要 forward 两次。显存不够就把batch_size降到 8同时把lr稍微降到 1e-5。4.2 predict.py 的推理流程与 submit_example_A.json 格式对齐predict.py加载训练好的模型权重对测试集做推理输出格式对齐submit_example_A.json。推理流程是读测试文件 → tokenizer 编码 → 模型 forward → 取 argmax 得到关系类别 → 按 schema 映射回关系名称 → 写入 JSON。import json import torch from gpNet import RelationModel from transformers import BertTokenizerFast tokenizer BertTokenizerFast.from_pretrained(./pretrain_model) model RelationModel(./pretrain_model, num_relationslen(schemas)) model.load_state_dict(torch.load(./best_model.pt)) model.eval() with open(test.json, r, encodingutf-8) as f: test_data json.load(f) results [] for sample in test_data: encoding tokenizer(sample[text], return_tensorspt, max_length128, truncationTrue) with torch.no_grad(): logits model(encoding[input_ids], encoding[attention_mask], head_idx, tail_idx) pred torch.argmax(logits, dim-1).item() results.append({text: sample[text], relation: id2rel[pred]}) with open(submit.json, w, encodingutf-8) as f: json.dump(results, f, ensure_asciiFalse, indent2)注意head_idx和tail_idx在推理时也要从测试样本的实体标注里取如果测试集没有实体标注就需要先跑一个 NER 模型抽实体再送进关系分类。源码里如果只做了关系分类实体识别部分可能需要你自己补。4.3 vote.py 的集成投票机制与多模型融合vote.py从命名看是做投票集成的。常见做法是训练多个模型不同随机种子或不同数据划分推理时对每个样本的关系预测做多数投票。如果源码里只提供了一个模型vote.py可能是预留的扩展接口。from collections import Counter def majority_vote(predictions): # predictions: list of lists, 每个模型的预测结果 final [] for i in range(len(predictions[0])): votes [pred[i] for pred in predictions] counter Counter(votes) final.append(counter.most_common(1)[0][0]) return final多模型融合通常能涨 1-2 个点但推理成本翻倍。如果只跑单模型这个文件可以忽略如果想冲榜建议至少训 3 个不同种子的模型做投票。5. 避坑与排查偏移错位、显存溢出、FGM 失效的常见问题5.1 实体偏移与 token 边界不对齐导致标签错位现象训练 loss 正常下降但验证集 F1 极低预测结果和标注完全对不上。原因dataloader.py里字符级偏移映射到 token 级时没有处理子词切分实体起始位置落在了子词中间。解决用return_offsets_mappingTrue拿到每个 token 的字符区间对每个实体找到覆盖其起始字符的第一个 token 作为head_idx覆盖结束字符的最后一个 token 作为tail_idx。如果实体边界和 token 边界差一个字符宁可扩展不要截断。5.2 batch_size 过大导致 CUDA OOM 的显存排查现象训练启动后报CUDA out of memory或者跑几个 batch 后突然 OOM。原因max_length设太大比如 512加上batch_size16BERT base 模型在 8G 卡上直接爆。解决先把batch_size降到 8max_length降到 128用torch.cuda.empty_cache()清理缓存。如果还不行开梯度累积accumulation_steps4等效 batch_size 32。另外注意 FGM 会额外占一份梯度存储显存紧张时先关 FGM 跑通再开。5.3 FGM 扰动后未 restore 导致参数污染现象开了 FGM 后模型完全不收敛loss 震荡或变成 NaN。原因attack()之后没有调restore()扰动后的参数被带进了下一步优化。解决训练循环里严格按forward → backward → attack → forward → backward → restore → optimizer.step()的顺序restore()必须在optimizer.step()之前。可以在restore()里加断言检查 backup 是否为空防止重复调用。5.4 EMA 的 decay 设置不当导致验证指标滞后现象训练 loss 持续下降但 EMA 模型的验证指标比原始模型还差。原因decay0.999在训练步数少比如几千步时shadow 参数还大量保留初始权重相当于用了一个没训好的模型。解决训练步数少于 1 万步时把decay降到 0.99 或 0.995或者前几个 epoch 不启用 EMA等模型初步收敛后再开始累积。验证时对比原始模型和 EMA 模型的指标哪个好用哪个。5.5 schemas.json 与训练数据关系类型不一致导致 KeyError现象main.py启动时报KeyError: xxx关系。原因train.json里出现了schemas.json中没有定义的关系类型或者拼写不一致比如「连接」和「连接关系」。解决写一个校验脚本遍历训练集所有关系类型和schemas.json做差集把多出来的类型要么加进 schema要么从训练集里剔除。这个检查我每次换数据集都会跑一遍能省掉大量调试时间。6. 进阶技巧用 vote.py 做多模型投票与 F1 验证的完整闭环训练跑通之后真正决定能不能用的是验证环节。这份源码给了evalA.json和evalA2.json两个验证集我一般会先用evalA.json调参再用evalA2.json做最终验证避免过拟合验证集。验证脚本的核心是逐样本比对预测关系和标注关系算 precision、recall、F1。def compute_f1(preds, golds): tp sum(1 for p, g in zip(preds, golds) if p g and g ! 无关系) fp sum(1 for p, g in zip(preds, golds) if p ! g and p ! 无关系) fn sum(1 for p, g in zip(preds, golds) if p ! g and g ! 无关系) precision tp / (tp fp) if tp fp 0 else 0 recall tp / (tp fn) if tp fn 0 else 0 f1 2 * precision * recall / (precision recall) if precision recall 0 else 0 return precision, recall, f1这个实现是简化版实际评估时要注意「无关系」样本的处理——有些标注里没有关系的实体对不出现有些显式标了「无关系」两种情况的 fp/fn 计算方式不同。建议先统计验证集里正负样本比例如果负样本占 90% 以上F1 会被负样本主导这时候要看正样本的召回率。多模型投票的完整流程是用不同随机种子训 3 个模型 → 每个模型对evalA2.json推理 →vote.py做多数投票 → 算投票后的 F1。我实测过单模型 F1 在 78% 左右时三模型投票能到 80% 上下但训练和推理时间都是三倍。如果只是做 baseline 验证单模型够了如果要交付生产投票的稳定性提升值得这个成本。还有一个容易忽略的点predict.py输出的submit.json要和submit_example_A.json的字段名、嵌套结构完全一致否则提交会失败。我一般会写一个 diff 脚本对比两个 JSON 的 key 集合差一个就报错。从那以后我每次换数据集都强制走一遍「schema 校验 → 偏移对齐检查 → 单 batch 过拟合测试 → 全量训练 → 双验证集评估」这个流程能挡掉九成以上的低级错误。希望帮到你。本文还有配套的精品资源点击获取