使用 T5 进行总结

简介

T5 (Text-To-Text Transfer Transformer) 是一种 Transformer-based 的预训练模型，由谷歌研发并于 2019 年发布，是一个多任务的通用文本与序列生成模型。通过将各种序列到序列的任务转换为单一的序列到序列问题，可以在不同数据集上进行预训练和微调，可以为多种 NLP 任务带来显著的收益。

安装

T5 可以通过 Hugging Face 的 transformers Python 库进行安装：

pip install transformers

训练与微调

数据准备

在进行训练之前，需要将数据转换为适合 T5 的格式。T5 使用将输入和输出拼接在一起的方式进行训练。例如，我们将要将一句话进行翻译，那么数据需要转换成：

"translate German to English: Wie heißt du?\tWhat's your name?"

其中，"\t" 前的部分是输入，"\t" 后的部分是输出。

预处理

将数据准备完后，需要使用 T5Tokenizer 进行预处理。T5Tokenizer 是为 T5 模型设计的特殊 tokenizer，能够将输入和输出进行拼接，并将其编码为数字 ID。

from transformers import T5Tokenizer

model_name = 't5-small'
tokenizer = T5Tokenizer.from_pretrained(model_name)

input_text = "translate German to English: Wie heißt du?"
output_text = "What's your name?"

# tokenization
inputs = tokenizer(input_text, return_tensors='pt', padding=True)
outputs = tokenizer(output_text, return_tensors='pt', padding=True)

# concatenation
input_ids = inputs['input_ids'].squeeze()
output_ids = outputs['input_ids'].squeeze()

src_text = tokenizer.decode(input_ids)
tgt_text = tokenizer.decode(output_ids)
print(f'src_text: {src_text}')
print(f'tgt_text: {tgt_text}')

定义模型

定义 T5 模型需要指定模型的名称和预训练的 checkpoint 路径。

from transformers import T5ForConditionalGeneration

model_name = 't5-small'
model = T5ForConditionalGeneration.from_pretrained(model_name)

# move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

训练模型

定义好模型后，我们需要使用使用 Sequence-to-Sequence 框架进行训练。在训练之前，需要设置训练参数（例如 batch size, learning rate, epoch），并定义损失函数和优化器。此处略过。

最后，使用训练数据进行模型训练。

for epoch in range(epochs):
  for i, (input_ids, target_ids) in enumerate(dataloader):
    input_ids = input_ids.to(device)
    target_ids = target_ids.to(device)

    output = model(input_ids=input_ids, labels=target_ids)
    loss, logits = output['loss'], output['logits']

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

微调模型

微调过程和训练过程类似，只需要更换数据集即可。在进行微调之前，需要先下载预训练的 checkpoint。

from transformers import T5ForConditionalGeneration, T5Tokenizer

checkpoint_path = 't5-base-checkpoints/'
model_name = 't5-base'

tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(checkpoint_path)

# move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

微调过程中，需要指定微调数据的格式和路径，并进行预处理。T5 会在训练过程中动态调整自己的词典，所以需要重新进行编码。

from torch.utils.data import DataLoader, Dataset

class CustomDataset(Dataset):
    def __init__(self, file_path, tokenizer):
        self.examples = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                src, tgt = line.strip().split('\t')
                self.examples.append((src, tgt))
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        src, tgt = self.examples[idx]
        input_ids = tokenizer.encode(src, add_special_tokens=True, padding=True)
        target_ids = tokenizer.encode(tgt, add_special_tokens=True, padding=True)
        return torch.tensor(input_ids), torch.tensor(target_ids)

train_dataset = CustomDataset('train.tsv', tokenizer)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
criterion = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

然后，使用微调数据进行微调。微调后，可以使用模型进行测试。

for epoch in range(epochs):
  for i, (input_ids, target_ids) in enumerate(train_loader):
    input_ids = input_ids.to(device)
    target_ids = target_ids.to(device)

    output = model(input_ids=input_ids, labels=target_ids)
    loss, logits = output['loss'], output['logits']

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

model.eval()
input_text = "translate German to English: Wie heißt du?"
input_ids = tokenizer.encode(input_text, add_special_tokens=True, return_tensors='pt').to(device)

outputs = model.generate(input_ids=input_ids)
output_text = tokenizer.decode(outputs[0])
print(f'output_text: {output_text}')

总结

T5 是一个多任务的通用文本与序列生成模型，在各种序列到序列的任务上表现非常优秀。我们可以通过训练和微调来使用 T5 模型，它会自动调节自身的词典以适应不同的任务需求。它现在已经得到了很多机器翻译、摘要、问题回答等任务的广泛应用，并且在模型准确度、效率和可扩展性等方面都具有非常大的优势。