Transformers 文本摘要实战指南:基于 BillSum 数据集微调 T5 模型与推理部署
【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers
摘要(Summarization)是自然语言处理中的核心任务之一,它从长文档或文章中提炼出保留全部关键信息的短版本。在 🤗 Transformers 生态中,摘要任务通常被建模为序列到序列(sequence-to-sequence,seq2seq)问题,与机器翻译同属一类。本文以日文版官方任务文档 docs/source/ja/tasks/summarization.md 为主体,结合当前仓库源码,完整讲解如何利用 T5 在 BillSum 数据集的加州法案子集上微调抽象式摘要模型,并给出从环境准备、数据加载、预处理、评估到训练与推理的端到端方案。读完本文,你将掌握用Trainer/Seq2SeqTrainer微调 seq2seq 模型的标准流程,以及用pipeline与model.generate()两种方式部署摘要模型。
摘要任务的两种形态
摘要任务可按生成方式分为两类:
- 抽取式(Extractive):直接从原文中挑选最具相关性的句子或片段进行拼接,不产生新的表述。
- 抽象式(Abstractive):理解原文语义后,生成全新的、概括性的文本,通常更贴近人类写作,也是本文 T5 微调方案所针对的目标。
在 T5 这类"文本到文本"统一架构中,摘要与翻译一样,被当作标准的 seq2seq 任务处理:编码器读取长文本,解码器自回归地生成短摘要。
环境准备与前置依赖
开始之前,需要安装以下 Python 库:
pip install transformers datasets evaluate rouge_score各库职责如下:
transformers:模型架构、Trainer训练框架与推理组件(对应本仓库 src/transformers);datasets:加载与预处理 BillSum 数据集;evaluate:加载并计算 ROUGE 指标;rouge_score:ROUGE 指标的后端计算依赖(基于 Google Research 的 rouge-score 实现)。
如需把模型上传到 Hugging Face Hub 与社区共享,建议先登录账号,按提示输入 token:
>>> from huggingface_hub import notebook_login >>> notebook_login()加载 BillSum 数据集
使用 🤗 Datasets 加载 BillSum 数据集中较小的加州法案子集(ca_test划分):
>>> from datasets import load_dataset >>> billsum = load_dataset("billsum", split="ca_test")再用train_test_split方法按 20% 比例切分出训练集与测试集:
>>> billsum = billsum.train_test_split(test_size=0.2)查看一条样本,观察数据结构:
>>> billsum["train"][0] {'summary': 'Existing law authorizes state agencies to enter into contracts for the acquisition of goods or services upon approval by the Department of General Services. ...', 'text': 'The people of the State of California do enact as follows:\n\n\nSECTION 1.\nSection 10295.35 is added to the Public Contract Code, to read:\n...', 'title': 'An act to add Section 10295.35 to the Public Contract Code, relating to public contracts.'}样本中真正用于训练的两个字段是:
text:法案正文,作为模型的输入;summary:text的浓缩版本,作为模型的目标(label)。
title字段在本流程中不使用。
预处理:提示词前缀与标签编码
加载 T5 分词器:
>>> from transformers import AutoTokenizer >>> checkpoint = "google-t5/t5-small" >>> tokenizer = AutoTokenizer.from_pretrained(checkpoint)预处理函数需要完成三件事:
- 加任务前缀:在输入前拼接
"summarize: ",让 T5 识别当前是摘要任务。T5 支持多任务,不同任务依赖不同的提示词来触发对应能力; - 用
text_target编码标签:编码目标文本时使用text_target关键字参数,而不是重复传入inputs; - 截断序列:通过
max_length限制序列长度,防止超出模型上下文。
>>> prefix = "summarize: " >>> def preprocess_function(examples): ... inputs = [prefix + doc for doc in examples["text"]] ... model_inputs = tokenizer(inputs, max_length=1024, truncation=True) ... labels = tokenizer(text_target=examples["summary"], max_length=128, truncation=True) ... model_inputs["labels"] = labels["input_ids"] ... return model_inputs将预处理函数应用到整个数据集,batched=True可一次处理多个样本以加速map:
>>> tokenized_billsum = billsum.map(preprocess_function, batched=True)动态填充:DataCollatorForSeq2Seq
直接用DataCollatorForSeq2Seq组装批次。它的核心价值是动态填充(dynamic padding):在 collation 阶段把同批样本填充到该批最长长度,而非把整个数据集统一填充到最大长度,从而显著减少无效计算与显存占用。
>>> from transformers import DataCollatorForSeq2Seq >>> data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=checkpoint)从源码 src/transformers/data/data_collator.py 可以看到其内部实现逻辑:
- 输入部分交由 tokenizer 按
padding=True(即'longest')策略填充; - 标签部分手动填充:由于标签的填充符必须是
-100(PyTorch 损失函数会自动忽略该值),无法直接复用 tokenizer 的 pad token,因此 DataCollatorForSeq2Seq 将每个标签补齐到批内最大长度,并用label_pad_token_id=-100填充; - 若传入的模型实现了
prepare_decoder_input_ids_from_labels,还会自动生成decoder_input_ids,避免重复计算 decoder 输入——这对开启label_smoothing的训练尤其有用。
T5 恰好实现了该方法(modeling_t5.py),其内部调用_shift_right(labels)将标签右移一位并填充起始 token,从而构造解码器的自回归输入。
评估:用 ROUGE 衡量摘要质量
训练过程中引入指标有助于评估模型表现。使用 🤗 Evaluate 库快速加载 ROUGE:
>>> import evaluate >>> rouge = evaluate.load("rouge")编写compute_metrics函数,将预测与标签解码为文本后计算 ROUGE:
>>> import numpy as np >>> def compute_metrics(eval_pred): ... predictions, labels = eval_pred ... decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True) ... labels = np.where(labels != -100, labels, tokenizer.pad_token_id) ... decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) ... result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True) ... prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions] ... result["gen_len"] = np.mean(prediction_lens) ... return {k: round(v, 4) for k, v in result.items()}该函数的关键细节:
- 用
skip_special_tokens=True解码,剔除<pad>、<eos>等特殊 token; - 将标签中的
-100替换为tokenizer.pad_token_id,否则这些占位符会被解码成无意义 token,污染指标; use_stemmer=True启用词干还原,使 ROUGE 能匹配同词根变体(如 "running"/"run");- 额外统计
gen_len(生成平均长度),方便观察摘要长度是否符合预期。
ROUGE 家族包含 ROUGE-1、ROUGE-2、ROUGE-L 等子指标,分别衡量单字(unigram)、双字(bigram)与最长公共子序列层面的 n-gram 重叠度,是摘要与机器翻译领域的事实标准。
训练:Seq2SeqTrainer 微调 T5
先用AutoModelForSeq2SeqLM加载 T5 模型:
>>> from transformers import AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer >>> model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)随后只需三步:
- 在
Seq2SeqTrainingArguments中定义训练超参数(唯一必填项是output_dir); - 把训练参数连同模型、数据集、分词器、数据整理器与
compute_metrics一起传给Seq2SeqTrainer; - 调用
trainer.train()开始微调。
>>> training_args = Seq2SeqTrainingArguments( ... output_dir="my_awesome_billsum_model", ... eval_strategy="epoch", ... learning_rate=2e-5, ... per_device_train_batch_size=16, ... per_device_eval_batch_size=16, ... weight_decay=0.01, ... save_total_limit=3, ... num_train_epochs=4, ... predict_with_generate=True, ... fp16=True, # 在 XPU 设备上请改为 bf16=True ... push_to_hub=True, ... ) >>> trainer = Seq2SeqTrainer( ... model=model, ... args=training_args, ... train_dataset=tokenized_billsum["train"], ... eval_dataset=tokenized_billsum["test"], ... processing_class=tokenizer, ... data_collator=data_collator, ... compute_metrics=compute_metrics, ... ) >>> trainer.train()各超参数含义与取值参考:
| 参数 | 值 | 说明 |
|---|---|---|
output_dir | my_awesome_billsum_model | 模型保存目录(必填) |
eval_strategy | epoch | 每个 epoch 结束时评估一次并保存 checkpoint |
learning_rate | 2e-5 | Adam 优化器初始学习率,微调场景常用 1e-5~5e-5 |
per_device_train/eval_batch_size | 16 | 每设备批大小,需根据显存调整 |
weight_decay | 0.01 | L2 权重衰减系数 |
save_total_limit | 3 | 最多保留 3 个 checkpoint,防止磁盘占满 |
num_train_epochs | 4 | 训练轮数 |
predict_with_generate | True | 评估时用generate生成摘要来计算 ROUGE/BLEU 等生成式指标(关键) |
fp16/bf16 | True | 混合精度加速;XPU 设备改用bf16=True |
关于predict_with_generate,从 training_args_seq2seq.py 的源码可知:该参数决定评估阶段是否调用生成接口来计算生成式指标,同时配套提供generation_max_length(默认取模型配置的max_length)与generation_num_beams(默认取模型配置的num_beams)两个可选参数,用于精细控制评估时的解码策略;此外还支持通过generation_config直接传入一个GenerationConfig对象或路径。训练完成后的显式解码参数(如max_new_tokens)仅影响推理,不影响训练。
训练结束后,用push_to_hub把模型共享到 Hub:
>>> trainer.push_to_hub()进阶示例脚本
若想深入完整流程,可参考仓库自带的 examples/pytorch/summarization 目录:
- run_summarization.py:基于
Seq2SeqTrainer的完整训练脚本,支持从命令行传入模型名、数据集名、--max_source_length、--max_target_length、--num_beams、--source_prefix等参数; - run_summarization_no_trainer.py:不依赖
Trainer、使用原生 PyTorch 训练循环的版本; - README.md:两种脚本的详细参数说明与用法示例;
- requirements.txt:脚本运行所需依赖。
推理:两种部署方式
微调完成后即可用于推理。准备一段待摘要文本——注意 T5 的输入同样需要"summarize: "前缀:
>>> text = "summarize: The Inflation Reduction Act lowers prescription drug costs, health care costs, and energy costs. It's the most aggressive action on tackling the climate crisis in American history, which will lift up American workers and create good-paying, union jobs across the country. It'll lower the deficit and ask the ultra-wealthy and corporations to pay their fair share. And no one making under $400,000 per year will pay a penny more in taxes."方式一:pipeline 一行推理
将微调后的模型封装进摘要pipeline(将stevhliu/my_awesome_billsum_model替换为你自己的模型仓库名):
>>> from transformers import pipeline >>> summarizer = pipeline("summarization", model="stevhliu/my_awesome_billsum_model") >>> summarizer(text) [{"summary_text": "The Inflation Reduction Act lowers prescription drug costs, health care costs, and energy costs. It's the most aggressive action on tackling the climate crisis in American history, which will lift up American workers and create good-paying, union jobs across the country."}]pipeline内部自动完成了分词、生成与解码的完整链路,适合快速验证与轻量部署。
方式二:手动调用 generate
需要细粒度控制生成行为时,可手动完成"分词 → 生成 → 解码"三步。
先分词并返回 PyTorch 张量:
>>> from transformers import AutoTokenizer >>> tokenizer = AutoTokenizer.from_pretrained("stevhliu/my_awesome_billsum_model") >>> inputs = tokenizer(text, return_tensors="pt").input_ids再调用model.generate生成摘要,max_new_tokens=100限制新增 token 数,do_sample=False使用贪心解码:
>>> from transformers import AutoModelForSeq2SeqLM >>> model = AutoModelForSeq2SeqLM.from_pretrained("stevhliu/my_awesome_billsum_model") >>> outputs = model.generate(inputs, max_new_tokens=100, do_sample=False)最后把生成的 token id 解码回文本:
>>> tokenizer.decode(outputs[0], skip_special_tokens=True) 'the inflation reduction act lowers prescription drug costs, health care costs, and energy costs. it's the most aggressive action on tackling the climate crisis in american history. it will ask the ultra-wealthy and corporations to pay their fair share.'generate背后是仓库 src/transformers/generation 模块实现的GenerationMixin,支持贪心、束搜索(beam search)、采样(sampling)、对比搜索等丰富解码策略;调整num_beams、temperature、top_k、top_p等参数可改变生成质量与多样性,具体可参考仓库中的文本生成 API 文档 docs/source/en/main_classes/text_generation.md。
小结
本文以 BillSum 加州法案子集为例,走通了"数据加载 → 前缀与标签预处理 → 动态填充 → ROUGE 评估 →Seq2SeqTrainer微调 →pipeline/generate推理"的完整摘要模型落地链路。核心要点可概括为:
- 摘要可划分为抽取式与抽象式,T5 方案将抽象式摘要建模为 seq2seq 生成任务;
- T5 等多任务模型要求输入带任务前缀(
"summarize: "),标签用text_target编码; DataCollatorForSeq2Seq以-100填充标签并自动生成decoder_input_ids,兼顾效率与正确性;predict_with_generate=True是让评估阶段真正生成摘要并计算 ROUGE 的关键开关;- 推理既可用
pipeline快速落地,也可用model.generate()获得细粒度控制。
【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考