Upload modeling_bert.py with huggingface_hub

modeling_bert.py

@@ -256,8 +256,19 @@ class BertPooler(nn.Module):
         return self.activation(self.dense(hidden_states[:, 0]))
 
 
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.LayerNorm(F.gelu(self.dense(hidden_states)))
+
+
 class BertModel(PreTrainedModel):
     config_class = BertUpdatedConfig
+    base_model_prefix = "bert"
     _supports_sdpa = True
     _supports_flash_attn_2 = True
 
@@ -314,12 +325,14 @@ class BertModel(PreTrainedModel):
 
 class BertForMaskedLM(PreTrainedModel):
     config_class = BertUpdatedConfig
+    base_model_prefix = "bert"
     _supports_sdpa = True
     _supports_flash_attn_2 = True
 
     def __init__(self, config):
         super().__init__(config)
         self.bert = BertModel(config)
+        self.transform = BertPredictionHeadTransform(config)
         self.cls = nn.Linear(config.hidden_size, config.vocab_size)
         self.post_init()
 
@@ -343,7 +356,7 @@ class BertForMaskedLM(PreTrainedModel):
             output_hidden_states=output_hidden_states, output_attentions=output_attentions,
             return_dict=True,
         )
-        logits = self.cls(outputs.last_hidden_state)
+        logits = self.cls(self.transform(outputs.last_hidden_state))
 
         loss = None
         if labels is not None: