Chapter 21
基于T5的文本摘要
Notebooktransformers29 cells
基于T5的文本摘要
Step1 导入相关包
In [ ]python · cell 3
python
import torch
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArgumentsStep2 加载数据集
In [ ]python · cell 5
python
ds = Dataset.load_from_disk("./nlpcc_2017/")
dsIn [ ]python · cell 6
python
ds = ds.train_test_split(100, seed=42)
dsIn [ ]python · cell 7
python
ds["train"][0]Step3 数据处理
In [ ]python · cell 9
python
tokenizer = AutoTokenizer.from_pretrained("Langboat/mengzi-t5-base")In [ ]python · cell 10
python
def process_func(exmaples):
contents = ["摘要生成: \n" + e for e in exmaples["content"]]
inputs = tokenizer(contents, max_length=384, truncation=True)
labels = tokenizer(text_target=exmaples["title"], max_length=64, truncation=True)
inputs["labels"] = labels["input_ids"]
return inputsIn [ ]python · cell 11
python
tokenized_ds = ds.map(process_func, batched=True)
tokenized_dsIn [ ]python · cell 12
python
tokenizer.decode(tokenized_ds["train"][0]["input_ids"])In [ ]python · cell 13
python
tokenizer.decode(tokenized_ds["train"][0]["labels"])Step4 创建模型
In [ ]python · cell 15
python
model = AutoModelForSeq2SeqLM.from_pretrained("Langboat/mengzi-t5-base")Step5 创建评估函数
In [ ]python · cell 17
python
import numpy as np
from rouge_chinese import Rouge
rouge = Rouge()
def compute_metric(evalPred):
predictions, labels = evalPred
decode_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
decode_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
decode_preds = [" ".join(p) for p in decode_preds]
decode_labels = [" ".join(l) for l in decode_labels]
scores = rouge.get_scores(decode_preds, decode_labels, avg=True)
return {
"rouge-1": scores["rouge-1"]["f"],
"rouge-2": scores["rouge-2"]["f"],
"rouge-l": scores["rouge-l"]["f"],
}Step6 配置训练参数
In [ ]python · cell 19
python
args = Seq2SeqTrainingArguments(
output_dir="./summary",
per_device_train_batch_size=4,
per_device_eval_batch_size=8,
gradient_accumulation_steps=8,
logging_steps=8,
eval_strategy="epoch",
save_strategy="epoch",
metric_for_best_model="rouge-l",
predict_with_generate=True
)Step7 创建训练器
In [ ]python · cell 21
python
trainer = Seq2SeqTrainer(
args=args,
model=model,
train_dataset=tokenized_ds["train"],
eval_dataset=tokenized_ds["test"],
compute_metrics=compute_metric,
tokenizer=tokenizer,
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer)
)Step8 模型训练
In [ ]python · cell 23
python
trainer.train()Step9 模型推理
In [ ]python · cell 25
python
from transformers import pipelineIn [ ]python · cell 26
python
pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=0)In [ ]python · cell 27
python
pipe("摘要生成:\n" + ds["test"][-1]["content"], max_length=64, do_sample=True)In [ ]python · cell 28
python
ds["test"][-1]["title"]In [ ]python · cell 29
python