NLP（九十四）transformers模块中的DataCollator

本文将会详细介绍HuggingFace开源的transformers模块中的DataCollator.

引言

Data Collator是HuggingFace开源的transformers模块进行数据处理的重要部分。它的输入是由数据集元素组成的列表，将其组装成批次，其中数据集元素为相同数据类型的train_dataset或者eval_dataset。

为了组装成数据批次，Data collators会应用某些处理（比如padding），有些（比如DataCollatorForLanguageModeling）还会在数据批次上应用随机数据增强（比如随机masking）。

Data collators是为了特定任务而设计的，如下：

Causal language modeling (CLM)
Masking language modeling (MLM)
Sequence classification
Seq2Seq
Token classification

以sequence classification任务为例，data collator只需将所有序列填充到一个小批次中，以确保它们具有相同的长度。

使用例子进行解读

在transformers模块源码中，data collators源码位于data/data_collator.py脚本中，所有data collator的父类为DataCollatorMixin，代码如下：

class DataCollatorMixin:
    def __call__(self, features, return_tensors=None):
        if return_tensors is None:
            return_tensors = self.return_tensors
        if return_tensors == "tf":
            return self.tf_call(features)
        elif return_tensors == "pt":
            return self.torch_call(features)
        elif return_tensors == "np":
            return self.numpy_call(features)
        else:
            raise ValueError(f"Framework '{return_tensors}' not recognized!")

这个类根据输入return_tensors决定处理哪种矩阵torch、tensorflow、numpy，一般不直接使用。

以下是其派生的子类，我们将一一进行了解。

DefaultDataCollator

DefaultDataCollator类为默认DataCollator类，它不进行任何padding或者truncation，并且假设所有的输入样本拥有相同的长度。如果输入样本长度不一样，则会报错。

该类一般不直接使用。

DataCollatorWithPadding

DataCollatorWithPadding类将输入的input_ids、attention_mask等向量做padding处理，截断或补充padding，使得它们的长度保持统一。

from transformers import DataCollatorWithPadding, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-7B")
dc = DataCollatorWithPadding(tokenizer)
raw_tokens = tokenizer(['南京', '南京市长', '南京市长江大桥'])
print(raw_tokens)
print(dc(raw_tokens))

输出结果如下：

{'input_ids': [[102034], [102034, 102975], [102034, 102975, 69177, 106936]], 'attention_mask': [[1], [1, 1], [1, 1, 1, 1]]}
{'input_ids': tensor([[102034, 151643, 151643, 151643],
        [102034, 102975, 151643, 151643],
        [102034, 102975,  69177, 106936]]), 'attention_mask': tensor([[1, 0, 0, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 1]])}

这里使用的transformers模块版本为4.39.3, tokenizers模块版本为0.15.2。

我们为Qwen1.5-7B模型的tokenier为例。从上面的输出中可以看到，原始的数据经tokenize后input_ids和attention_mask长度并不一样，使用DataCollatorWithPadding后它们的长度一样，使用151643数字（对应token为<|endoftext|>）进行填充。

DataCollatorWithPadding类还有其他参数，比如padding, max_length, pad_to_multiple_of, return_tensors等，我们可以指定序列最大长度，比如：

1	`DataCollatorWithPadding(tokenizer, max_length=10, padding='max_length')`

DataCollatorForTokenClassification

DataCollatorForTokenClassification类适用于序列标注任务（token classification），比如命名实体识别（NER），每个token都会对应一个预测label，当一个序列的token比label多时，对额外的label位置添加-100处理，使得在计算交叉熵损失时，-100位置的label损失为0不用考虑。

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-7B")
dc = DataCollatorForTokenClassification(tokenizer)
data = [('南京', [1]), ('南京市长', [1, 0]), ('南京市长江大桥', [1, 0])]
features = []
for text, label in data:
    feature = tokenizer(text)
    feature['label'] = label
    features.append(feature)

print(dc(features))

输出结果为:

{'input_ids': tensor([[102034, 151643, 151643, 151643],
        [102034, 102975, 151643, 151643],
        [102034, 102975,  69177, 106936]]), 'attention_mask': tensor([[1, 0, 0, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 1]]), 'label': tensor([[   1, -100, -100, -100],
        [   1,    0, -100, -100],
        [   1,    0, -100, -100]])}

可以看到，此时仍对序列进行了填充，使得所有序列保持同一长度，同时对token多于label的场景使用-100进行填充。

DataCollatorForSeq2Seq

DataCollatorForSeq2Seq类适用于Seq2Seq任务，包括机器翻译，文本生成任务等。该类根据前面的序列来预测后面序列，后面序列为labels，一个batch内多条后面序列长度不同时，用-100来填充。同样地，在计算损失时，-100位置的损失不予考虑。

from transformers import DataCollatorForSeq2Seq, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-7B")
dc = DataCollatorForSeq2Seq(tokenizer)
data = [('南京', '市长江大桥'), ('南京市', '长江大桥'), ('南京市长江', '大桥')]

features = []
for text, label in data:
    feature = tokenizer(text)
    feature['labels'] = tokenizer(label)['input_ids']
    features.append(feature)

print(dc(features))

输出结果为：

{'input_ids': tensor([[102034, 151643, 151643],
        [112891, 151643, 151643],
        [102034, 102975,  69177]]), 'attention_mask': tensor([[1, 0, 0],
        [1, 0, 0],
        [1, 1, 1]]), 'labels': tensor([[102975,  69177, 106936],
        [104924, 106936,   -100],
        [106936,   -100,   -100]])}

可以看到，后面序列labels长度不一致，使用-100填充至batch内最大长度。

DataCollatorForLanguageModeling

DataCollatorForLanguageModeling类适用于语言模型，比如BERT系列等。

该类有个mlm参数，其默认值为True, 取值如下：

如果设置为False，则labels与inputs一致，使用-100进行填充
如果设置为True，则labels对于non-masked tokens设置为-100，对于masked token，其处理方式同BERT模型：80% MASK（MASK值）, 10% random（随机替换成其它token值）, 10% original（原始值）。

mlm_probability参数默认值为0.15，即对input_ids中的token会有15%的概率进行mask。

from transformers import DataCollatorForLanguageModeling, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
dc = DataCollatorForLanguageModeling(tokenizer)
data = ['南京', '南京市', '南京市长江']
raw_tokens = [tokenizer(text) for text in data]
print(dc(raw_tokens))

输出结果为：

{'input_ids': tensor([[ 101, 1298,  776,  102,    0,    0,    0],
        [ 101, 1298,  776, 2356,  102,    0,    0],
        [ 101, 1298,  776, 2356, 7270, 3736,  102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1]]), 'labels': tensor([[-100, -100, -100, -100, -100, -100, -100],
        [-100, -100, -100, -100, -100, -100, -100],
        [-100, -100,  776, -100, -100, -100, -100]])}

注意: 我们这里使用模型为bert-base-chinese，基座为BERT模型。

输出的labels列表中第三个序列中第三个元素进行了mask。

DataCollatorForWholeWordMask

DataCollatorForWholeWordMask类继承了DataCollatorForLanguageModeling类，区别在于其全词掩盖功能（WholeWordMask），该遮词方法需要在tokenizer分词的时候就对期望连续遮住的词汇的非第一个字前加上##标记。

DataCollatorForSOP

DataCollatorForSOP适用于句子顺序预测任务，将会在后续版本中进行移除，这里不再介绍。

DataCollatorForPermutationLanguageModeling

DataCollatorForPermutationLanguageModeling适用于置换语言模型（permutation language modeling），比如XLNet模型等，可以参考XLNet模型的相关解读，这里不予介绍。

NER实战

在文章NLP（六十六）使用HuggingFace中的Trainer进行BERT模型微调中，我们完成文本分类任务时使用的Data Collator为DataCollatorWithPadding。

我们使用DataCollatorForTokenClassification类来完成命名实体识别任务，数据集用peoples_daily_ner，模型采用bert-base-chinese。

训练脚本如下：

import datasets
from transformers import AutoTokenizer, DataCollatorForTokenClassification
from transformers import Trainer, TrainingArguments
from transformers import AutoModelForTokenClassification


model_id = 'bert-base-chinese'
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForTokenClassification.from_pretrained(model_id, num_labels=7)
raw_datasets = datasets.load_dataset('peoples_daily_ner')
max_length=128

def add_text(sample):
    text = []
    for item in sample['tokens']:
        text.append(''.join(item))
    return {"text": text}

processed_dataset = raw_datasets.map(add_text, batched=True)

def tokenize_function(sample):
    features = tokenizer(sample['text'], max_length=max_length, padding=True, truncation=True)
    features['labels'] = []
    for i in range(len(sample['ner_tags'])):
        features['labels'].append([-100] + sample['ner_tags'][i][:max_length-2])
    return features

tokenized_datasets = processed_dataset.map(tokenize_function, batched=True)
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

training_args = TrainingArguments(output_dir='ner_model',
                                  evaluation_strategy="epoch",
                                  per_device_train_batch_size=64,
                                  per_device_eval_batch_size=32,
                                  learning_rate=1e-4,
                                  num_train_epochs=10,
                                  warmup_ratio=0.2,
                                  logging_dir='ner_log',
                                  logging_strategy="epoch",
                                  save_strategy="epoch",
                                  report_to="tensorboard") 

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    data_collator=data_collator,
    tokenizer=tokenizer
)

trainer.train()

这里我们的数据预处理方式比较粗糙，样本的token与label并没有严格对应，不过对于bert-base-chinese模型，大多数样本已经满足要求了。

我们对测试集进行评估：

import datasets
import numpy as np
from transformers import AutoTokenizer
from transformers import AutoModelForTokenClassification
from seqeval.metrics import classification_report
from tqdm import tqdm

max_length = 128
checkpoint = "./ner_model/checkpoint-3270"
model = AutoModelForTokenClassification.from_pretrained(checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
label_list = ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']

def single_predict(text):
    encoding = tokenizer(text, max_length=max_length, padding=True, truncation=True, return_tensors='pt')
    logits = model(**encoding).logits
    prediction = np.argmax(logits.detach().numpy(), axis=2)[0]
    return [label_list[_] for _ in prediction[1:-1]]

text = "王安石的另一位好友叫吴充。"
t = single_predict(text=text)
print(t)

raw_datasets = datasets.load_dataset('peoples_daily_ner')['test']
true_tag_list, pred_tag_list = [], []
for i in tqdm(range(len(raw_datasets))):
    test_text = ''.join(raw_datasets[i]['tokens'])
    true_tag = [label_list[_] for _ in raw_datasets[i]['ner_tags']]
    pred_tag = single_predict(text=test_text)
    true_tag_list.append(true_tag)
    if len(true_tag) <= len(pred_tag):
        pred_tag_list.append(pred_tag[:len(true_tag)])
    else:
        pred_tag_list.append(pred_tag+["O"]*(len(true_tag)-len(pred_tag)))

print(classification_report(true_tag_list, pred_tag_list, digits=4))

评估指标如下：

              precision    recall  f1-score   support

         LOC     0.9179    0.8650    0.8906      3658
         ORG     0.8596    0.8801    0.8697      2185
         PER     0.9514    0.9340    0.9426      1864

   micro avg     0.9087    0.8859    0.8972      7707
   macro avg     0.9096    0.8930    0.9010      7707
weighted avg     0.9095    0.8859    0.8973      7707