用huggingface.transformers.AutoModelForSeq2SeqLM实现文本生成任务

诸神缄默不语

已于 2024-03-13 16:26:05 修改

阅读量3.1k

点赞数 26

CC 4.0 BY-SA版权

分类专栏：人工智能学习笔记文章标签：人工智能算法深度学习自然语言处理文本生成 transformers seq2seq

于 2024-03-07 14:39:56 首次发布

本文链接：https://blog.youkuaiyun.com/PolarisRisingWar/article/details/136531105

人工智能学习笔记专栏收录该内容

271 篇文章

订阅专栏

本文详细介绍了如何使用BART模型和transformers库的Seq2SeqTrainer进行文本生成，包括设置参数、数据预处理、训练过程以及生成和推理方法的演示。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

诸神缄默不语-个人优快云博文目录

本文介绍的是用BART（别的s2s模型也差不多）+ transformers官方训练框架（Seq2SeqTrainer）运行文本生成代码，先介绍前置条件，然后介绍训练代码，最后会介绍使用generate()函数或者使用pipeline来进行推理的方法。

1. 导入包

import json,random,copy,os,re

from transformers import AutoTokenizer,AutoModelForSeq2SeqLM,DataCollatorForSeq2Seq,Seq2SeqTrainingArguments,\
                        Seq2SeqTrainer,EarlyStoppingCallback

from datasets import Dataset,DatasetDict

2. 设置参数

model_path="bart-base-chinese"
output_path="data/my_checkpoints/bart_202403071330"
source_max_length=1024
target_max_length=128  #这2个都是模型限长

3. 导入分词器和模型

tokenizer=AutoTokenizer.from_pretrained(model_path)
model=AutoModelForSeq2SeqLM.from_pretrained(model_path)

4. 导入数据集

具体部分略，总之最后要是datasets.Dataset的格式，输入键是text，输出键是summary
然后进行如下数据预处理工作：

def preprocess_function(examples):
    model_inputs = tokenizer(examples['text'],max_length=source_max_length,truncation=True)

    labels = tokenizer(text_target=examples["summary"],max_length=target_max_length,truncation=True)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_dataset = datasets.map(preprocess_function, batched=True)

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

datasets包的使用文档可参考huggingface.datasets使用说明

5. 训练

training_args = Seq2SeqTrainingArguments(
    output_dir=output_path,
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=32,
    weight_decay=0.01,
    save_strategy="epoch",
    save_steps=1,
    save_total_limit=3,
    num_train_epochs=100,
    predict_with_generate=True,
    fp16=True,
    push_to_hub=False,
    metric_for_best_model="eval_loss",
    load_best_model_at_end=True
)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["valid"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    callbacks=[EarlyStoppingCallback(10)],
)

trainer.train()

6. 测试

在output_path中会保留save_total_limit个checkpoint，一般来说建议用checkpoint数最大的那个文件夹。这个文件夹称为final_checkpoint_folder，从这个文件夹里面调用模型权重：model=AutoModelForSeq2SeqLM.from_pretrained(final_checkpoint_folder)
以下代码中的model指的是这个model

6.1 使用`generate()`函数

函数文档：https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.generate

outputs=model.generate(tokenizer.encode(input_text,return_tensors="pt",max_length=source_max_length),max_new_tokens=target_max_length,do_sample=False)
output_text=tokenizer.decode(outputs[0],skip_special_tokens=True).replace(" ","")

（这个会在CPU上跑，如果把模型和数据都放到GPU上就能在GPU上跑）

之所以要删空格是因为BART输出会在token之间自动添加空格，就跟英文一样。

generate()函数的使用可以参考这2篇博文：基于 transformers 的 generate() 方法实现多样化文本生成：参数含义和算法原理解读_length_penalty-优快云博客和 LLM（大语言模型）解码时是怎么生成文本的？ - 知乎

6.2 使用pipeline

函数源码：https://github.com/huggingface/transformers/blob/main/src/transformers/pipelines/text2text_generation.py

函数文档：https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.Text2TextGenerationPipeline

需要导入pipeline函数：

from transformers import pipeline

generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
outputs = generator(input_text, max_new_tokens=128, do_sample=False)
output_text = outputs[0]["generated_text"].replace(" ", "")