【AI】使用 Hugging Face Transformers 进行文本摘要实现

2025-04-14T13:39:48+08:00 | 6分钟阅读 | 更新于 2025-04-14T13:39:48+08:00

Macro Zhao

【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  
    })  
})

每个数据条目包含一个 iddialogue 和相应的 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.

数据验证

确保数据集中的所有必需文件都存在。特别是,检查traintestvalidation文件是否可用。

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,则可以进入数据转换阶段。否则,需要重新访问和评估数据引入组件。

数据转换

在这个阶段,我们需要转换数据以满足我们的需求。原始数据包含三个列:iddialoguesummary。然而,在训练期间我们不能直接使用原始文本作为输入。 首先,我们需要对对话和摘要进行标记化。然后,我们使用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_idsattention_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模型。我们逐步介绍了数据准备、模型训练和评估的过程。在当今信息密集的世界中,文本摘要的能力是一种强大的工具,其潜在应用非常广泛。

© 2011 - 2025 Macro Zhao的分享站

关于我

如遇到加载502错误,请尝试刷新😄

Hi,欢迎访问 Macro Zhao 的博客。Macro Zhao(或 Macro)是我在互联网上经常使用的名字。

我是一个热衷于技术探索和分享的IT工程师,在这里我会记录分享一些关于技术、工作和生活上的事情。

我的CSDN博客:
https://macro-zhao.blog.csdn.net/

欢迎你通过评论或者邮件与我交流。
Mail Me

推荐好玩(You'll Like)
  • AI 动·画

    • 这是一款有趣·免费的能让您画的画中的角色动起来的AI工具。
    • 支持几十种动作生成。
  • AI 识字

    • 遇到不认识的字,写在这里,会自动识别所写汉字。
    • 还能对其进行意义查询。
  • 中日假期日历

    • 方便查询2025日本红日子。
    • 对日打工族必备工具。
  • 在线架子鼓

    • 简易但功能齐全的架子鼓。
    • 代码敲累了,你就敲一敲它吧。
  • 微信公众号编辑器

    • 简易但功能丰富的免费!公众号编辑器。
    • 还在不断完善中,喜欢的可以持续关注.
我的项目(My Projects)
  • 爱学习网

  • 小乙日语App

    • 这是一个帮助日语学习者学习日语的App。
      (当然初衷也是为了自用😄)
    • 界面干净,简洁,漂亮!
    • 其中包含 N1 + N2 的全部单词和语法。
    • 不需注册,更不需要订阅!完全免费!
  • 小乙日文阅读器

    • 词汇不够?照样能读日语名著!
    • 越读积累越多,积跬步致千里!
    • 哪里不会点哪里!妈妈再也不担心我读不了原版读物了!
赞助我(Sponsor Me)

如果你喜欢我的作品或者发现它们对你有所帮助,可以考虑给我买一杯咖啡 ☕️。这将激励我在未来创作和分享更多的项目和技术。🦾

👉 请我喝一杯咖啡

If you like my works or find them helpful, please consider buying me a cup of coffee ☕️. It inspires me to create and share more projects in the future. 🦾

👉 Buy me a coffee