---
title: "Transformer 学习之路 - 基于前缀模型的文本摘要技术"
date: 2024-04-19T17:35:18+08:00
draft: false
description: "深入解析基于前缀模型的文本摘要技术，结合代码示例，详细分析其技术原理与应用场景。"
categories: ["Python", "Transformer"]
---

# Transformer 学习之路 - 基于前缀模型的文本摘要技术

在自然语言处理（NLP）领域，文本摘要是一项重要的任务，它旨在从长篇文章中提取出关键信息，生成简洁的摘要。随着 Transformer 模型的兴起，基于前缀模型（Prefix Model）的文本摘要技术逐渐成为研究热点。本文将深入探讨前缀模型的原理、应用场景及其实现细节，并结合代码示例帮助读者理解如何应用这一技术。

## 前缀模型的基本概念

前缀模型是一种特殊的文本生成模型，通过在输入序列前添加“前缀”来引导模型生成符合特定需求的输出。其核心思想是通过给输入文本添加适当的“前缀”来调整模型的生成方向。例如，如果想要生成一个故事摘要，我们可以添加“摘要:”作为前缀，这样模型会理解它应该输出的是摘要而不是继续讲述故事。

### 前缀模型的典型应用场景

前缀模型在多种生成式任务中都有广泛应用，包括但不限于：

- **文本摘要**：在输入文本前加上“总结：”或“摘要：”这样的前缀，引导模型输出总结性的句子。
- **机器翻译**：在输入前添加特定语言标签，例如在英文文本前添加“[Translate to French]”，模型将生成法语的翻译。
- **对话生成**：在输入前加入角色标签或话题前缀，模型可以生成具有特定风格的对话。
- **代码生成**：添加“生成代码：”等指示语，使模型知道要生成符合指令的代码。

## 前缀模型的结构和实现

前缀模型通常基于序列到序列（Seq2Seq）架构，前缀信息可以直接添加到输入序列中。以文本摘要为例，其实现步骤如下：

1. **定义前缀**：在输入文章前添加“摘要：”，即构造输入`“摘要：[原文内容]”`。
2. **编码阶段**：
   - 编码器会处理整个带前缀的输入序列，把“摘要：”这一部分信息编码为隐藏状态。
   - 编码器将这个包含“摘要”提示的隐藏状态传递给解码器。
3. **解码阶段**：
   - 解码器根据编码器的输出，结合“摘要：”前缀的暗示，从而生成一段总结性内容。

### 实现示例

假设我们使用 Hugging Face 提供的 `transformers` 库的一个预训练模型来实现前缀模型，以下是一个简单的代码示例：

```python
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 加载模型和分词器，例如BART或T5
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")

# 定义输入内容和前缀
text = "The weather forecast indicates heavy rain this weekend in several regions."
prefix = "Summarize: "  # 用于指引模型生成摘要

# 在输入文本前添加前缀
input_text = prefix + text

# 编码输入
inputs = tokenizer(input_text, return_tensors="pt")

# 生成摘要
summary_ids = model.generate(inputs["input_ids"], max_length=50, num_beams=5, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

print("Generated Summary:", summary)

from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments

ds = Dataset.load_from_disk("/content/drive/MyDrive/ai-learning/2.NLP Task/08-text_summarization/nlpcc_2017/")
ds = ds.train_test_split(100, seed=42)

tokenizer = AutoTokenizer.from_pretrained("THUDM/glm-large-chinese", trust_remote_code=True)

def process_func(examples):
    contents = ["摘要生成: \n" + e + tokenizer.mask_token for e in examples["content"]]
    inputs = tokenizer(contents, max_length=384, truncation=True, padding="max_length", return_tensors="pt")
    inputs = tokenizer.build_inputs_for_generation(inputs, targets=examples['title'], padding=True, max_gen_length=64)
    return inputs

tokenized_ds = ds.map(process_func, batched=True, remove_columns=ds["train"].column_names)

model = AutoModelForSeq2SeqLM.from_pretrained("THUDM/glm-large-chinese", trust_remote_code=True)

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"],
    }

import os

os.environ["WANDB_DISABLED"] = "true"
args = Seq2SeqTrainingArguments(
    output_dir="./summary_glm",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=8,
    logging_steps=8,
    num_train_epochs=1
)

trainer = Seq2SeqTrainer(
    args=args,
    model=model,
    train_dataset=tokenized_ds["train"],
    tokenizer=tokenizer,
)

trainer.train()

input_text = ds["test"][-1]["content"]
inputs = tokenizer("摘要生成: \n" + input_text + tokenizer.mask_token, return_tensors="pt")
inputs = tokenizer.build_inputs_for_generation(inputs, max_gen_length=64)
inputs = inputs.to("cuda")
output = model.generate(**inputs, max_new_tokens=64, eos_token_id=tokenizer.eop_token_id, do_sample=True)
tokenizer.decode(output[0].tolist())

model.save_pretrained("/content/drive/MyDrive/ai-learning/2.NLP Task/08-text_summarization/summary_glm")