| 488 | |
| 489 | |
| 490 | class BertLMPredictionHead(nn.Module): |
| 491 | def __init__(self, config): |
| 492 | super().__init__() |
| 493 | self.transform = BertPredictionHeadTransform(config) |
| 494 | |
| 495 | # The output weights are the same as the input embeddings, but there is |
| 496 | # an output-only bias for each token. |
| 497 | self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| 498 | |
| 499 | self.bias = nn.Parameter(torch.zeros(config.vocab_size)) |
| 500 | |
| 501 | # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings` |
| 502 | self.decoder.bias = self.bias |
| 503 | |
| 504 | def forward(self, hidden_states): |
| 505 | hidden_states = self.transform(hidden_states) |
| 506 | hidden_states = self.decoder(hidden_states) |
| 507 | return hidden_states |
| 508 | |
| 509 | |
| 510 | class BertOnlyMLMHead(nn.Module): |