Huggingface | 使用WMT16数据集微调BART训练新的标记进行翻译

发布时间 2023-07-06 20:57:14作者: 张Zong在修行

BART模型是用来预训练seq-to-seq模型的降噪自动编码器(autoencoder)。它是一个序列到序列的模型,具有对损坏文本的双向编码器和一个从左到右的自回归解码器,所以它可以完美的执行翻译任务。

如果你想在翻译任务上测试一个新的体系结构,比如在自定义数据集上训练一个新的标记,那么处理起来会很麻烦,所以在本文中,我将介绍添加新标记的预处理步骤,并介绍如何进行模型微调

因为Huggingface Hub有很多预训练过的模型,可以很容易地找到预训练标记器。但是我们要添加一个标记可能就会有些棘手,下面我们来完整的介绍如何实现它,首先加载和预处理数据集。

加载数据集

我们使用WMT16数据集及其罗马尼亚语-英语子集。load_dataset()函数将从Huggingface下载并加载任何可用的数据集。

import datasets
 
dataset = datasets.load_dataset("stas/wmt16-en-ro-pre-processed", cache_dir="./wmt16-en_ro")
print(dataset['train'][0])
# {'translation': {'en': 'Membership of Parliament: see Minutes', 'ro': 'Componenţa Parlamentului: a se vedea procesul-verbal'}}

在上面的运行结果中可以看到数据集内容。我们需要将其“压平”,这样可以更好的访问数据,让后将其保存到硬盘中。

print(dataset)
'''
DatasetDict({
    train: Dataset({
        features: ['translation'],
        num_rows: 610320
    })
    validation: Dataset({
        features: ['translation'],
        num_rows: 1999
    })
    test: Dataset({
        features: ['translation'],
        num_rows: 1999
    })
})
'''

上面的运行结果中可以看到数据集的训练集测试集和验证集的特征和条数。

def flatten(batch):
    batch['en'] = batch['translation']['en']
    batch['ro'] = batch['translation']['ro']
    
    return batch

train = dataset['train'].map(flatten)
'''
100%|███████████████████████████████████████████████████████████████████████| 610320/610320 [00:57<00:00, 10573.06ex/s]
'''

print("en => ", train[0]['en'])
print("ro => ", train[0]['ro'])
# en =>  Membership of Parliament: see Minutes
# ro =>  Componenţa Parlamentului: a se vedea procesul-verbal

上面的代码中,flatten()函数把数据的最外一层的'translation'去掉,只保留最里面的'en''ro'属性。train[0]['en']的意思是我们取训练数据的第一条的'en'属性,打印出来是:Membership of Parliament: see Minutes

test = dataset['test'].map(flatten)
validation = dataset['validation'].map(flatten)

测试集和验证集也是同样处理。

train.save_to_disk("./dataset/train")
test.save_to_disk("./dataset/test")
validation.save_to_disk("./dataset/validation")

最后就是保存处理过的数据集到我们指定的文件位置。

标记器

标记器提供了训练标记器所需的所有工作。它由四个基本组成部分:(但这四个部分不是所有的都是必要的)

models:标记器将如何分解每个单词。例如,给定单词“playing”:i) BPE模型将其分解为“play”+“ing”两个标记,ii) WordLevel将其视为一个标记。

normalizers:需要在文本上发生的一些转换。有一些过滤器可以更改Unicode、小写字母或删除内容。

pre-tokenizers:为操作文本提供更大灵活性处理的函数。例如,如何处理数字。数字100应该被认为是“100”还是“1”、“0”、“0”?

Post-Processors:后处理具体情况取决于预训练模型的选择。例如,将 [BOS](句首)或 [EOS](句尾)标记添加到 BERT 输入。

下面的代码使用BPE模型、小写Normalizers和空白Pre-Tokenizers。然后用默认值初始化训练器对象,主要包括

1、词汇量大小使用50265以与BART的英语标记器一致

2、特殊标记,如<s><pad>

3、初始词汇量,这是每个模型启动过程的预定义列表。

# 导入依赖包
from tokenizers import normalizers, pre_tokenizers, Tokenizer, models, trainers
import datasets

# 加载数据集
train = datasets.load_from_disk("./dataset/train")

# 测试是否加载进来
print(len(train)) # 610320
print(train[0]['en']) # Membership of Parliament: see Minutes
# Build a tokenizer
bpe_tokenizer = Tokenizer(models.BPE())
bpe_tokenizer.normalizer = normalizers.Lowercase()
bpe_tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()

trainer = trainers.BpeTrainer(
    vocab_size=50265,
    special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"],
    initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
)

使用Huggingface的最后一步是连接Trainer和BPE模型,并传递数据集。根据数据的来源,可以使用不同的训练函数。我们将使用train_from_iterator()

def batch_iterator():
    batch_length = 1000
    for i in range(0, len(train), batch_length):
        yield train[i : i + batch_length]["ro"]

bpe_tokenizer.train_from_iterator( batch_iterator(), length=len(train), trainer=trainer )
'''
[00:00:18] Pre-processing sequences                 ███████████████████████████████████████████████ 610320   /   610320
[00:00:00] Tokenize words                           ███████████████████████████████████████████████ 124659   /   124659
[00:00:00] Count pairs                              ███████████████████████████████████████████████ 124659   /   124659
[00:00:15] Compute merges                           ███████████████████████████████████████████████ 49900    /    49900
'''

bpe_tokenizer.save("./ro_tokenizer.json")

最后把ro的tokenizer保存成json数据。

BART微调

现在可以使用使用新的标记器了。

# 导入依赖包
from transformers import BartForConditionalGeneration, AutoTokenizer, PreTrainedTokenizerFast
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer
from datasets import load_from_disk

# 加载模型
model = BartForConditionalGeneration.from_pretrained("./model")

# 加载 Tokenizers
en_tokenizer = AutoTokenizer.from_pretrained("./model")
ro_tokenizer = PreTrainedTokenizerFast.from_pretrained( "./ro_tokenizer.json" )

ro_tokenizer.pad_token = en_tokenizer.pad_token

# 加载数据集
train = load_from_disk("./dataset/train")
test = load_from_disk("./dataset/test")
validation = load_from_disk("./dataset/validation")

# 模型输入的 Tokenizers 设置
def tokenize_dataset(sample):
    input = en_tokenizer(sample['en'], padding='max_length', max_length=120, truncation=True)
    label = ro_tokenizer(sample['ro'], padding='max_length', max_length=120, truncation=True)

    input["decoder_input_ids"] = label["input_ids"]
    input["decoder_attention_mask"] = label["attention_mask"]
    input["labels"] = label["input_ids"]

    return input

train = train.select(range(2000))
test = test.select(range(100))
validation = validation.select(range(100))

# 训练集、测试集和验证集的 tokenized 输入
train_tokenized = train.map(tokenize_dataset, batched=True)
test_tokenized = test.map(tokenize_dataset, batched=True)
validation_tokenized = validation.map(tokenize_dataset, batched=True)

上面代码的ro_tokenizer.pad_token = en_tokenizer.pad_token,为罗马尼亚语的标记器设置填充标记是非常必要的。因为它将在label = ro_tokenizer(sample['ro'], padding='max_length', max_length=120, truncation=True)使用,标记器使用填充可以使所有输入都具有相同的大小。

训练的过程:

# set training arguments - these params are not really tuned, feel free to change
training_args = Seq2SeqTrainingArguments(
    output_dir="./",
    evaluation_strategy="steps",
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    predict_with_generate=True,
    logging_steps=2,  # set to 1000 for full training
    save_steps=64,  # set to 500 for full training
    eval_steps=64,  # set to 8000 for full training
    warmup_steps=1,  # set to 2000 for full training
    max_steps=128, # delete for full training
    overwrite_output_dir=True,
    save_total_limit=3,
    fp16=True, # True if GPU
)


# instantiate trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_tokenized,
    eval_dataset=validation_tokenized,
)

# training
trainer.train()

过程也非常简单,设置训练参数,使用Trainer对象绑定所有内容,并启动流程。上述超参数都是测试目的,所以如果要得到最好的结果还需要进行超参数的设置,我们使用这些参数是可以运行的。

运行结果:

{'loss': 15.8848, 'learning_rate': 0.0, 'epoch': 0.0}                                                                  
{'loss': 15.8666, 'learning_rate': 0.0, 'epoch': 0.0}                                                                  
{'loss': 15.1562, 'learning_rate': 4.960629921259843e-05, 'epoch': 0.01}                                               
{'loss': 13.8136, 'learning_rate': 4.881889763779528e-05, 'epoch': 0.01}                                               
{'loss': 12.8159, 'learning_rate': 4.8031496062992124e-05, 'epoch': 0.01}                                              
{'loss': 11.8634, 'learning_rate': 4.724409448818898e-05, 'epoch': 0.01}                                               
{'loss': 11.1089, 'learning_rate': 4.6850393700787405e-05, 'epoch': 0.01}                                              
{'loss': 10.7867, 'learning_rate': 4.606299212598425e-05, 'epoch': 0.02}                                               
{'loss': 10.0939, 'learning_rate': 4.52755905511811e-05, 'epoch': 0.02}                                                
{'loss': 9.3424, 'learning_rate': 4.4488188976377954e-05, 'epoch': 0.02}                                               
{'loss': 8.5162, 'learning_rate': 4.370078740157481e-05, 'epoch': 0.02}                                                
{'loss': 7.8301, 'learning_rate': 4.330708661417323e-05, 'epoch': 0.02}                                                
{'loss': 7.4714, 'learning_rate': 4.251968503937008e-05, 'epoch': 0.03}                                                
{'loss': 7.8587, 'learning_rate': 4.173228346456693e-05, 'epoch': 0.03}                                                
{'loss': 6.469, 'learning_rate': 4.0944881889763784e-05, 'epoch': 0.03}                                                
{'loss': 5.8621, 'learning_rate': 4.015748031496063e-05, 'epoch': 0.03}                                                
{'loss': 5.5784, 'learning_rate': 3.9370078740157485e-05, 'epoch': 0.03}                                               
{'loss': 5.0736, 'learning_rate': 3.858267716535433e-05, 'epoch': 0.04}                                                
{'loss': 4.9796, 'learning_rate': 3.779527559055118e-05, 'epoch': 0.04}                                                
{'loss': 4.7809, 'learning_rate': 3.7007874015748034e-05, 'epoch': 0.04}                                               
{'loss': 4.4241, 'learning_rate': 3.622047244094489e-05, 'epoch': 0.04}                                                
{'loss': 4.1273, 'learning_rate': 3.5433070866141735e-05, 'epoch': 0.04}                                               
{'loss': 4.1239, 'learning_rate': 3.464566929133858e-05, 'epoch': 0.05}                                                
{'loss': 3.9724, 'learning_rate': 3.385826771653544e-05, 'epoch': 0.05}                                                
{'loss': 3.5906, 'learning_rate': 3.3070866141732284e-05, 'epoch': 0.05}                                               
{'loss': 3.3749, 'learning_rate': 3.228346456692913e-05, 'epoch': 0.05}                                                
{'loss': 3.5252, 'learning_rate': 3.1496062992125985e-05, 'epoch': 0.05}                                               
{'loss': 7.1243, 'learning_rate': 3.070866141732284e-05, 'epoch': 0.06}                                                
{'loss': 2.9568, 'learning_rate': 2.992125984251969e-05, 'epoch': 0.06}                                                
{'loss': 2.9735, 'learning_rate': 2.9133858267716534e-05, 'epoch': 0.06}                                               
{'loss': 5.8592, 'learning_rate': 2.8346456692913388e-05, 'epoch': 0.06}                                               
{'loss': 2.75, 'learning_rate': 2.755905511811024e-05, 'epoch': 0.06}                                                  
{'eval_loss': 2.8496463298797607, 'eval_runtime': 1.3008, 'eval_samples_per_second': 76.875, 'eval_steps_per_second': 38.438, 'epoch': 0.06}                                                                                                  
{'loss': 2.6724, 'learning_rate': 2.677165354330709e-05, 'epoch': 0.07}                                                
{'loss': 2.3622, 'learning_rate': 2.5984251968503937e-05, 'epoch': 0.07}                                               
{'loss': 2.3444, 'learning_rate': 2.5196850393700788e-05, 'epoch': 0.07}                                               
{'loss': 2.3187, 'learning_rate': 2.440944881889764e-05, 'epoch': 0.07}                                                
{'loss': 2.1952, 'learning_rate': 2.362204724409449e-05, 'epoch': 0.07}                                                
{'loss': 2.0657, 'learning_rate': 2.283464566929134e-05, 'epoch': 0.08}                                                
{'loss': 5.019, 'learning_rate': 2.204724409448819e-05, 'epoch': 0.08}                                                 
{'loss': 1.9412, 'learning_rate': 2.125984251968504e-05, 'epoch': 0.08}                                                
{'loss': 1.8045, 'learning_rate': 2.0472440944881892e-05, 'epoch': 0.08}                                               
{'loss': 1.8103, 'learning_rate': 1.9685039370078743e-05, 'epoch': 0.08}                                               
{'loss': 2.0013, 'learning_rate': 1.889763779527559e-05, 'epoch': 0.09}                                                
{'loss': 1.6938, 'learning_rate': 1.8110236220472444e-05, 'epoch': 0.09}                                               
{'loss': 4.2392, 'learning_rate': 1.732283464566929e-05, 'epoch': 0.09}                                                
{'loss': 1.2889, 'learning_rate': 1.6535433070866142e-05, 'epoch': 0.09}                                               
{'loss': 1.5864, 'learning_rate': 1.5748031496062993e-05, 'epoch': 0.09}                                               
{'loss': 1.4123, 'learning_rate': 1.4960629921259845e-05, 'epoch': 0.1}                                                
{'loss': 1.3579, 'learning_rate': 1.4173228346456694e-05, 'epoch': 0.1}                                                
{'loss': 1.3563, 'learning_rate': 1.3385826771653545e-05, 'epoch': 0.1}                                                
{'loss': 1.6602, 'learning_rate': 1.2598425196850394e-05, 'epoch': 0.1}                                                
{'loss': 1.3247, 'learning_rate': 1.1811023622047245e-05, 'epoch': 0.1}                                                
{'loss': 1.4064, 'learning_rate': 1.1023622047244095e-05, 'epoch': 0.11}                                               
{'loss': 0.8697, 'learning_rate': 1.0236220472440946e-05, 'epoch': 0.11}                                               
{'loss': 0.9604, 'learning_rate': 9.448818897637795e-06, 'epoch': 0.11}                                                
{'loss': 1.0761, 'learning_rate': 8.661417322834646e-06, 'epoch': 0.11}                                                
{'loss': 1.0938, 'learning_rate': 7.874015748031496e-06, 'epoch': 0.11}                                                
{'loss': 1.0781, 'learning_rate': 7.086614173228347e-06, 'epoch': 0.12}                                                
{'loss': 3.0479, 'learning_rate': 6.299212598425197e-06, 'epoch': 0.12}                                                
{'loss': 3.4292, 'learning_rate': 5.511811023622048e-06, 'epoch': 0.12}                                                
{'loss': 0.9718, 'learning_rate': 4.7244094488188975e-06, 'epoch': 0.12}                                               
{'loss': 3.2809, 'learning_rate': 3.937007874015748e-06, 'epoch': 0.12}                                                
{'loss': 0.9576, 'learning_rate': 3.1496062992125985e-06, 'epoch': 0.13}                                               
{'loss': 0.806, 'learning_rate': 2.3622047244094487e-06, 'epoch': 0.13}                                                
{'eval_loss': 1.642179012298584, 'eval_runtime': 1.5686, 'eval_samples_per_second': 63.749, 'eval_steps_per_second': 31.875, 'epoch': 0.13}                                                                                                   
{'train_runtime': 42.8357, 'train_samples_per_second': 5.976, 'train_steps_per_second': 2.988, 'train_loss': 4.709173169918358, 'epoch': 0.13}                                                                                                
100%|████████████████████████████████████████████████████████████████████████████████| 128/128 [00:42<00:00,  2.99it/s]

最后在output_dir设置的输出文件夹下,生成了checkpoint。

每个文件夹下保存的是中间训练的模型的各种参数信息等等。如下图所示:

推理

推理过程也很简单,加载经过微调的模型并使用generate()方法进行转换就可以了,但是需要注意的是对源 (En) 和目标 (RO) 序列使用适当的分词器。

总结

虽然在使用自然语言处理(NLP)时,标记化似乎是一个基本操作,但它是一个不应忽视的关键步骤。HuggingFace的出现可以方便的让我们使用,这使得我们很容易忘记标记化的基本原理,而仅仅依赖预先训练好的模型。但是当我们希望自己训练新模型时,了解标记化过程及其对下游任务的影响是必不可少的,所以熟悉和掌握这个基本的操作是非常有必要的。