transformers DataCollator介绍

本博客主要介绍 transformers DataCollator的使用

from transformers import AutoTokenizer, AutoModel, \
    DataCollatorForSeq2Seq, DataCollatorWithPadding, \
    DataCollatorForTokenClassification, DefaultDataCollator, DataCollatorForLanguageModeling

PRETRAIN_MODEL = "E:\pythonWork\models\chinese-roberta-wwm-ext"
tokenizer = AutoTokenizer.from_pretrained(PRETRAIN_MODEL)
model = AutoModel.from_pretrained(PRETRAIN_MODEL)

texts = ['今天天气真好。', "我爱你"]
encodings = tokenizer(texts)


labels = [list(range(len(each))) for each in texts]

inputs = [{"input_ids":t, "labels": l} for t,l in zip(encodings['input_ids'], labels)]



dc1 = DefaultDataCollator()
dc2 = DataCollatorForTokenClassification(tokenizer=tokenizer)
dc3 = DataCollatorWithPadding(tokenizer=tokenizer)
dc4 = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
d5 = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
d6 = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)


print('DataCollatorForTokenClassification')
print(dc2(inputs))



print('DataCollatorWithPadding')
print(dc3(encodings))

print('DataCollatorForSeq2Seq')
print(dc4(inputs))


print(123)

DataCollatorForTokenClassification

观察如下输出，token分类任务中，每个token都应该有一个标签，所以存在以下数量关系：

长度：ids==labels

ids进行了填充，

labels进行了填充

attention_mask进行了填充

DataCollatorForTokenClassification
{'input_ids': tensor([[ 101,  791, 1921, 1921, 3698, 4696, 1962,  511,  102],
        [ 101, 2769, 4263,  872,  102,    0,    0,    0,    0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 0, 0, 0, 0]]), 'labels': tensor([[   0,    1,    2,    3,    4,    5,    6, -100, -100],
        [   0,    1,    2, -100, -100, -100, -100, -100, -100]])}

DataCollatorWithPadding

这个collator一般使用在文本分类，文本分类中，label和输入的ids长度之间没有什么关系。

ids进行了填充，

labels进行了填充

attention_mask进行了填充

DataCollatorWithPadding
{'input_ids': tensor([[ 101,  791, 1921, 1921, 3698, 4696, 1962,  511,  102],
        [ 101, 2769, 4263,  872,  102,    0,    0,    0,    0]]), 'token_type_ids': tensor([[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, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 0, 0, 0, 0]])}

DataCollatorForSeq2Seq

长度：ids 和labels之间的长度不存在关系。  （注意和DataCollatorForTokenClassification进行区分）

ids进行了填充，

labels进行了填充

attention_mask进行了填充

注意：这个示例通常是encoder-decoder(或者称之为condition generation)模型使用的输入,ids和label长度不需要一样，encoder模型输入ids，抽取特征，decoder进行解码，预测labels。此时的mask是针对ids的填充进行mask(ids中的0即是填充)

DataCollatorForSeq2Seq
{'input_ids': tensor([[ 101,  791, 1921, 1921, 3698, 4696, 1962,  511,  102],
        [ 101, 2769, 4263,  872,  102,    0,    0,    0,    0]]), 'labels': tensor([[   0,    1,    2,    3,    4,    5,    6],
        [   0,    1,    2, -100, -100, -100, -100]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 0, 0, 0, 0]])}

还有另一种情况，就是gpt，chatglm,llama这种decoder模型中，也会使用DataCollatorForSeq2Seq，但是此时的labels和ids必是相同长度的然后。原因是这样的，

此时的输入=ids+labels

此时的标签（labels）= len(ids)*[-100] + labels (-100不计算损失)

举个例子：

假设输入input='我爱' ，对应的ids是[1,2]  ,我希望模型预测出标签label='人民',对应的label ids是[3,4]

那么模型输入=[1,2,3,4]

模型要预测的标签=[-100,-100,3,4]   

如果我们考虑多个例子，算上填充，那么情况是下面这样的。

假设我们还有另一条数据是input='今天天气',  label='真好'

把这两条样本组合成一个batch，则

输入=[

[1,2,3,4,0,0],  # 填充0，保持长度一样

[11,12,13,14,15,16]

]

标签=[

[-100,-100,3,4,-100,-100] ,  # 末尾填充了两个-100，保证batch内长度一致，并且-100不计算损失

[-100,-100,-100,-100,15,16]     # 前面4个是-100，因为‘今天天气’四个字是已知的，我们只想要预测“真好”两个字

]