| 471 | |
| 472 | |
| 473 | class BertPredictionHeadTransform(nn.Module): |
| 474 | def __init__(self, config): |
| 475 | super().__init__() |
| 476 | self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| 477 | if isinstance(config.hidden_act, str): |
| 478 | self.transform_act_fn = ACT2FN[config.hidden_act] |
| 479 | else: |
| 480 | self.transform_act_fn = config.hidden_act |
| 481 | self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
| 482 | |
| 483 | def forward(self, hidden_states): |
| 484 | hidden_states = self.dense(hidden_states) |
| 485 | hidden_states = self.transform_act_fn(hidden_states) |
| 486 | hidden_states = self.LayerNorm(hidden_states) |
| 487 | return hidden_states |
| 488 | |
| 489 | |
| 490 | class BertLMPredictionHead(nn.Module): |