| 528 | |
| 529 | |
| 530 | class BertPreTrainingHeads(nn.Module): |
| 531 | def __init__(self, config): |
| 532 | super().__init__() |
| 533 | self.predictions = BertLMPredictionHead(config) |
| 534 | self.seq_relationship = nn.Linear(config.hidden_size, 2) |
| 535 | |
| 536 | def forward(self, sequence_output, pooled_output): |
| 537 | prediction_scores = self.predictions(sequence_output) |
| 538 | seq_relationship_score = self.seq_relationship(pooled_output) |
| 539 | return prediction_scores, seq_relationship_score |
| 540 | |
| 541 | |
| 542 | class BertPreTrainedModel(PreTrainedModel): |