transformers DataCollator介绍

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

from transformers import AutoTokenizer, AutoModel, \DataCollatorForSeq2Seq, DataCollatorWithPadding, \DataCollatorForTokenClassification, DefaultDataCollator, DataCollatorForLanguageModelingPRETRAIN_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

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进行了填充

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]])}