首页 / 资讯中心 / 文章详情

微调BERT实现抽取式摘要:完整代码工程解析与实战避坑指南

微调BERT实现抽取式摘要:完整代码工程解析与实战避坑指南 ★ FEATURED ARTICLE
简介这是一份基于BERT微调实现抽取式摘要生成的完整论文代码适合有一定深度学习基础、希望掌握Transformer微调与序列到序列建模的自然语言处理开发者。压缩包共36个文件包含20个Python脚本、7个文本文件及若干配置与说明文档分别承担核心逻辑、数据映射、仓库配置与使用指南等作用整体大小约14.99MB目录结构清晰便于对照论文逐步复现。目前已有1523人学习下载。项目中内置了CNN/DailyMail数据集的采样样本及映射文件并实现了基于Hugging Face Transformers的预训练模型加载、任务特定层构建、ROUGE指标评估等关键环节。通过阅读和运行这些代码开发者可以快速上手BERT在摘要生成任务中的完整流程深入理解自注意力机制、微调策略与解码器设计为后续开展自然语言生成相关研究打下扎实基础。1. 微调BERT做摘要提取一套论文代码的完整打开方式做NLP的多少都遇到过这种情况论文里说BERT加个分类头就能做抽取式摘要指标很好看可自己复现时要么显存爆掉要么ROUGE分数差一大截。这套代码就是干这个用的——它把“微调BERT提取摘要”这件事做成了完整的工程闭环从原始文本预处理、BERT训练、验证集评估到生成摘要一条命令跑通。代码面向的是想复现论文实验的开发者或者正在做文本摘要毕业设计、需要一套能改能跑的基线系统的同学。它的核心价值不在某个花哨的网络结构而在于把论文里模糊的“用BERT做抽取式摘要”翻译成了可运行的训练脚本、可观测的loss曲线、可比较的ROUGE分数。换句话说看完这篇拆解你能稳住训练流程知道每一步代码在做什么也能动手改出自己的版本。2. 架构总览一套代码如何组织训练、评估与推理2.1 仓库文件的职责划分这套代码不是单文件脚本而是按模块划分的标准论文工程结构。拿到了解压后的文件夹先别急着跑花五分钟看目录能省掉后面大量排查时间。常规模拟结构如下project_root/ ├── run_summarization.py # 主入口训练/评估/预测统一入口 ├── modeling.py # 模型定义BERTSUM核心结构 ├── data_loader.py # 数据读取与预处理 ├── rouge_eval.py # ROUGE评估脚本 ├── config.py # 超参数配置 ├── data/ │ ├── train.json │ ├── val.json │ └── test.json └── outputs/ ├── checkpoints/ └── predictions/run_summarization.py承担了所有模式的控制逻辑通过--do_train、--do_eval、--do_predict三个开关决定当前运行阶段。这种设计借鉴了Transformers官方示例的风格好处是训练和预测共享同一套数据加载逻辑不会出现训练集和测试集预处理不一致的问题。modeling.py里实现的不是简单的BERT Linear而是基于BertSum论文的经典结构——BERT顶部加一个sentence-level的Transformer层再对每个句子做二分类。config.py集中管理所有可调参数这算是我比较欣赏的设计。论文实验的复现最怕的就是参数散落在代码各处想调一个学习率要翻三个文件。集中在一个配置文件里用argparse逐项展开既方便命令行覆盖也方便做消融实验。文件里几乎每个参数都有英文注释说明默认值和调整方向对刚入门的开发者比较友好。2.2 数据流从原始文本到摘要标签抽取式摘要的任务定义是把“从原文中挑哪些句子组成摘要”建模为句子级二分类。数据集中每篇文档已经预先切好了句子并且标注了哪些句子是摘要句。data_loader.py读取JSON格式的原始数据后会按文档为单位做如下处理将整篇文档的句子拼接为BERT可接受的长度通常截断到512 token为每个句子生成一个标签向量1表示该句被选入摘要0表示不选按固定batch size组织成tensor喂给模型一个值得注意的细节是数据集中的“摘要句”标签在论文里通常用两种方式构造一种是直接使用人工标注的摘要句与原文句子的贪婪匹配结果另一种是用ROUGE分数自动挑选。这套代码两种模式都支持在config.py里能看到--label_mode这个参数可选值为greedy或rouge。默认使用greedy模式它更贴近原始论文的做法生成的标签更稳定不会因为ROUGE计算库的版本差异影响训练标签的一致性。2.3 训练与评估的交互模式训练流程的完整生命周期是加载预训练BERT权重 → 在摘要数据集上微调 → 每个epoch结束跑一次验证集评估 → 保存最优checkpoint。run_summarization.py中的主循环负责把整个流程串起来。验证阶段比较特殊的点是它不仅计算loss还会用当前的模型权重对验证集做一次贪婪解码即每句话独立判断是否入选然后直接算ROUGE分数。这样做的好处是让开发者直接看到ROUGE指标而不是只盯着loss曲线避免出现loss下降但实际摘要质量变差的“off-target”情况。我在复现类似项目时也习惯在验证阶段就做完整的解码评估虽然会慢一些但对判断训练状态非常有价值。代码里默认每两个epoch做一次完整评估如果想更快看到趋势可以改成每个epoch一次。3. 核心实现拆解模型结构、训练循环与ROUGE评估3.1 BERTSUM模型在BERT之上加一层句子Transformermodeling.py中模型实现的巧妙之处在于它将BERT输出的每个句子的首个[CLS]向量抽取出来再经过一个positional encoding和一层Transformer encoder最后接一个二分类层输出得分。这样做能让模型建模句子之间的上下文关系而不是独立判断每个句子。class BertSumModel(nn.Module): def __init__(self, bert_model, hidden_size768, num_labels2): super().__init__() self.bert bert_model # 句子级别的Transformer层捕捉句间依赖 self.sentence_transformer nn.TransformerEncoder( nn.TransformerEncoderLayer( d_modelhidden_size, nhead8, dim_feedforward2048, dropout0.1, activationgelu ), num_layers2 ) # 位置编码让模型知道句子顺序 self.position_embedding nn.Embedding(512, hidden_size) self.classifier nn.Linear(hidden_size, num_labels) def forward(self, input_ids, attention_mask, sentence_spans): # 获取BERT的隐层输出 bert_outputs self.bert( input_idsinput_ids, attention_maskattention_mask, output_hidden_statesTrue ) hidden_states bert_outputs.last_hidden_state # 抽取每个句子的CLS向量 sentence_vectors [] for span in sentence_spans: cls_idx span[0] # 每个句子的起始位置 sentence_vectors.append(hidden_states[:, cls_idx, :]) sentence_vectors torch.stack(sentence_vectors, dim1) # 加上位置编码后过句子Transformer positions torch.arange(sentence_vectors.size(1)).to(sentence_vectors.device) sentence_vectors sentence_vectors self.position_embedding(positions) sentence_features self.sentence_transformer(sentence_vectors) # 二分类得到每个句子的入选概率 logits self.classifier(sentence_features) return logits这段代码的逻辑可以拆成四步理解。第一步把整篇文档过BERT拿到每个token的向量表示。第二步通过sentence_spans这个参数定位到每个句子的起始token位取该位置的hidden state作为句子向量。这是抽取式摘要里最常用的做法——直接用句子首token的表示指代整个句子。第三步把句子向量加上可学习的位置编码再过两层Transformer encoder让每个句子的表示融入全篇句子的上下文信息。第四步接一个线性层做二分类。参数说明hidden_size768对应BERT-base的隐层维度如果换成BERT-large需要改成1024。nhead8是多头注意力头数dim_feedforward是前馈网络中间层维度。如果你发现显存不足优先砍掉sentence_transformer的层数从2改成1或者减小dim_feedforward对ROUGE分数的影响通常不大。3.2 训练循环从数据装载到损失计算训练循环写在run_summarization.py里核心逻辑并不复杂但有几个容易被忽视的细节。比如标签的构造方式、mask的处理、以及loss在句级别上的聚合方式。for batch in train_dataloader: input_ids batch[input_ids].to(device) attention_mask batch[attention_mask].to(device) labels batch[labels].to(device) sentence_spans batch[sentence_spans] logits model(input_ids, attention_mask, sentence_spans) # 计算每个句子的分类损失 loss_fct nn.CrossEntropyLoss() loss loss_fct( logits.view(-1, 2), # [batch_size * num_sentences, 2] labels.view(-1) # [batch_size * num_sentences] ) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() # 梯度裁剪防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm1.0) total_loss loss.item() if step % 100 0: logging.info(fstep {step}: loss {loss.item():.4f})这个循环里最有价值的是两处。第一处是loss_fct用的是普通的CrossEntropyLoss没有对正负样本做加权。抽取式摘要的数据集里负样本非摘要句数量远多于正样本但这套代码默认不做类别平衡——原因是论文里发现BERT的预训练表示已经足够强直接微调就能达到不错的效果加类别权重反而可能让模型变得保守。如果你想处理类别不平衡可以在config.py里加一个--pos_weight参数来控制。第二处是梯度裁剪clip_grad_norm_设1.0是BertSum系列实验的常用配置。BERT微调时后期容易出现loss震荡梯度裁剪能从一定程度上遏制这个问题。如果训练时发现loss曲线出现尖峰第一反应不是调学习率而是检查梯度裁剪的阈值是否合理。3.3 ROUGE评估实现细节与常见偏差来源rouge_eval.py实现了标准ROUGE-1、ROUGE-2、ROUGE-L的计算。这段代码我单独拎出来讲是因为它在复现论文时最容易被坑。很多论文报告的分数是基于特定tokenization规则和摘要后处理步骤的直接调第三方库得到的分数不一定可比。def compute_rouge(generated, reference): 生成摘要和参考摘要的ROUGE计算 Args: generated: 模型生成的摘要句子列表 reference: 参考摘要句子列表 gen_str .join(generated) ref_str .join(reference) # 简易ROUGE-1计算unigram的F1 gen_tokens gen_str.split() ref_tokens ref_str.split() overlap len(set(gen_tokens) set(ref_tokens)) precision overlap / len(gen_tokens) if gen_tokens else 0 recall overlap / len(ref_tokens) if ref_tokens else 0 f1 2 * precision * recall / (precision recall) if (precision recall) else 0 # ROUGE-L使用LCS算法 def lcs(x, y): m, n len(x), len(y) dp [[0] * (n 1) for _ in range(m 1)] for i in range(1, m 1): for j in range(1, n 1): if x[i-1] y[j-1]: dp[i][j] dp[i-1][j-1] 1 else: dp[i][j] max(dp[i-1][j], dp[i][j-1]) return dp[m][n] lcs_len lcs(gen_tokens, ref_tokens) rouge_l_precision lcs_len / len(gen_tokens) if gen_tokens else 0 rouge_l_recall lcs_len / len(ref_tokens) if ref_tokens else 0 rouge_l_f1 2 * rouge_l_precision * rouge_l_recall / \ (rouge_l_precision rouge_l_recall) if (rouge_l_precision rouge_l_recall) else 0 return { rouge-1-f1: f1, rouge-l-f1: rouge_l_f1 }这段实现是简化的教学版本对比rouge库时你会发现它有若干差异点。比如这里的ROUGE-1没有做stopwords过滤而某些论文里的指标是过滤掉停用词后计算的。再比如ROUGE-L的LCS算法在这个实现里用的是标准动态规划复杂度为O(mn)当摘要句子数量大时运行会比较慢。实际项目里建议直接使用rouge_score或py-rouge库但要记得在论文复现时明确说明采用的是哪种计算口径。还要注意一个细节这套评估代码在输出摘要时会按原始句子顺序排列而不是按模型预测分数从高到低排列。这在抽取式摘要任务里是一个需要明确的点——摘要要求可读性保持原文顺序比按重要性排序更合理。如果你在论文里看到“预测摘要”是按分数排序生成的那多半是另一个版本的实现。4. 数据预处理链路长文本截断、句子切分与标签构造4.1 超过512 token的长文档怎么处理BERT的输入上限是512个token但新闻文档经常远超这个长度。这套代码给出的方案是滑动窗口式截断——保留文档开头和结尾的部分从中间截取同时保证句子完整。具体实现逻辑是按句号、问号、感叹号把文本切成句子然后从文档开头按顺序拼接句子直到剩余空间不足以放下下一句为止。def truncate_document(sentences, max_tokens512, tokenizerNone): 将长文档截断到BERT可处理的最大长度 策略优先保留文档开头保证前文完整 truncated_sentences [] current_length 0 for sent in sentences: token_len len(tokenizer.encode(sent, add_special_tokensFalse)) # 预留[CLS]和[SEP]的位置 if current_length token_len 2 max_tokens: break truncated_sentences.append(sent) current_length token_len return truncated_sentences这个截断策略虽简单但有效新闻类文本的关键信息通常分布在前几段保留开头信息损失最小。max_tokens512是BERT-base的硬限制如果你的场景需要更长的上下文可以考虑改用Longformer或BigBird的权重替换BERT部分但这套代码本身不直接支持需要额外改modeling.py。一个常被忽略的参数是add_special_tokensFalse——编码句子时先不添加[CLS]和[SEP]而是在最终拼接整个文档时统一添加。这样能精确控制总长度不超过512。如果你在代码里直接tokenizer.encode(sent)忘记了关掉特殊token截断效果会偏差导致实际能容纳的句子数减少。4.2 句子边界定位与标签构造data_loader.py中最关键的输出是sentence_spans——它记录每个句子在文档token序列中的起始位置。构造过程是先对整个文档做tokenize得到整篇的input_ids序列然后逐一找到每个句子的第一个token在整篇序列中的下标位置。这个位置的定位依赖于tokenizer的offset mapping能力。def build_sentence_spans(tokenized_doc, sentence_tokens_list): 返回每个句子的起始token位置下标 spans [] start_idx 1 # 跳过开头的[CLS] for sent_tokens in sentence_tokens_list: sent_len len(sent_tokens) spans.append([start_idx, start_idx sent_len - 1]) start_idx sent_len return spans标签的构造则取决于--label_mode参数。greedy模式下代码遍历参考摘要中的每个句子与文档中的每个句子做贪婪匹配若ROUGE-2分数超过阈值则将该句标记为摘要句。这种方式生成的标签密度相对稀疏通常一篇文档只有两三个正样本。4.3 训练集与验证集的预处理一致性这是复现实验时最容易翻车的地方——训练时用的截断方式、句子切分规则如果和验证时不一致模型效果会明显下降。这套代码在data_loader.py中统一封装了build_dataset_from_file函数无论是训练集、验证集还是测试集都走同一条预处理管线。切句规则全部基于spacy或nltk的sent_tokenizetokenizer全部使用同一个BERT tokenizer。这样基本杜绝了“训练用的是A切句逻辑验证用的是B切句逻辑”的错位问题。我在自己的项目里吃过这方面的亏——某次为了省时间验证集预处理写了个简易版按句号切句结果ROUGE分数掉了四个点后来发现是预处理不一致导致的。从那以后我每次改数据流程都会强制走一遍验证集的样例输出确认和训练集完全一致后再跑实验。5. 常见问题与避坑复现论文实验的五个经典场景5.1 显存不足导致batch size跑不动现象训练脚本启动后CUDA out of memory程序直接退出。原因BERT的参数量非常大加上sentence_transformer层和梯度显存占用远超预期。特别是你机器上如果跑的是BERT-largebatch size设成8基本会爆显存。解决先把batch_size降到2或4如果还爆就检查max_length是否设成了512而实际数据没那么长。另一个常见做法是打开--gradient_accumulation_steps参数设成4或8用小batch多次累加梯度效果等价于大batch。我在复现时通常会先用batch size 1跑通一个step确认显存占用再逐步调大。5.2 ROUGE分数比论文低很多现象训练完的模型在验证集上ROUGE-1只有35左右论文里报告40。原因排查方向有三个。第一ROUGE计算口径不同论文可能用了带stopwords过滤或stemming的版本第二验证集的截断方式不同有些实现直接把超长文档从中间砍掉有些保留开头第三检查标签构造的阈值greedy模式里ROUGE-2阈值设成0.5还是0.7对最终标签密度影响很大。解决先从代码里找到验证集的评估函数确认它的ROUGE实现是对着哪个参考库写的。然后用论文的测试集样例如果有公开的话跑一遍预测对比输出摘要的长度和句子数是否接近。如果预测摘要普遍比参考摘要短多半是标签阈值偏高调低阈值重新训练。5.3 训练loss下降但验证集ROUGE不涨现象loss从2.1降到0.8验证集ROUGE分数停滞在初始水平附近。原因这是典型的过拟合信号模型在训练集上学会了记忆但没有泛化。另一个可能是学习率设置过大BERT部分的权重被破坏得太厉害。解决首先确认BERT部分的学习率设置——很多实现会为BERT层单独设置一个较小的学习率如2e-5而分类头用较大的学习率如1e-4。如果整个模型共用一个比较大的学习率很容易把BERT预训练的知识冲掉。其次检查是否加weight decay和warmup比例这两项在Transformer微调中缺一不可。我一般会在config.py里把--learning_rate设成3e-5--warmup_proportion设成0.1跑几个epoch观察曲线。5.4 tokenizer加载慢且每次启动都重新下载现象每次运行脚本from_pretrained都像是重新下载模型权重网络差的时候卡半天。原因HuggingFace的from_pretrained方法默认从远程拉取预训练权重如果本地没有缓存。解决显式指定本地缓存路径。在代码里把from_pretrained(bert-base-uncased)改成先下载到本地目录再指定路径加载。具体做法是运行一次transformers-cli download bert-base-uncased或者在config.py里设置--model_path参数指向你本地已下载好的模型目录。这样不仅启动快还能避开断网或代理问题。5.5 预测阶段出现重复句子或摘要过短现象模型对每句话的打分普遍偏低最终入选的句子只有一句摘要不完整。原因预测阶段没有限制最低摘要长度模型倾向于保守预测。解决在run_summarization.py的预测逻辑中加一个后处理如果入选句子数少于--min_sentence_num比如设为3则按分数从高到低补充缺少的句子到原文中。这个后处理不改变模型本身但对最终生成摘要的可用性提升很明显。还有一个可选做法是加入--extract_ratio参数控制提取句子比例占全文的比例上限比如0.3防止选太多句子导致摘要冗长。6. 二次开发实战给模型加上句间去重能力并验证效果如果你已经跑通了这套代码想在此基础上做点自己的改进我建议从“句间冗余去除”这个角度切入。抽取式摘要一个常见问题是选出的句子各自独立可能包含重复信息。BERTSUM的num_layers2的Transformer只说捕捉上下文信息但没有显式的去重机制。一个轻量级方案是在后处理阶段增加MMR最大边际相关过滤原理是同时考虑句子与文档的相关性ROUGE分数和句子间的相似度余弦相似度在选出摘要句时对已选句子相似度太高的候选进行惩罚。from sklearn.metrics.pairwise import cosine_similarity def mmr_filter(sentence_embs, scores, lambda_val0.7, max_len3): MMR后处理在保持相关性的同时减少冗余 Args: sentence_embs: 句子向量矩阵shape [N, dim] scores: 模型输出的每个句子的分数 lambda_val: 权衡相关性和冗余度的参数 max_len: 最大摘要句子数 selected_idx [] remaining list(range(len(scores))) for _ in range(max_len): if not remaining: break mmr_scores [] for idx in remaining: # 与已选句子的最大相似度 if selected_idx: max_sim max( cosine_similarity([sentence_embs[idx]], [sentence_embs[j]])[0][0] for j in selected_idx ) mmr lambda_val * scores[idx] - (1 - lambda_val) * max_sim else: mmr scores[idx] mmr_scores.append(mmr) best_idx remaining[mmr_scores.index(max(mmr_scores))] selected_idx.append(best_idx) remaining.remove(best_idx) return selected_idx这段代码的核心在lambda_val参数上它的值越大越看重句子和文档的相关性越接近0越强调候选句子与已选句子的差异性。验证改进效果的方法很简单在验证集上同时跑原版后处理和改进版本对比ROUGE-2分数——ROUGE-2对n-gram重叠敏感能反映摘要中是否存在重复的短语。通常加入MMR后ROUGE-2会小幅提升因为冗余减少后单位字数里携带的信息量增加了。接入方式是在run_summarization.py的预测阶段把原始的logits threshold判断改成先取top-k候选k设大一点比如8再用MMR筛选到最终长度。句子的向量可以从modeling.py中sentence_vectors的输出拿——具体做法是让模型的前向推理返回特征向量而不是最终logits。我在自己复现类似项目时犯过一个错误直接对logits做MMR而不是对句子特征向量做cosine相似度。logits是分类得分维度只有2做相似度计算完全没有区分度。换回特征向量后冗余去除效果才真正体现出来。从那以后我每次做摘要相关改造都会强制走一遍“特征向量和分类头分开输出”的验证流程确认拿到的向量是有语义信息的。回看这套代码它的价值不在于算法多超前而在于把抽取式摘要的完整流程组织得清清楚楚数据怎么处理、模型怎么训练、分数怎么评估、坑在哪里。从第一次跑通到今天能按自己的需求改后处理逻辑我对BERT微调的理解基本都是从这套代码里长出来的。希望这篇拆解能帮你在自己的实验里省下排查的时间把精力花在该花的地方——比如想清楚你的摘要任务到底要优化ROUGE还是可读性。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助?
咨询建站