Making auto-etc work better

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)
 

config.json

@@ -8,8 +8,8 @@
     "bidirectional": true,
     "auto_map": {
         "AutoConfig": "configuration_bert.BertConfig",
-        "AutoModelForSequenceClassification": "bert_layers.BertForTextEncoding",
-        "AutoTokenizer": "bert-base-uncased"
+	"AutoModel": "bert_layers.BertModel",
+        "AutoModelForSequenceClassification": "bert_layers.BertForTextEncoding"
     },
     "classifier_dropout": null,
     "gradient_checkpointing": false,

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 2048,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}