transformers DataCollatorForWholeWordMask类

构造方法

DataCollatorForWholeWordMask( 
	tokenizer: PreTrainedTokenizerBase,
	mlm: bool = True,
	mlm_probability: float = 0.15,
	pad_to_multiple_of: typing.Optional[int] = None,
	tf_experimental_compile: bool = False,
	return_tensors: str = 'pt' )
1
2
3
4
5
6
7

在构建MLM时，15%的Word piece会被mask，这15%中，80%会直接替换为[Mask]，10%将其替换为其他任意单词，10%保留原始token。

而WWM(Whole Word Masking)将整词mask，改变了原来预训练阶段的训练样本生成策略。原有的基于wordPiece的分词方式将一个完整的词切分为若干子词，预训练阶段字词被分开随机mask。而www采取的措施是，一个完整的词被mask，同属于该词的其他部分也会被mask，即全词mask。

参数tokenizer表示对句子编码的分词器。

参数mlm表示是否使用MLM模型，也就是会随机的Mask掉序列中的一部分Token。该参数设置为True时，也就是使用MLM模型时，会将输入数据的labels中没有被mask掉的token的值设置为-100，被mask掉的token的值设置为原来的值。

参数mlm_probability表示被mask掉的词，而非字，占序列的长度。也可以说是mask掉词的概率。

参数pad_to_multiple_of表示填充的序列的倍数。

参数return_tensors表示返回数据的类型，有三个可选项，分别是"tf"、“pt”、“np”，分别表示tensorflow可以处理的数据类型，pytorch可以处理的数据类型以及numpy数据类型。

使用示例

def preprocess_fn(data):
    data = {k: sum(data[k], []) for k in data.keys()}
    length = len(data["input_ids"]) // 128 * 128
    result = {k: [v[i: i + 128] for i in range(0, length, 128)] for k, v in data.items()}
    result["labels"] = result["input_ids"].copy()
    return result


dataset = datasets.load_dataset("wikitext", "wikitext-2-raw-v1")
tokenizer = transformers.AutoTokenizer.from_pretrained("bert-wwm")
data_collator = transformers.DataCollatorForWholeWordMask(tokenizer=tokenizer,
                                                          mlm_probability=0.15,
                                                          return_tensors="tf")
dataset = dataset.map(lambda data: tokenizer(data["text"], truncation=True),
                      batched=True,
                      batch_size=1000,
                      remove_columns=["text"])
dataset = dataset.map(preprocess_fn,
                      batched=True,
                      batch_size=1000)
train_dataset = dataset["train"].to_tf_dataset(columns=["input_ids", "attention_mask", "labels"],
                                               batch_size=16,
                                               shuffle=True,
                                               collate_fn=data_collator)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24