gbure

masked_lm.py (2472B)

---

 from typing import Callable
 
 import torch
 import transformers
 
 import gbure.utils
 
 
 class MaskedLM(torch.nn.Module):
     """
     Masked language model to be used on top of a transformer.
 
     This class is only useful for unsupervised pre-training, Soares et al. keep the BERT loss alongside their "matching the blanks" loss.
 
     Config:
         transformer_model: Which transformer to use (e.g. bert-large-uncased).
     """
 
     def __init__(self, config: gbure.utils.dotdict, tokenizer: transformers.PreTrainedTokenizer) -> None:
         """
         Instantiate a masked language model.
 
         Args:
             config: global config object
             tokenizer: tokenizer used to create the vocabulary
         """
         super().__init__()
 
         self.config: gbure.utils.dotdict = config
         self.tokenizer: transformers.PreTrainedTokenizer = tokenizer
 
         self.transformer: transformers.PreTrainedModel
         if self.config.get("load") or self.config.get("pretrained"):
             # TODO introduce a config parameter to change the initialization of <tags> embeddings
             transformer_config = transformers.AutoConfig.from_pretrained(self.config.transformer_model)
             transformer_config.vocab_size = len(tokenizer)
             self.transformer = transformers.AutoModelForMaskedLM(transformer_config)
         else:
             self.transformer = transformers.AutoModelForMaskedLM.from_pretrained(self.config.transformer_model)
             self.transformer.resize_token_embeddings(len(tokenizer))
 
     @property
     def encoder(self) -> transformers.PreTrainedModel:
         if isinstance(self.transformer, transformers.BertForMaskedLM):
             return self.transformer.bert
         elif isinstance(self.transformer, transformers.DistilBertForMaskedLM):
             return self.transformer.distilbert
         else:
             raise RuntimeError("Unknown transformer model, can't split masked language model off")
 
     def forward(self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         """
         Compute masked language model loss from transformer's output.
 
         Sadly, huggingface does not provide an unified interface, so we need to do some copy-pasting.
         """
         masked_target: torch.Tensor = target * mask - 100 * (1 - mask.long())
         output = self.transformer(input_ids=input, attention_mask=mask, labels=masked_target, return_dict=True)
         return output.loss