transformers DataCollatorWithPadding类

构造方法

DataCollat​​orWithPadding(tokenizer:PreTrainedTokenizerBase
						padding:typing.Union[bool, str, transformers.utils.generic.PaddingStrategy] = True
						max_length : typing.Optional[int] = None
						pad_to_multiple_of : typing.Optional[int] = None
						return_tensors : str = 'pt ' )
1
2
3
4
5

在transfomers中，定义了一个DataCollator类，该类用于将数据集的单个元素打包成一批数据。DataCollatorWithPadding类是DataCollator类的一个实现类，该类在打包时将动态填充输入的数据。

参数tokenizer表示输入的分词器。参数padding可以为bool类型，True表示填充，False表示不填充；也可以为字符串，表示填充策略，"longest"表示根据输入数据中最长的数据来进行填充，"max_length"表示填充至参数max_length设置的长度，“do_not_pad"表示不填充。参数pad_to_multiple_of表示填充的数据的倍数。参数return_tensors表示返回的数据类型，可以为"pt”，pytorch数据类型；“tf”，tensorflow数据类型；“np”，"numpy"数据类型。

使用示例

>>> import transformers
>>> import datasets
>>> dataset = datasets.load_dataset("glue", "cola", split="train")
>>> dataset = dataset.map(lambda data: tokenizer(data["sentence"],padding=True), batched=True)
>>> dataset
Dataset({
    features: ['sentence', 'label', 'idx', 'input_ids', 'token_type_ids', 'attention_mask'],
    num_rows: 8551
})
>>> tokenizer = transformers.BertTokenizer.from_pretrained("bert-base-uncased")
>>> data_collator = transformers.DataCollatorWithPadding(tokenizer, 
						   								 padding="max_length",
						   								 max_length=12,
						   								 return_tensors="tf")
>>> dataset = dataset.to_tf_dataset(columns=["label", "input_ids"], batch_size=16, shuffle=False, collate_fn=data_collator)
>>> dataset
<PrefetchDataset element_spec={'input_ids': TensorSpec(shape=(None, None), dtype=tf.int64, name=None), 'attention_mask': TensorSpec(shape=(None, None), dtype=tf.int64, name=None), 'labels': TensorSpec(shape=(None,), dtype=tf.int64, name=None)}>
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17