class BertPretrainingLoss(BertPreTrainedModel):
def __init__(self, bert_encoder, config):
super(BertPretrainingLoss, self).__init__(config)
self.bert = bert_encoder
self.cls = BertPreTrainingHeads(
config, self.bert.embeddings.word_embeddings.weight)
self.cls.apply(self.init_bert_weights)
def forward(self,
input_ids,
token_type_ids=None,
attention_mask=None,
masked_lm_labels=None,
next_sentence_label=None):
sequence_output, pooled_output = self.bert(
input_ids,
token_type_ids,
attention_mask,
output_all_encoded_layers=False)
prediction_scores, seq_relationship_score = self.cls(
sequence_output, pooled_output)
if masked_lm_labels is not None and next_sentence_label is not None:
loss_fct = CrossEntropyLoss(ignore_index=-1)
next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2),
next_sentence_label.view(-1))
masked_lm_loss = loss_fct(
prediction_scores.view(-1, self.config.vocab_size),
masked_lm_labels.view(-1))
total_loss = masked_lm_loss + next_sentence_loss
return total_loss
else:
return prediction_scores, seq_relationship_score
def __init__(self, config, num_labels=2, num_rel=0, num_sentlvl_labels=0, no_nsp=False): super(BertForPreTrainingLossMask, self).__init__(config) self.bert = BertModel(config) self.cls = BertPreTrainingHeads( config, self.bert.embeddings.word_embeddings.weight, num_labels=num_labels) self.num_sentlvl_labels = num_sentlvl_labels self.cls2 = None if self.num_sentlvl_labels > 0: self.secondary_pred_proj = nn.Embedding( num_sentlvl_labels, config.hidden_size) self.cls2 = BertPreTrainingHeads( config, self.secondary_pred_proj.weight, num_labels=num_sentlvl_labels) self.crit_mask_lm = nn.CrossEntropyLoss(reduction='none') if no_nsp: self.crit_next_sent = None else: self.crit_next_sent = nn.CrossEntropyLoss(ignore_index=-1) self.num_labels = num_labels self.num_rel = num_rel if self.num_rel > 0: self.crit_pair_rel = BertPreTrainingPairRel( config, num_rel=num_rel) 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.apply(self.init_bert_weights) self.bert.rescale_some_parameters()
def __init__(self, config):
super(BertForMTPostTraining, self).__init__(config)
self.bert = BertModel(config)
self.cls = BertPreTrainingHeads(
config, self.bert.embeddings.word_embeddings.weight)
self.qa_outputs = torch.nn.Linear(config.hidden_size, 2)
self.apply(self.init_bert_weights)
def test_BertPreTrainingHeads():
input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
config = BertConfig(vocab_size_or_config_json_file=32000,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072)
embeddings = BertEmbeddings(config)
model = BertPreTrainingHeads(config, embeddings.word_embeddings.weight)
embedding_output = embeddings(input_ids, token_type_ids)
print(model(embedding_output, embedding_output))
def __init__(self, bert_encoder, config):
super(BertPretrainingLoss, self).__init__(config)
self.bert = bert_encoder
self.cls = BertPreTrainingHeads(
config, self.bert.embeddings.word_embeddings.weight)
self.cls.apply(self.init_bert_weights)
def __init__(self, config):
super(BertForPreTraining, self).__init__(config)
self.bert = BertModel(config)
self.cls = BertPreTrainingHeads(
config, self.bert.embeddings.word_embeddings.weight)
self.apply(self.init_bert_weights)