【AI】使用 Hugging Face Transformers 进行文本摘要实现
推荐超级课程:
@TOC
我们将构建一个用于总结摘要的模型
文本摘要是 Hugging Face Transformers 提供的一个强大功能。它允许我们从大量文本中生成简洁的摘要。这在处理大量文本数据并需要以简短易读的格式提取核心思想时特别有用。
Hugging Face Transformers 概述
transformers 库是 TensorFlow 2.0 和 PyTorch 上的最先进的自然语言处理 (NLP) 库。它提供了数千个预训练模型来执行文本上的各种任务,如分类、信息提取、摘要、翻译、文本生成等。
利用预训练模型可以显著降低计算成本,节省宝贵的时间和资源,这些资源和时间原本会用于从头开始训练模型。这些模型支持跨不同模态的各种任务,包括自然语言处理 (NLP)、计算机视觉、音频和多模态任务。
在本教程中,我们将使用 google/pegasus-cnn_dailymail
模型。
为什么选择 google/pegasus-cnn_dailymail
模型?
Pegasus 模型是在 CNN/DailyMail 数据集上训练的,专为抽象文本摘要而设计。CNN/DailyMail 数据集是文本摘要任务的流行选择。
这个模型标识符可以使用 Hugging Face Transformers 库提供的 from_pretrained
方法来加载预训练模型。
model = "google/pegasus-cnn_dailymail"
tokenizer = AutoTokenizer.from_pretrained(model)
model_pegasus = AutoModelForSeq2SeqLM.from_pretrained(model).to(device)
先决条件
在开始之前,请确保已安装必要的库。在终端中运行以下命令:
pip install transformers[sentencepiece] datasets sacrebleu rouge_score py7zr -q
pip install --upgrade accelerate
pip uninstall -y transformers accelerate
pip install transformers accelerate
accelerate
库是一个轻量级的解决方案,用于在多种类型的硬件上训练和部署机器学习模型。它提供了一个简单的 API,用于将模型训练卸载到 GPU,并且与 PyTorch 和 TensorFlow 兼容。
本教程需要以下 Python 库:
from transformers import pipeline, set_seed
from datasets import load_dataset, load_from_disk, load_metric
import matplotlib.pyplot as plt
from datasets import load_dataset
import pandas as pd
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import nltk
from nltk.tokenize import sent_tokenize
from tqdm import tqdm
import torch
nltk.download("punkt")
from transformers import DataCollatorForSeq2Seq
from transformers import TrainingArguments, Trainer
load_dataset
:这是 datasets
库提供的一个函数,允许您加载一个数据集。您可以使用它从 Hugging Face Hub 或本地文件加载数据集。
load_from_disk
:这是 datasets
库提供的一个函数,允许您加载已保存到磁盘的数据集。如果您的数据集太大,无法全部加载到内存中,这可能会很有用,因为 datasets
库使用内存映射,允许您在不将整个数据集加载到内存的情况下加载数据集。
AutoModelForSeq2SeqLM
:这用于将一个序列转换为另一个长度相同的序列的任务(例如,翻译、摘要、文本生成)。
AutoTokenizer
:这用于对文本进行标记化。
如果您有 GPU,您可以选择使用 CUDA 来执行训练和评估过程。如果没有 GPU,系统将默认使用 CPU。
device = "cuda" if torch.cuda.is_available() else "cpu"
数据摄取
首先,我们需要从 Hugging Face 下载 SamSum 数据集。您可以使用以下链接下载它:
samsum · Hugging Face 上的数据集
下载后,解压数据并使用 load_from_disk
库加载数据集。记得指定数据集的路径。
dataset_samsum = load_from_disk('samsum_dataset')
dataset_samsum
def download_file():
if not os.path.exists(local_data_file):
filename, headers = request.urlretrieve(
url = source_URL,
filename = local_data_file
)
def extract_zip_file():
unzip_path = unzip_dir
os.makedirs(unzip_path, exist_ok = True)
with zipfile.ZipFile(local_data_file, 'r') as zip_ref:
zip_ref.extractall(unzip_path)
这个数据集用于训练文本摘要模型。它包含以下结构:
DatasetDict({
train: Dataset({
features: ['id', 'dialogue', 'summary'],
num_rows: 14732
})
test: Dataset({
features: ['id', 'dialogue', 'summary'],
num_rows: 819
})
validation: Dataset({
features: ['id', 'dialogue', 'summary'],
num_rows: 818
})
})
每个数据条目包含一个 id
、dialogue
和相应的 summary
。例如:
Dialogue:
Eric: MACHINE!
Rob: That's so gr8!
Eric: I know! And shows how Americans see Russians ;)
Rob: And it's really funny!
Eric: I know! I especially like the train part!
Rob: Hahaha! No one talks to the machine like that!
Eric: Is this his only stand-up?
Rob: Idk. I'll check.
Eric: Sure.
Rob: Turns out no! There are some of his stand-ups on YouTube.
Eric: Gr8! I'll watch
Summary:
Eric and Rob are going to watch a stand-up on YouTube.
数据验证
确保数据集中的所有必需文件都存在。特别是,检查train
、test
和validation
文件是否可用。
def validate_all_files_exist():
try:
validation_status = None
all_files = os.listdir(<数据集的路径>)
for file in all_files:
if file not in ["train", "test", "validation"]:
validation_status = False
else:
validation_status = True
return validation_status
except Exception as e:
raise e
如果validation_status
设置为True
,则可以进入数据转换阶段。否则,需要重新访问和评估数据引入组件。
数据转换
在这个阶段,我们需要转换数据以满足我们的需求。原始数据包含三个列:id
、dialogue
和summary
。然而,在训练期间我们不能直接使用原始文本作为输入。
首先,我们需要对对话和摘要进行标记化。然后,我们使用map
函数将这种转换应用到数据集中。
def convert_examples_to_features(example_batch):
input_encodings = tokenizer(example_batch['dialogue'] , max_length = 1024, truncation = True )
with tokenizer.as_target_tokenizer():
target_encodings = tokenizer(example_batch['summary'], max_length = 128, truncation = True )
return {
'input_ids' : input_encodings['input_ids'],
'attention_mask': input_encodings['attention_mask'],
'labels': target_encodings['input_ids']
}
dataset_samsum_pt = dataset_samsum.map(convert_examples_to_features, batched = True)
因此,转换后的数据集中将添加三个附加列:input_ids
、attention_mask
(用于区分填充标记)和labels
。
Dataset({
features: ['id', 'dialogue', 'summary', 'input_ids', 'attention_mask', 'labels'],
num_rows: 14732
})
例如,预处理数据的显示如下:
{'id': '13728867',
'dialogue': 'Olivia: Who are you voting for in this election? rnOliver: Liberals as always.rnOlivia: Me too!!rnOliver: Great',
'summary': 'Olivia and Olivier are voting for liberals in this election. ',
'input_ids': [18038,151,2632,127,119,6228,118,115,136,2974,152,10463,151,35884,130,329,107,18038,151,2587,314, 1242,10463,151,1509,1],
'attention_mask': [1,1,1,1,1,1,1,1,1,1,1,1, 1,1,1,1,1,1,1,1,1,1,1,1,1,1],
'labels': [18038, 111, 34296, 127, 6228, 118, 33195, 115, 136, 2974, 107, 1]}
模型训练
现在,我们可以开始训练过程。最初,我们使用 Hugging Face Transformers 的DataCollatorForSeq2Seq
库来收集适合序列到序列模型的序列批次。我们创建一个特定于标记化和模型的DataCollatorForSeq2Seq
实例。标记化器用于处理输入序列的标记化,模型用于确定填充标记的 ID。
from transformers import DataCollatorForSeq2Seq
seq2seq_data_collator = DataCollatorForSeq2Seq(tokenizer, model=model_pegasus)
Afterward, we provide the training arguments and start the training process.
from transformers import TrainingArguments, Trainer
dataset_samsum_pt = load_from_disk(self.config.data_path)
trainer_args = TrainingArguments(
output_dir='pegasus-samsum',
num_train_epochs=1,
warmup_steps=500,
per_device_train_batch_size=1,
per_device_eval_batch_size=1,
weight_decay=0.01,
logging_steps=10,
evaluation_strategy='steps',
eval_steps=500,
save_steps=1e6,
gradient_accumulation_steps=16
)
trainer = Trainer(
model=model_pegasus,
args=trainer_args,
tokenizer=tokenizer,
data_collator=seq2seq_data_collator,
train_dataset=dataset_samsum_pt["train"],
eval_dataset=dataset_samsum_pt["validation"])
trainer.train()
## Save model
model_pegasus.save_pretrained(os.path.join(self.config.root_dir,"pegasus-samsum-model"))
## Save tokenizer
tokenizer.save_pretrained(os.path.join(self.config.root_dir,"tokenizer"))
模型评估
一旦训练过程完成,我们需要评估模型并检查它总结对话的效果。虽然 ROUGE 和 BLEU 分数通常用于机器翻译和摘要,但本教程将使用 ROUGE 分数。
BLEU (Bilingual Evaluation Understudy) Score 是在机器翻译任务中广泛使用的突出指标。BLEU 分数注重精确度。这些任务的主要目标是自动将文本从一种语言翻译成另一种语言。BLEU 分数量化了机器翻译的文本与参考翻译之间的相似度。这种测量使用n-gram,即连续的“n”个单词序列。 ROUGE (Recall-Oriented Understudy for Gisting Evaluation) Score 是用于评估自动文本摘要和机器翻译的一套指标。**ROUGE 分数注重召回率。**它将自动生成的摘要或翻译与参考或一组参考进行比较。ROUGE 分数范围从 0 到 1,反映了机器生成的摘要与参考之间的相似度,分数越高表示相似度越高。
def generate_batch_sized_chunks(self, list_of_elements, batch_size):
"""
将数据集分割成我们可以同时处理的小批次
从 list_of_elements 中产生连续的批次大小的块。
"""
for i in range(0, len(list_of_elements), batch_size):
yield list_of_elements[i : i + batch_size]
def calculate_metric_on_test_ds(self, dataset, metric, model, tokenizer,
batch_size=16, device="cuda" if torch.cuda.is_available() else "cpu",
column_text="article",
column_summary="highlights"):
article_batches = list(self.generate_batch_sized_chunks(dataset[column_text], batch_size))
target_batches = list(self.generate_batch_sized_chunks(dataset[column_summary], batch_size))
for article_batch, target_batch in tqdm(
zip(article_batches, target_batches), total=len(article_batches)):
inputs = tokenizer(article_batch, max_length=1024, truncation=True,
padding="max_length", return_tensors="pt")
summaries = model.generate(input_ids=inputs["input_ids"].to(device),
attention_mask=inputs["attention_mask"].to(device),
length_penalty=0.8, num_beams=8, max_length=128)
''' 参数 length_penalty 确保模型不会生成过长的序列。 '''
# 最后,我们解码生成的文本,
# 替换填充标记,并将解码的文本与参考添加到指标中。
decoded_summaries = [tokenizer.decode(s, skip_special_tokens=True,
clean_up_tokenization_spaces=True)
for s in summaries]
decoded_summaries = [d.replace("", " ") for d in decoded_summaries]
metric.add_batch(predictions=decoded_summaries, references=target_batch)
# 最后计算并返回 ROUGE 分数。
score = metric.compute()
return score
def evaluation(self):
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(self.config.tokenizer_path)
model_pegasus = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path).to(device)
# 加载数据
dataset_samsum_pt = load_from_disk(self.config.data_path)
rouge_names = ["rouge1", "rouge2", "rougeL", "rougeLsum"]
rouge_metric = load_metric('rouge')
score = self.calculate_metric_on_test_ds(
dataset_samsum_pt['test'][0:10], rouge_metric, model_pegasus, tokenizer, batch_size = 2, column_text = 'dialogue', column_summary= 'summary')
rouge_dict = dict((rn, score[rn].mid.fmeasure ) for rn in rouge_names )
return rouge_dict
预测
完成整个训练和评估过程后,我们可以使用我们的模型来总结随机的对话。
from transformers import AutoTokenizer
from transformers import pipeline
def predict(text):
tokenizer = AutoTokenizer.from_pretrained(<tokenizer_path>)
model = <model_path>
kwargs = {"length_penalty": 0.8, "num_beams":8, "max_length": 128}
Pipeline = pipeline("summarization", model= model, tokenizer=tokenizer)
summary = Pipeline(text, **kwargs)[0]["summary_text"]
return summary
pipeline
:这是transformers
库提供的一个高级函数,允许您使用预训练模型轻松生成预测。您可以使用它进行文本分类、命名实体识别和文本生成等任务。
结论
在本教程中,我们探讨了使用 Hugging Face Transformers 进行文本摘要,特别是google/pegasus-cnn_dailymail
模型。我们逐步介绍了数据准备、模型训练和评估的过程。在当今信息密集的世界中,文本摘要的能力是一种强大的工具,其潜在应用非常广泛。