当前位置:   article > 正文

transformers DataCollatorForLanguageModeling类

datacollatorforlanguagemodeling
构造方法
DataCollatorForLanguageModeling(
	tokenizer: PreTrainedTokenizerBase,
    mlm: bool = True,
    mlm_probability: float = 0.15,
    pad_to_multiple_of: Optional[int] = None,
    return_tensors: str = "pt")
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6

在构建语言模型或者说是进行MLM任务时需要使用的数据收集器,该数据收集器会以一定概率(由参数mlm_probability控制)将序列中的Token替换成Mask标签。不同于DataCollatorWithPadding、DataCollatorForTokenClassification和DataCollatorForTokenClassification,该数据收集器只会将序列填充到最长序列长度。

参数tokenzier表示用于编码数据的分词器。

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

参数mlm_probability表示随机替换Token的概率。值越大,被替换的Token的数量也就越多。默认为0.15,这也就是BERT模型中预训练MLM任务中使用的概率值。

参数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("distilroberta-base")
data_collator = transformers.DataCollatorForLanguageModeling(tokenizer=tokenizer,
                                                             mlm=True,
                                                             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
  • 25
本文内容由网友自发贡献,转载请注明出处:【wpsshop博客】
推荐阅读
相关标签
  

闽ICP备14008679号