Upload modeling_bert.py with huggingface_hub

modeling_bert.py

@@ -271,12 +271,13 @@ class BertModel(PreTrainedModel):
     base_model_prefix = "bert"
     _supports_sdpa = True
     _supports_flash_attn_2 = True
+    _keys_to_ignore_on_load_missing = [r"pooler\."]
 
-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
         super().__init__(config)
         self.embeddings = BertEmbeddings(config)
         self.encoder = BertEncoder(config)
-        self.pooler = BertPooler(config)
+        self.pooler = BertPooler(config) if add_pooling_layer else None
         self.post_init()
 
     def get_input_embeddings(self):
@@ -310,7 +311,7 @@ class BertModel(PreTrainedModel):
             output_hidden_states=output_hidden_states,
             output_attentions=output_attentions,
         )
-        pooled = self.pooler(last_hidden_state)
+        pooled = self.pooler(last_hidden_state) if self.pooler is not None else None
 
         if not return_dict:
             return tuple(v for v in [last_hidden_state, pooled, all_hidden_states, all_attentions] if v is not None)