add missing parts

automodel.py

@@ -5,6 +5,7 @@ import torch
 import torch.nn as nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+from transformers.activations import ACT2FN
 from transformers import BertPreTrainedModel
 from transformers.modeling_outputs import (MaskedLMOutput,
                                            SequenceClassifierOutput)
@@ -233,6 +234,24 @@ class ClinicalMosaicForSequenceClassification(BertPreTrainedModel):
             attentions=None,
         )
 
+
+class BertPredictionHeadTransform(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=1e-12)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
 class BertLMPredictionHead(nn.Module):
 
     def __init__(self, config, bert_model_embedding_weights):