简介基于Transformer的生成式文本摘要Python源码专为自然语言处理开发者与研究者设计帮助快速搭建并理解生成式摘要系统适用于新闻提炼、报告概括等长文档压缩场景。压缩包共10个文件包含5个Python脚本、1个TSV数据文件、1个依赖清单、1个说明文档及图片等包体仅201KBPython脚本覆盖数据预处理、模型定义、训练与推理等完整流程README说明环境配置与目录结构requirements.txt列明依赖便于直接上手。生成式摘要依托自注意力机制捕捉长距离依赖相比抽取式方法可生成更具连贯性与概括性的文本源码为此提供了简洁可运行的最小实现。已有119人学习参考适合具备一定深度学习基础、希望深入Transformer摘要原理的读者通过阅读和修改源码可掌握序列到序列训练细节与注意力计算过程并在此基础上进行微调或扩展快速验证模型改进思路节约从零搭建的时间成本。1. 生成式摘要的难点和这份源码的真实价值这份基于 Transformer 的生成式文本摘要 python 源码包不是那种给你一个训练好的模型、跑一遍predict.py就完事的黑盒封装而是一个从数据预处理、模型定义、训练循环到 beam search 解码都能逐行改动的完整框架。真正写过摘要算法的人会有同感抽取式摘要只需要判断每个句子是否保留但生成式摘要要求模型重新组织语言要解决词汇重复、OOV未登录词、句子不通顺三个老问题。Transformer 架构的出现把自注意力机制变成了标配并行度高长距离依赖处理能力强所以现在开源社区里的生成式摘要项目几乎都围绕它展开。这个源码适合两类人一是想搞清楚 encoder-decoder 内部每一步张量变化的 NLP 学习者二是想把摘要能力接到内部文档系统、没有预算去调预训练大模型 API 的工程师。下面我按“原理 → 代码拆解 → 调参 → 上线”的顺序把它讲透。2. Transformer 解码器是怎么“写”出摘要的从自注意力到生成式文本摘要2.1 编码器-解码器结构中的自注意力Transformer 把摘要任务当作序列到序列问题输入原文是 source输出摘要文本是 target。编码器读全文解码器每次生成一个 token同时在每步能“回看”原文和解码器已经生成的部分。这里的关键是自注意力self-attention。一个 token 的最终表示不是只看它自己而是把它与序列中所有其他 token 做相似度加权求和。这个机制直接建立了任意两个位置之间的路径长度为一的依赖不像 LSTM 需要信息经过一个 token 一个 token 地传递才能到达远端。为了让模型从多个角度理解语义源码里通常实现多头注意力multi-head attention把 d_model 维向量切分成多组每组各自计算注意力再拼接。下面是这个源码包里最常见的 scaled dot-product attention 的 python 实现片段def attention(query, key, value, maskNone, dropoutNone): d_k query.size(-1) scores torch.matmul(query, key.transpose(-2, -1)) scores scores / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) p_attn scores.softmax(dim-1) if dropout is not None: p_attn dropout(p_attn) return torch.matmul(p_attn, value), p_attn这段代码里为什么除法要除以sqrt(d_k)当 d_k 比较大时点积结果的方差会随维度增大softmax 容易进入梯度极小的饱和区所以用缩放因子把方差拉回来。masked_fill是解码器用来屏蔽未来信息的生成第 t 个 token 时注意力只能看到前 t-1 个 token不能偷看之后的位置否则训练和推理就出现不一致。这也是生成式摘要和抽取式摘要一个重要的实现差异。2.1.1 位置编码的必要性因为自注意力本身不包含顺序信息必须给每个 token 叠加一个位置向量才能让模型知道“哪些词在前哪些词在后”。源码里常见的实现是正弦位置编码但更重要的是要理解它为什么放在 embedding 之后、进编码器之前。如果把它忘掉模型会把 “A 打败 B” 和 “B 打败 A” 当成一样的输入摘要结果自然完全不可用。很多从 RNN 转过 Transformer 的人在第一次写摘要模型时就在这里栽过跟头。2.2 生成式摘要与抽取式摘要的本质区别抽取式摘要本质上是一个句子级别的分类任务把原文里的句子按重要性打分、按顺序拼接起来模型不需要产生语料里没出现过的词。生成式摘要就要难得多它要求模型理解原文语义并发出一个全新的序列。很多第一次接触的人容易踩坑把 ROUGE 分数低归结为模型不够复杂其实问题往往出在解码策略上。同样一个 Transformer 编码器如果只用贪心解码每步取概率最大的词生成的摘要很容易落入重复循环比如反复输出“经济增长 经济增长”。这是因为贪心解码只看局部最优一旦某个词形成高概率路径就停不下来。源码里的 beam search 会同时维护多个候选序列每个候选维护自己的累积对数概率然后从所有候选的分支中选择 top-k避免了一味跟着一个死胡同走到黑。这里还要注意解码器和编码器是用两层掩码分别处理 padding 和未来信息的这在训练阶段就要做对否则 inference 阶段的时间偏移问题会让你怎么调都调不出流畅通顺的摘要。2.3 从 RNN 到 Transformer 的选型理由更早的摘要系统用 LSTM 做编码器和解码器当前 hidden state 必须携带整个历史信息无法处理特别长的新闻或技术文档。Transformer 把逐位置的前馈网络和注意力相结合每个 token 的初始表示通过 positional encoding 带上位置信息然后经过多层叠加逐步融合全局上下文。它的另一个好处是训练可以并行计算整个序列的注意力矩阵而 RNN 只能按时间步串行。代价是复杂度随着序列长度平方增长所以源码里通常把最大原文长度限制在 512 或 1024并在数据预处理阶段截断这属于工程上非常务实的做法。用这个理由去理解为什么项目里选择 Transformer 而不是一开始就挂上预训练语言模型预训练模型确实效果更好但加载成本和推理延迟都会高一个量级。这份源码走的是从零构建路线更容易看懂每一步在计算什么后续接入私有大模型时你至少能准确说出“我现在的瓶颈是 beam search 的解码长度还是 encoder 的序列长度”而不至于把全部问题都丢给 prompt。3. 源码拆解数据流与训练流程3.1 项目目录结构与数据预处理解压后目录结构如下是一个典型的 PyTorch 工程布局. ├── code │ ├── data_loader.py │ ├── model.py │ ├── train.py │ ├── translate.py │ └── utils.py ├── data │ ├── train.jsonl │ ├── valid.jsonl │ └── test.jsonl ├── img ├── requirements.txt ├── .gitignore └── README.mddata目录里的jsonl文件是目前摘要项目里比较通用的格式每行一个 JSON 对象包含原始文本和参考摘要两个字段。我一般会把它转成 torch 的 Dataset核心代码如下class SummarizationDataset(Dataset): def __init__(self, path, tokenizer, max_src_len512, max_tgt_len128): self.samples [] with open(path, r, encodingutf-8) as f: for line in f: obj json.loads(line) src tokenizer.encode(obj[src])[:max_src_len] tgt tokenizer.encode(obj[tgt])[:max_tgt_len] self.samples.append((src, tgt)) def __len__(self): return len(self.samples) def __getitem__(self, idx): return self.samples[idx]这里有个细节tokenizer是决定模型能不能收敛的重要因素。如果用单字切分代码最简单但中文里模型很难学到词级别语义。这个源码包里用的是 byte-pair-encodingBPE这类子词切分方式会把“自然语言处理”切成“自然”“语言”“处理”等更细的 token缓解 OOV 问题。max_src_len和max_tgt_len要分别设置原文长度 512、摘要长度 128 是典型的初始配置。如果你的数据是英文新闻摘要可以直接用英文 BPE 分词中文的话建议先用 jieba 预分词后训练 BPE或者直接用已有的中文分词器别用空格分词糊弄。3.1.1 collate_fn 里的 padding 处理在一个 batch 里每句话长度不一样不能直接堆成 tensor所以要在collate_fn中做 padding。这段代码虽然不在所有摘要项目里都长得一致但思路是通用的把每个 batch 内部的 src 和 tgt 分别 pad 到当前 batch 最长长度同时生成对应的src_mask和tgt_mask再按长度降序排列。如果不做这个步骤后面 Transformer 的 attention 计算会因为维度不一致直接报错。3.2 手写 Transformer 构建代码关键实现片段model.py 里主要包含EncoderLayer、DecoderLayer和完整Transformer类。以多头注意力为例它的维度操作是这样的把句子的(batch, seq_len, d_model)输入通过三个权重矩阵分别投影得到(batch, num_heads, seq_len, d_k)然后交换维度做注意力计算。模型定义里要特别留意两个掩码src_mask把 padding 位置改成 0防止注意力集中在无效的blanktoken 上tgt_mask同时处理 padding 和未来信息的屏蔽用一个上三角矩阵把未来位置挡住。这里给出手写 Transformer 里多头注意力的一个关键初始化片段方便对照源码看self.linear_q nn.Linear(d_model, d_model) self.linear_k nn.Linear(d_model, d_model) self.linear_v nn.Linear(d_model, d_model) self.linear_out nn.Linear(d_model, d_model) def forward(self, query, key, value, maskNone): batch_size query.size(0) q self.linear_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k self.linear_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v self.linear_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) x, _ attention(q, k, v, mask) x x.transpose(1, 2).contiguous().view(batch_size, -1, d_model) return self.linear_out(x)注意view后为什么还要transpose因为要把num_heads调到batch和seq_len之间后续注意力计算才能让每个头独立算自己的 QK^T。这个 reshape 顺序错了结果不会直接报错但每个头会算到别的头的向量上模型性能会异常差很多人会在这一步浪费不少时间。3.3 训练循环与教师强制训练阶段采用 teacher forcing解码器每一步的输入不用上一步预测出的 token而是用参考摘要里真实的 token。这样模型收敛快很多但代价是训练和推理之间存在 gap。源码里训练循环通常长这样for src, tgt in train_loader: src, tgt src.to(device), tgt.to(device) tgt_input tgt[:, :-1] tgt_output tgt[:, 1:] logits model(src, tgt_input, src_mask, tgt_mask) loss criterion(logits.reshape(-1, vocab_size), tgt_output.reshape(-1)) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step()为什么tgt_input和tgt_output要错开一位因为模型在位置 i 的任务是预测位置 i1 的 token所以输入序列和标签序列整体右移一位这也是生成式序列任务最常见的实现约定。损失函数用的是交叉熵但要设置ignore_indexpad_idx否则 padding 位置的空白 label 也会参与 loss让模型浪费精力学习输出pad。梯度裁剪max_norm1.0是防止 loss spike 时梯度爆炸的保险这在摘要任务里很常见特别是模型刚起步时。3.3.1 混合精度和断点续训如果训练数据稍大建议在训练脚本里加上 AMP自动混合精度把torch.cuda.amp.autocast()和GradScaler包住 forward 和 backward。这样做显存占用能降低 30% 到 40%同时训练速度明显提升。源码里不一定默认开启但自己加上并不难。另外保存 checkpoint 时最好把 optimizer 的 state dict 一起存下来断点续训时才能恢复到和中断前一致的 momentum 和学习率状态而不是只恢复模型权重。3.4 推理beam search 生成摘要到了 translate.py模型就切换到自回归生成模式。每步生成时把已经生成的序列重新喂给 decoder预测下一个 token 的概率分布。为了提升效果常用 beam search 代替贪心解码维护一个大小为 beam_size 的候选列表。简化后的核心逻辑如下def beam_search(model, src, beam_size4, max_len128): src_mask (src ! pad_idx).unsqueeze(1) encoder_output model.encode(src, src_mask) beams [(0.0, [bos_idx], encoder_output)] for _ in range(max_len): new_beams [] for score, seq, mem in beams: tgt_input torch.tensor([seq]).to(device) logits model.decode(mem, tgt_input, tgt_mask) probs logits[:, -1, :].log_softmax(dim-1) topk probs.topk(beam_size) for i in range(beam_size): new_score score topk.values[0][i].item() new_beams.append((new_score, seq [topk.indices[0][i].item()], mem)) beams sorted(new_beams, keylambda x: x[0], reverseTrue)[:beam_size] return beams[0][1]这个实现为了好读把不少工程优化省掉了比如已经生成eos_idx的序列要单独处理并保留但思路是一样的每条路径的分数是各步对数概率之和beam 宽度越大搜索空间越大生成质量通常越好但解码时间也越长。源码里一般还会在分数上除以len(seq) ** length_penalty给短句一个倾向性避免模型一味生成很长的摘要去稀释概率值。4. 超参数设置与模型调优4.1 核心超参数一览我把这个源码里最常见的超参数整理成一张表并给出在 8G 到 16G 显存条件下比较稳妥的初始值超参数含义建议值说明d_model词向量与注意力维度256 或 512256 适合小数据512 适合数据量充足的场景n_heads多头注意力头数8需要能被 d_model 整除否则报错num_layers编码器/解码器层数4 或 6层数翻倍训练时间几乎翻倍max_src_len原始文本最大长度512超过部分截断太长会显存溢出max_tgt_len摘要最大长度128根据业务需要调整batch_size每个 batch 的句子数16 到 64按显存调OOM 就减半learning_rateAdam 初始学习率0.0003 到 0.001配合 warmup 使用beam_size解码候选数4训练阶段不用推理阶段使用length_penalty长度惩罚系数0.6 或 1.0数值越小生成越倾向于短句表格里n_heads必须整除d_model这是很多人写自定义 Transformer 时第一个报错的地方因为d_k d_model // n_heads是整数除法除不尽就会出现维度不匹配。learning_rate我习惯用 Noam 式学习率调度先线性上升到指定步数再按步数倒数衰减。这种方案比固定学习率稳定得多能明显减少训练中后期的 loss 震荡。4.2 小数据上做摘要的调参策略很多使用者一上来就按照原论文设d_model512、num_layers6、batch_size64结果自己的文本摘要数据集只有几万条甚至几千条模型很容易过拟合训练 loss 不断下降验证 ROUGE 却停滞不动。这时候要先降模型容量把d_model降到 256num_layers降到 4dropout 提到 0.3同时把max_src_len截短到 384减少 padding 带来的无效计算。对于数据量更少的场景我用过一个更实用的做法保留这个源码的生成解码头和训练流程但把编码器权重替换成已经预训练好的中文 BERT 或者 RoBERTa 前几层这样原文理解能力可以直接借用预训练知识。这个源码本身没有内置这部分需要自己写一个兼容的模型加载钩子但改动量不大主要工作量在把预训练模型的 embedding 输出对齐到 Transformer decoder 的输入维度。如果追求最稳定的效果也可以直接跑 src 里的 baseline加上 label smoothing 0.1对生成质量和泛化都有正面效果。另一个容易被忽视的点是早停和模型保存。训练时不要只看最后一步的 checkpoint而是每隔一个 epoch 在验证集上算一次 ROUGE-L把分数最高的模型单独保存。很多摘要项目用训练 loss 判断是否收敛效果其实不如直接监控验证集 ROUGE因为 loss 下降不代表摘要可读性变好。4.3 评价指标 ROUGE 的计算与注意事项生成式摘要评测基本都看 ROUGE 家族。ROUGE-1 看单字/词重合ROUGE-2 看二元词组的重合ROUGE-L 看最长公共子序列。安装依赖后用rouge-score库一条命令就能算pip install rouge-score rouge-score --refs data/test.ref.txt --preds data/test.pred.txt更常见的是在脚本里逐条统计然后对整批结果求均值。这里有几个容易忽略的细节英文摘要记得把标点和[SEP]、[PAD]等特殊 token 全部滤掉中文摘要我建议先按字计算 ROUGE因为按词分会产生分词器不一致带来的误差如果你的数据是新闻标题按字算更稳定。另外ROUGE 与人工评价的相关性有限它只能告诉你模型生成的内容和参考摘要的重合程度不能告诉你摘要是否通顺、是否抓住核心信息。所以调参时不要只看一个指标至少要同时看 ROUGE-1、ROUGE-L 和你自己定义的“关键实体是否出现的精确率”。5. 排错与扩展从 demo 到可用的摘要接口5.1 显存不足和 batch size 的处理训练时报CUDA out of memory是最常见的问题。第一步把batch_size减半第二步把max_src_len从 512 降到 256大部分场景都能解决。还不行就用梯度累积每跑两步或四步才做一次反向传播模拟更大的 batch既不增加显存又能维持较大的有效批次。accumulation_steps 4 for step, (src, tgt) in enumerate(train_loader): loss compute_loss(src, tgt) / accumulation_steps loss.backward() if (step 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这里要把原始 loss 除以累积步数否则等价于把学习率凭空放大训练会不稳定。5.2 摘要重复和 OOV 问题处理生成摘要经常出现整段重复。除了调大beam_size更直接的办法是启用repetition_penalty或no_repeat_ngram_size。源码的生成方法里如果没有这个参数可以在 beam search 的下一步剪枝时把已经出现过 n-gram 的候选全部过滤掉只保留那些新组合不在已生成序列中重复的候选。即便这样会略微降低 ROUGE 分数但可读性会好很多。OOV 是由词表不够大导致的。BPE 能解决一部分但碰到人名、产品名仍然无能为力。实际工程中可以在后处理阶段加一个“原文实体替换”的规则先用实体识别抽取出摘要里的占位符对应的原文实体再替换回真实字符串或者干脆在解码时允许从原文复制 token这就是 pointer-network 的做法。这个源码包没内置但它的模型输出打分结构支持再加一个 copy gate改动量不大。5.3 把源码改造成 HTTP 服务最终要让别人用最简单的方式是用 FastAPI 包一层import uvicorn from fastapi import FastAPI, Request app FastAPI() def load_model(): from model import Transformer model Transformer(...) model.load_state_dict(torch.load(checkpoint.pt)) model.eval() return model model load_model() app.post(/summarize) async def summarize(request: Request): body await request.json() text body.get(text, ) summary model.generate(text, max_len128) return {summary: summary} uvicorn.run(app, host0.0.0.0, port8000)局部变量model在模块加载时初始化一次避免每个请求都重新加载权重。这一步做完你手里的就是一个可调、可测、可部署的摘要最小系统。要注意的是一次请求里文本长度超过编码器上限时直接返回 400 而不是硬截断把错误暴露给调用方而不是让模型在截断后的文本上生成莫名其妙的摘要。本文还有配套的精品资源点击获取