def __init__(self, config):
super(UnilmForLM, self).__init__(config)
self.bert = UnilmModel(config)
self.cls = BertOnlyMLMHead(config)
self.crit_mask_lm = nn.CrossEntropyLoss(reduction='none')
if hasattr(config, 'label_smoothing') and config.label_smoothing:
self.crit_mask_lm_smoothed = LabelSmoothingLoss(
config.label_smoothing, config.vocab_size, ignore_index=0, reduction='none')
else:
self.crit_mask_lm_smoothed = None
self.num_labels = 2
self.cls2 = BertOnlyNSPHead(config)
self.crit_next_sent = nn.CrossEntropyLoss(ignore_index=-1)
self.init_weights()
self.tie_weights()