Huggingface T5模型代码笔记

0 前言

本博客主要记录如何使用T5模型在自己的Seq2seq模型上进行Fine-tune。

1 文档介绍

本文档介绍来源于Huggingface官方文档，参考T5。

1.1 概述

T5模型是由Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu.在论文 Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer中提出的。

该论文摘要如下：

迁移学习在自然语言处理(NLP)中已经成为一种强大的技术。迁移学习是指，模型首先在数据丰富的任务上进行预训练，然后再对下游任务进行微调。迁移学习的有效性引起了不同的方法、方法和实践。在本文中，我们通过引入一个统一的框架，将每个语言问题转换成文本到文本的格式，来探索自然语言处理的迁移学习技术。我们的系统研究比较了数十项语言理解任务的训练前目标、架构、未标记数据集、迁移方法和其他因素。通过将我们的探索与规模和我们的新"Colossal Clean Crawled Corpus"相结合，我们在总结、问答、文本分类等许多基准测试中取得了最先进的结果。为了促进NLP迁移学习的未来工作，我们发布了我们的数据集、预训练模型和代码。

提示：

• T5是一种编码器-解码器模型，在非常多的无监督和有监督任务上进行多任务训练，并将每个任务转换为文本到文本格式。T5可以很好地处理各种任务，通过在每个任务对应的输入前加上不同的前缀，例如:translate English to German: ...，摘要:summarize: ....（注：这里是指输入的文本可以通过加入特定前缀的方式指定是那种text-to-text任务，比如你想让T5做翻译任务，那么给他的输入就可以是"translate English to German: What is your name?"）。关于使用哪个前缀的更多信息，原论文的附录D给出了所有的前缀。对于序列到序列的生成，建议使用generate()函数。这种方法负责通过cross-attention层将编码输入到解码器，并自动回归生成解码器输出。T5使用相对标量嵌入。编码器输入填充可以在左边和右边完成。

1.2 训练

T5是一个编码器-解码器模型，并将所有NLP问题转换为文本到文本的形式。它是通过teacher forcing（如果不熟悉，可以参考What is Teacher Forcing for Recurrent Neural Networks?）的方式进行训练的。这意味着对于训练，我们总是需要一个输入序列和一个目标序列。输入序列通过input_ids喂给模型的Encoder。目标序列在其右边，即跟在一个start-sequence token之后，通过decoder_input_ids喂给模型的Decoder。通过teacher forcing的方式，目标序列最后会加一个EOS token（表示结束End Of Sentence），并作为标签。T5可以在监督和非监督的方式下进行训练/微调。

1.2.1 无监督去噪训练

在该设置下，输入序列的范围被所谓的哨点标记(sentinel tokens，也就是唯一的掩码标记)屏蔽，而输出序列则由相同的哨点标记和真实掩码标记的串联组成。每个哨兵标记代表这个句子的唯一掩码标记，应该从<extra_id_0>, <extra_id_1>, ... 直到<extra_id_99>。默认情况下，T5 Tokenizer中有100个sentinel token可用。

例如，句子"the cute dog walks in the park" 中掩盖"cute dog"和"the"应按如下方式处理：

from transformers import T5ForConditionalGeneration, T5Tokenizer
model = T5ForConditionalGeneration.from_pretrained("t5-small")
tokenizer = T5Tokenizer.from_pretrained("t5-small")

input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids
labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2>', return_tensors='pt').input_ids
# the forward function automatically creates the correct decoder_input_ids
loss = model(input_ids=input_ids, labels=labels).loss

1.2.2 有监督训练

在该设置下，输入序列和输出序列是标准的序列到序列的输入输出映射。在翻译中，例如输入序列"the house is wonderful."和输出序列"Das Haus ist wunderbar"，这些句子应按下列方法处理：

from transformers import T5ForConditionalGeneration, T5Tokenizer
model = T5ForConditionalGeneration.from_pretrained("t5-small")
tokenizer = T5Tokenizer.from_pretrained("t5-small")

input_ids = tokenizer('translate English to German: The house is wonderful.', return_tensors='pt').input_ids
labels = tokenizer('Das Haus ist wunderbar.', return_tensors='pt').input_ids
# the forward function automatically creates the correct decoder_input_ids
loss = model(input_ids=input_ids, labels=labels).loss

Tokenizer中的return_tensors的作用

若设置return_tensors="pt"，则返回为tensor，否则为正常list。

inputs = tokenizer("translate English to Spanish: Find all student's name in students table.")
print(inputs)
{'input_ids': [13959, 1566, 12, 5093, 10, 2588, 66, 1236, 31, 7, 564, 16, 481, 953, 5, 1], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}
inputs = tokenizer("translate English to Spanish: Find all student's name in students table.", return_tensors="pt")
print(inputs)
{'input_ids': tensor([[13959,  1566,    12,  5093,    10,  2588,    66,  1236,    31,     7,
           564,    16,   481,   953,     5,     1]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}

2 一个简单T5翻译示例

以下内容涉及一些Pytorch搭建神经网络模型的基础知识，不了解的朋友可以参考这个系列：莫凡Pytorch学习笔记（一）

2.1 模型搭建

我们借助T5ForConditionalGeneration模型来搭建一个简单的由文本到SQL语句的翻译模型。其Pytorch实现如下：

class T5ForTextToSQL(torch.nn.Module):
    '''
    A basic T5 model for Text-to-SQL task.
    '''
    def __init__(self):
        super(T5ForTextToSQL, self).__init__()
        self.t5 = T5ForConditionalGeneration.from_pretrained('t5-small')
        
    def forward(self, input_ids, labels):
        out = self.t5(input_ids=input_ids, labels=labels)
        
        return out
    
    def generate(self, input_ids):
        result = self.t5.generate(input_ids=input_ids)
        
        return result

该类共三个方法。其中__init__和forward方法是必须要实现的。这里的generate也是原模型的方法，我们用于验证。

注意：：这个类其实就是一个非常简单的T5ForConditionalGeneration类，但是我们这里把它封装为一个自己的类的目的主要有以下几点：

1. 方便后续拓展，例如增加一些自定义的模块
2. 仅保留T5ForConditionalGeneration的forward和generate函数给用户使用

2.2 数据集

这里我们使用一个自定义的数据集来快速开始。其定义如下：

class TextToSQL_Dataset(torch.utils.data.Dataset):
    '''
    A simple text-to-sql dataset example.
    '''
    def __init__(self, text_l, schema_l, sql_l, tokenizer, block_size=1):
        self.tokenizer = tokenizer
        self.max_len = block_size
        self.text = text_l
        self.scheme = schema_l
        self.sql = sql_l
        
        
    def _text_to_encoding(self, item):
        return self.tokenizer(item)
 

    def _text_to_item(self, text):
        try:
            if (text is not None):
                return self._text_to_encoding(text)
            else:
                return None
        except:
            return None
        
        
    def __len__(self):
        return len(self.sql)
    
    
    def __getitem__(self, _id):
        text = self.text[_id]
        sql = self.sql[_id]
        schema = self.scheme[_id]
        text_encodings = self._text_to_item("translate Text to SQL: " + text)
        sql_encodings = self._text_to_item(sql)
        schema_encodings = self._text_to_item(schema)
        item = dict()
        item['text_encodings'] = {key: torch.tensor(value) for key, value in text_encodings.items()}
        item['sql_encodings'] = {key: torch.tensor(value) for key, value in sql_encodings.items()}
        item['schema_encodings'] = {key: torch.tensor(value) for key, value in schema_encodings.items()}
        
        return item

我们的主要目的是跑通整个Fine-tune T5到Text-to-SQL任务上，所以为了不浪费更多的时间在构建数据集上面，这里我自已编写了一些自然语言文本到SQL语句的对应，用来快速开始。

其中，训练集和测试集如下：

# 以下为train_set

text_l = [
    "Find all student names in student database.",
    "Count student's number for class 1. ",
    "Given the max student age in class 1.",
    "Please find the minium student age in class 1.",
    "Tell me the number of classes.",
    "Who is the student that older than 15."
]
schema_l = [
    'Table: student$$header: name%%age%%class%%',
]*len(text_l)
sql_l = [
    "SELECT name FROM student",
    "SELECT COUNT(*) FROM student WHERE class=1",
    "SELECT MAX(age) FROM student WHERE class=1",
    "SELECT MIN(age) FROM student WHERE class=1",
    "SELECT COUNT(class) FROM student",
    "SELECT name FROM student WHERE age>15",
]

# 以下为test_set
test_text_l = [
    "Find all student ages in student database.",
    "Count student's number for class 3. ",
    "Given the min student age in class 2.",
    "Please find the maxium student age in class 2.",
    "Who is the student that younger than 14."
]
test_schema_l = [
    'Table: student$$header: name%%age%%class%%',
]*len(text_l)
test_sql_l = [
    "SELECT age FROM student",
    "SELECT COUNT(*) FROM student WHERE class=3",
    "SELECT MIN(age) FROM student WHERE class=2",
    "SELECT MAX(age) FROM student WHERE class=2",
    "SELECT name FROM student WHERE age<14",
]

这里的几个文本-SQL语句对都比较简单，只是为了用于测试流程完整性。

在构建好数据集文本后，我们定义和生成dataset以及dataloader对象：

train_dataset = TextToSQL_Dataset(text_l, schema_l, sql_l, tokenizer)
test_dataset = TextToSQL_Dataset(test_text_l, test_schema_l, test_sql_l, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True)

2.3 模型初始化

首先，初始化一个简单的模型：

model = T5ForTextToSQL()

为了验证训练的有效性，我们先来看看一个最初的不经过Fine-tune的原始T5模型在Text-to-SQL任务上的表现。

这里我们直接使用测试集进行评估测试，主要代码就是调用其generate函数来生成输出的Sequence并与ground truth进行对比。

device = torch.device('cuda:3') if torch.cuda.is_available() else torch.device('cpu')
model.eval()
model = model.to(device)
for i,batch in enumerate(test_loader):
    input_ids = batch['text_encodings']['input_ids'].to(device)
    sql_ids = batch['sql_encodings']['input_ids'].to(device)
    result = model.generate(input_ids)
    print("==="*20)
    print("Question:")
    print(tokenizer.decode(input_ids[0]))
    print("SQL:")
    print(tokenizer.decode(result[0]))
    print()

运行以上代码，得到输出如下：

============================================================
Question:
translate Text to SQL: Who is the student that younger than 14.</s>
SQL:
<pad> Text zu SQL: Wer ist der Student, der unter 14 Jahren ist?</s>

============================================================
Question:
translate Text to SQL: Find all student ages in student database.</s>
SQL:
<pad> Text in SQL: Finden Sie alle Studenten Alter in der Studentendatenbank.</s>

============================================================
Question:
translate Text to SQL: Given the min student age in class 2.</s>
SQL:
<pad> Text zu SQL: Angesichts des min Studenten Alters in der Klasse 2

============================================================
Question:
translate Text to SQL: Count student's number for class 3.</s>
SQL:
<pad> Text in SQL: Count student's number for class 3.</s>

============================================================
Question:
translate Text to SQL: Please find the maxium student age in class 2.</s>
SQL:
<pad> Text in SQL: Bitte finden Sie das maxium Studentenalter in der Klasse

可以看到，此时模型的输出，比如"Text zu SQL: Wer ist der Student, der unter 14 Jahren ist?"或"Text in SQL: Bitte finden Sie das maxium Studentenalter in der Klasse"好像是一些英语翻译成德语/法语的结果输入，并不能够输出SQL语句（无论是形式和内容，一点儿都不像）。

接下来，我们对这个简单的模型进行训练。

2.4 模型的训练

在进行模型训练时，我们首先要设置好优化器。这里我们采用T5在预训练时使用的AdamW又优化器，学习率设置为5e-5（这些超参数后面都可以仔细精调）。

optim = AdamW(model.parameters(), lr=5e-5)

device = torch.device('cuda:3') if torch.cuda.is_available() else torch.device('cpu')
model.train()
model = model.to(device)

for epoch in range(100):
    for i,batch in enumerate(train_loader):
        optim.zero_grad()
        input_ids = batch['text_encodings']['input_ids'].to(device)
        sql_ids = batch['sql_encodings']['input_ids'].to(device)
        loss = model(input_ids=input_ids, labels=sql_ids).loss
        loss.backward()
        optim.step()
        if epoch % 10 == 0 and i % 10 == 0:
            print("Epoch: ", epoch, " , step: ", i)
            print("training loss: ", loss.item())

运行以上代码，程序的输出如下：

Epoch:  0  , step:  0
training loss:  5.786584854125977
Epoch:  10  , step:  0
training loss:  3.1280531883239746
Epoch:  20  , step:  0
training loss:  1.795115351676941
Epoch:  30  , step:  0
training loss:  0.7517924308776855
Epoch:  40  , step:  0
training loss:  0.2508695125579834
Epoch:  50  , step:  0
training loss:  0.0881464034318924
Epoch:  60  , step:  0
training loss:  0.3708261251449585
Epoch:  70  , step:  0
training loss:  0.0828586220741272
Epoch:  80  , step:  0
training loss:  0.03668573126196861
Epoch:  90  , step:  0
training loss:  0.02559477463364601

可以看到，随着模型训练的进行，loss降低到一个比较低的水平后基本收敛。接下来，我们来验证一下模型是否真正学到了东西。

2.5 模型的评估

我们手动查看模型在测试集上的表现效果，并与2.3中的模型输出进行对比。

device = torch.device('cuda:3') if torch.cuda.is_available() else torch.device('cpu')
model.eval()
model = model.to(device)
for i,batch in enumerate(test_loader):
    input_ids = batch['text_encodings']['input_ids'].to(device)
    sql_ids = batch['sql_encodings']['input_ids'].to(device)
    result = model.generate(input_ids)
    print("==="*20)
    print("Question:")
    print(tokenizer.decode(input_ids[0]))
    print("SQL:")
    print(tokenizer.decode(result[0]))
    print()

运行以上代码，其输出如下：

============================================================
Question:
translate Text to SQL: Given the min student age in class 2.</s>
SQL:
<pad> SELECT MAX(age) FROM student WHERE class=2</s>

============================================================
Question:
translate Text to SQL: Who is the student that younger than 14.</s>
SQL:
<pad> SELECT name FROM student WHERE age>14</s>

============================================================
Question:
translate Text to SQL: Find all student ages in student database.</s>
SQL:
<pad> SELECT MIN(age) FROM student</s>

============================================================
Question:
translate Text to SQL: Please find the maxium student age in class 2.</s>
SQL:
<pad> SELECT MAX(age) FROM student WHERE class=2</s>

============================================================
Question:
translate Text to SQL: Count student's number for class 3.</s>
SQL:
<pad> SELECT COUNT(*) FROM student WHERE class=3</s>

可以看到，相比于最开始的2.3中的模型输出，此时的输出已经是一个简单的SQL语句形式了，而且对于问题4

• "translate Text to SQL: Please find the maxium student age in class 2."

模型的生成语句

• "SELECT MAX(age) FROM student WHERE class=2"
  完全正确。
  同样的，对于问题5
• "translate Text to SQL: Count student's number for class 3."，

模型的生成语句

• "SELECT COUNT(*) FROM student WHERE class=3"。
  也完全正确。

从上述结果分析来看，这样简单的模型已经能够对某些简单的自然语言问句给出正确的SQL语句回答了（当然，主要是我们的测试集太简单，这里也仅做示例而已）。

后续我们其实可以考虑扩展我们的模型以适应不同的应用场景。