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

rsvarma
Tue, 11 Aug 2020 03:12:49 +0000

In that case, I notice that in the docs for BertForPreTraining the labels that I am supposed to pass in are the following:

labels ( torch.LongTensor of shape (batch_size, sequence_length) , optional, defaults to None ) – Labels for computing the masked language modeling loss. Indices should be in [-100, 0, ..., config.vocab_size] (see input_ids docstring) Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size]

Which implies that I somehow need to have knowledge of which tokens were masked, but I am currently unclear on how to get this information as I am not sure whether the tokenizer is masking tokens or if I am expected to do so myself either before or after using the tokenizer. In that sense I suppose a more specific question would be this:
Does the BertTokenizer class have the capability to prepare masked inputs for the BertForPreTraining and BertForMaskedLM classes? If so, where in the code does this occur (while tracing through the code, I have seen that it takes in a mask_token parameter, presumably to associate an input id with it? but I have not encountered any code that actually replaces tokens with masked tokens). If not, how could I go about masking tokens myself?