| 650 | |
| 651 | |
| 652 | class BertOnlyMLMHead(nn.Module): |
| 653 | def __init__(self, config, bert_model_embedding_weights): |
| 654 | super(BertOnlyMLMHead, self).__init__() |
| 655 | self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights) |
| 656 | |
| 657 | def forward(self, sequence_output): |
| 658 | prediction_scores = self.predictions(sequence_output) |
| 659 | return prediction_scores |
| 660 | |
| 661 | |
| 662 | class BertOnlyNSPHead(nn.Module): |