Allow loading via AutoModelForSequenceClassification

bert_layers.py

@@ -29,6 +29,7 @@ from transformers.modeling_outputs import (MaskedLMOutput,
 from transformers.models.bert.modeling_bert import BertPreTrainedModel
 
 from .blockdiag_linear import BlockdiagLinear
+from .configuration_bert import BertConfig
 from .monarch_mixer_sequence_mixer import MonarchMixerSequenceMixing
 
 logger = logging.getLogger(__name__)
@@ -475,6 +476,8 @@ class BertModel(BertPreTrainedModel):
     ```
     """
 
+    config_class = BertConfig
+
     def __init__(self, config, add_pooling_layer=True):
         super(BertModel, self).__init__(config)
         self.embeddings = BertEmbeddings(config)
@@ -602,6 +605,8 @@ class BertOnlyNSPHead(nn.Module):
 #######################
 class BertForMaskedLM(BertPreTrainedModel):
 
+    config_class = BertConfig
+
     def __init__(self, config):
         super().__init__(config)
 
@@ -748,6 +753,8 @@ class BertForSequenceClassification(BertPreTrainedModel):
     e.g., GLUE tasks.
     """
 
+    config_class = BertConfig
+
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
@@ -873,6 +880,8 @@ class BertForSequenceClassification(BertPreTrainedModel):
 
 class BertForTextEncoding(BertPreTrainedModel):
 
+    config_class = BertConfig
+
     def __init__(self, config):
         super().__init__(config)