在这一部分中，我们将探讨如何使用 Transformer 模型将长文档压缩为摘要，这个任务被称为文本摘要。文本摘要是 NLP 领域中最具挑战性的任务之一，因为它需要理解长篇文本并生成连贯的总结，捕捉文档中的核心主题。然而，当文本摘要任务处理得当时，它能够极大地提高各种业务流程的效率，减轻专家阅读长篇文档的负担。
虽然 Hugging Face Hub 上已经有许多用于文本摘要的预训练模型，但这些模型几乎都只适用于英文文档。因此，在本节中，我们将增加一些变化，训练一个可以处理英语和西班牙语的双语模型。
1. 加载数据
准备多语言语料库
我们将使用多语言亚马逊评论语料库（Multilingual Amazon Reviews Corpus）来创建双语摘要模型。该语料库包含六种语言的亚马逊产品评论，通常用于多语言分类器的基准测试。由于每条评论都有一个简短的标题，我们可以将这些标题作为目标摘要，让模型从中学习！首先，让我们从 Hugging Face Hub 下载英语和西班牙语的子集：
import os
os.environ['http_proxy'] = ''
os.environ['https_proxy'] = '
from datasets import load_dataset
spanish_dataset = load_dataset("amazon_reviews_multi", "es")
english_dataset = load_dataset("amazon_reviews_multi", "en")
english_dataset
每种语言的训练集有 20 万条评论，验证集和测试集各有 5,000 条评论。我们感兴趣的评论信息包含在 review_body（评论内容）和 review_title（评论标题）列中。让我们定义一个函数，从训练集中随机抽取几条样本，看看评论和标题的样子：
def show_samples(dataset, num_samples=3, seed=42):
sample = dataset["train"].shuffle(seed=seed).select(range(num_samples))
for example in sample:
print(f"\n'>> Title: {example['review_title']}'")
print(f"'>> Review: {example['review_body']}'")
show_samples(english_dataset)
这样我们就可以看到评论和标题的多样性，评论内容从积极到消极不等。训练一个能够处理所有 40 万条评论的模型需要耗费大量时间，因此我们决定专注于为单一产品领域生成摘要。
我们将重点放在书籍评论领域，筛选出书籍（book）和电子书购买（digital_ebook_purchase）相关的评论：
def filter_books(example):
return (
example["product_category"] == "book"
or example["product_category"] == "digital_ebook_purchase"
)
spanish_books = spanish_dataset.filter(filter_books)
english_books = english_dataset.filter(filter_books)
show_samples(english_books)
接着，我们将英语和西班牙语的评论合并为一个 DatasetDict 对象，并随机打乱数据，以确保模型不会只关注某一种语言：
from datasets import concatenate_datasets, DatasetDict
books_dataset = DatasetDict()
for split in english_books.keys():
books_dataset[split] = concatenate_datasets(
[english_books[split], spanish_books[split]]
)
books_dataset[split] = books_dataset[split].shuffle(seed=42)
show_samples(books_dataset)
在处理文本摘要任务时，短参考摘要可能会使模型输出极简的结果，因此我们通过过滤掉标题过短的样本，确保模型可以生成更有意义的摘要：
books_dataset = books_dataset.filter(lambda x: len(x["review_title"].split()) > 2)
2. 数据预处理
文本摘要任务与机器翻译类似，我们需要将长文本“翻译”成一个更短的版本。大多数用于摘要任务的 Transformer 模型采用编码器-解码器架构，如 mT5，它是一种多语言 Transformer 模型。
首先，我们加载与预训练模型 mt5-small 相关的 tokenizer，并测试它是否能够正确处理输入文本：
from transformers import AutoTokenizer
model_checkpoint = "google/mt5-small"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
inputs = tokenizer("I loved reading the Hunger Games!")
inputs
接下来，为了确保模型不会处理过长的输入和输出，我们在对评论和标题进行标记时，对它们进行截断：
max_input_length = 512
max_target_length = 30
def preprocess_function(examples):
model_inputs = tokenizer(
examples["review_body"],
max_length=max_input_length,
truncation=True,
)
labels = tokenizer(
examples["review_title"], max_length=max_target_length, truncation=True
)
model_inputs["labels"] = labels["input_ids"]
return model_inputs
tokenized_datasets = books_dataset.map(preprocess_function, batched=True)
3. 文本摘要的评估指标
文本摘要任务的常用评估指标是 ROUGE 分数。它衡量生成摘要和参考摘要之间的重叠程度。我们可以使用 evaluate 库中的 rouge_score 来计算这些分数：
import evaluate
rouge_score = evaluate.load("rouge")
scores = rouge_score.compute(
predictions=["I absolutely loved reading the Hunger Games"],
references=["I loved reading the Hunger Games"]
)
scores
4. 创建强大的基线模型
一个常见的文本摘要基线是提取文章的前三个句子。我们可以使用 nltk 库来处理句子边界，创建一个简单的函数来提取前三个句子并计算 ROUGE 分数：
from nltk.tokenize import sent_tokenize
def three_sentence_summary(text):
return "\n".join(sent_tokenize(text)[:3])
def evaluate_baseline(dataset, metric):
summaries = [three_sentence_summary(text) for text in dataset["review_body"]]
return metric.compute(predictions=summaries, references=dataset["review_title"])
score = evaluate_baseline(books_dataset["validation"], rouge_score)
score
5. 模型
我们加载预训练的 mt5-small 模型，并使用 DataCollatorForSeq2Seq 来动态填充输入和标签。接下来，设置数据加载器和优化器，并将模型和数据加载器传递给 Accelerator.prepare() 进行训练：
from transformers import AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, AdamW, get_scheduler, Accelerator
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
train_dataloader = DataLoader(
tokenized_datasets["train"], shuffle=True, collate_fn=data_collator, batch_size=8
)
optimizer = AdamW(model.parameters(), lr=2e-5)
lr_scheduler = get_scheduler(
"linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=len(train_dataloader) * 3
)
accelerator = Accelerator()
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
model, optimizer, train_dataloader, eval_dataloader
)
6. 训练
我们实现训练循环，包括生成摘要、计算 ROUGE 分数和将结果上传到 Hugging Face Hub：
from tqdm.auto import tqdm
import numpy as np
progress_bar = tqdm(range(len(train_dataloader) * 3))
for epoch in range(3):
model.train()
for step, batch in enumerate(train_dataloader):
outputs = model(**batch)
loss = outputs.loss
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
model.eval()
generated_tokens = []
for step, batch in enumerate(eval_dataloader):
with torch.no_grad():
outputs = model.generate(batch["input_ids"], attention_mask=batch["attention_mask"])
generated_tokens.append(outputs)
rouge_score.add_batch(predictions=generated_tokens, references=batch["labels"])
print(f"Epoch {epoch}:", rouge_score.compute())
7. 推理
训练完成后，我们可以使用 pipeline 进行推理，生成书评的摘要：
from transformers import pipeline
summarizer = pipeline("summarization", model="Chaklam/test-summ-accelerate")
def print_summary(idx):
review = books_dataset["test"][idx]["review_body"]
title = books_dataset["test"][idx]["review_title"]
summary = summarizer(books_dataset["test"][idx]["review_body"])[0]["summary_text"]
print(f"'>>> Review: {review}'")
print(f"\n'>>> Title: {title}'")
print(f"\n'>>> Summary: {summary}'")
print_summary(100)
这样，我们就可以生成模型的摘要，并与测试集中未见过的书评进行比较。
结语
在本篇博文中，我们介绍了如何通过使用 mT5 Transformer 模型来完成文本摘要任务。我们从多语言亚马逊评论语料库中提取了英文和西班牙文的书籍评论，并将其标题作为摘要任务的参考。接着，通过预处理、使用 mT5 模型进行训练，并利用 ROUGE 分数来评估生成摘要的质量，我们成功构建了一个双语文本摘要系统。
在实践中，文本摘要能够帮助用户快速理解长文档的核心内容，而训练一个双语模型可以扩展其应用范围至多语言环境。随着训练的不断深入和数据集的多样性提升，该模型有望生成更加简洁和准确的摘要。这一流程不仅展示了 Transformer 模型在自然语言生成任务中的强大能力，也为其他多语言任务提供了借鉴和参考。
如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！
欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。
谢谢大家的支持！