Where in the code does masking of tokens happen when pretraining BERT

abdallah197
Mon, 17 Aug 2020 21:17:19 +0000

@leoapolonio guided me to this snippet of code, this might be helpful, what I understood is that the function mask_tokens in DataCollatorForLanguageModeling class is responsible for randomly masking the tokens.

github.com

huggingface/transformers/blob/master/src/transformers/data/data_collator.py#L157

147.     if are_tensors_same_length:


148.         return torch.stack(examples, dim=0)


149.     else:


150.         if self.tokenizer._pad_token is None:


151.             raise ValueError(


152.                 "You are attempting to pad samples but the tokenizer you are using"


153.                 f" ({self.tokenizer.__class__.__name__}) does not have one."


154.             )


155.         return pad_sequence(examples, batch_first=True, padding_value=self.tokenizer.pad_token_id)





157. def mask_tokens(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:


158.     """


159.     Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.


160.     """





162.     if self.tokenizer.mask_token is None:


163.         raise ValueError(


164.             "This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer."


165.         )





167.     labels = inputs.clone()