add model

config.json

@@ -1,12 +1,12 @@
 {
   "activation": "gelu",
   "architectures": [
-    "DistilBertForSequenceClassification_ANE"
+    "DistilBertForSequenceClassification"
   ],
   "attention_dropout": 0.1,
   "auto_map": {
-    "AutoConfig": "configuration_distilbert_ane.DistilBertConfig_ANE",
-    "AutoModelForSequenceClassification": "modeling_distilbert_ane.DistilBertForSequenceClassification_ANE"
+    "AutoConfig": "configuration_distilbert_ane.DistilBertConfig",
+    "AutoModelForSequenceClassification": "modeling_distilbert_ane.DistilBertForSequenceClassification"
   },
   "dim": 768,
   "dropout": 0.1,
@@ -22,7 +22,7 @@
     "POSITIVE": 1
   },
   "max_position_embeddings": 512,
-  "model_type": "distilbert_ane",
+  "model_type": "distilbert",
   "n_heads": 12,
   "n_layers": 6,
   "output_past": true,

configuration_distilbert_ane.py

@@ -1,7 +1,5 @@
-from transformers import DistilBertConfig
-
-class DistilBertConfig_ANE(DistilBertConfig):
-    model_type = "distilbert_ane"
+from transformers.models.distilbert import configuration_distilbert
 
+class DistilBertConfig(configuration_distilbert.DistilBertConfig):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)

modeling_distilbert_ane.py

@@ -43,7 +43,7 @@ import torch
 import torch.nn as nn
 
 from transformers.models.distilbert import modeling_distilbert
-from .configuration_distilbert_ane import DistilBertConfig_ANE
+from .configuration_distilbert_ane import DistilBertConfig
 
 # Note: Original implementation of distilbert uses an epsilon value of 1e-12
 # which is not friendly with the float16 precision that ANE uses by default
@@ -317,8 +317,8 @@ class Transformer(modeling_distilbert.Transformer):
                 [TransformerBlock(config) for _ in range(config.n_layers)]))
 
 
-class DistilBertModel_ANE(modeling_distilbert.DistilBertModel):
-    config_class = DistilBertConfig_ANE
+class DistilBertModel(modeling_distilbert.DistilBertModel):
+    config_class = DistilBertConfig
 
     def __init__(self, config):
         super().__init__(config)
@@ -332,14 +332,14 @@ class DistilBertModel_ANE(modeling_distilbert.DistilBertModel):
         raise NotImplementedError
 
 
-class DistilBertForMaskedLM_ANE(modeling_distilbert.DistilBertForMaskedLM):
-    config_class = DistilBertConfig_ANE
+class DistilBertForMaskedLM(modeling_distilbert.DistilBertForMaskedLM):
+    config_class = DistilBertConfig
 
     def __init__(self, config):
         super().__init__(config)
         from transformers.activations import get_activation
         setattr(self, 'activation', get_activation(config.activation))
-        setattr(self, 'distilbert', DistilBertModel_ANE(config))
+        setattr(self, 'distilbert', DistilBertModel(config))
         setattr(self, 'vocab_transform', nn.Conv2d(config.dim, config.dim, 1))
         setattr(self, 'vocab_layer_norm', LayerNormANE(config.dim, eps=EPS))
         setattr(self, 'vocab_projector',
@@ -390,13 +390,13 @@ class DistilBertForMaskedLM_ANE(modeling_distilbert.DistilBertForMaskedLM):
         return ((mlm_loss, ) + output) if mlm_loss is not None else output
 
 
-class DistilBertForSequenceClassification_ANE(
+class DistilBertForSequenceClassification(
         modeling_distilbert.DistilBertForSequenceClassification):
-    config_class = DistilBertConfig_ANE
+    config_class = DistilBertConfig
 
     def __init__(self, config):
         super().__init__(config)
-        setattr(self, 'distilbert', DistilBertModel_ANE(config))
+        setattr(self, 'distilbert', DistilBertModel(config))
         setattr(self, 'pre_classifier', nn.Conv2d(config.dim, config.dim, 1))
         setattr(self, 'classifier', nn.Conv2d(config.dim, config.num_labels,
                                               1))
@@ -441,13 +441,13 @@ class DistilBertForSequenceClassification_ANE(
         return ((loss, ) + output) if loss is not None else output
 
 
-class DistilBertForQuestionAnswering_ANE(
+class DistilBertForQuestionAnswering(
         modeling_distilbert.DistilBertForQuestionAnswering):
-    config_class = DistilBertConfig_ANE
+    config_class = DistilBertConfig
 
     def __init__(self, config):
         super().__init__(config)
-        setattr(self, 'distilbert', DistilBertModel_ANE(config))
+        setattr(self, 'distilbert', DistilBertModel(config))
         setattr(self, 'qa_outputs', nn.Conv2d(config.dim, config.num_labels,
                                               1))
 
@@ -497,12 +497,12 @@ class DistilBertForQuestionAnswering_ANE(
         return ((total_loss, ) + output) if total_loss is not None else output
 
 
-class DistilBertForTokenClassification_ANE(
+class DistilBertForTokenClassification(
         modeling_distilbert.DistilBertForTokenClassification):
 
     def __init__(self, config):
         super().__init__(config)
-        setattr(self, 'distilbert', DistilBertModel_ANE(config))
+        setattr(self, 'distilbert', DistilBertModel(config))
         setattr(self, 'classifier',
                 nn.Conv2d(config.hidden_size, config.num_labels, 1))
 
@@ -544,13 +544,13 @@ class DistilBertForTokenClassification_ANE(
         return ((loss, ) + output) if loss is not None else output
 
 
-class DistilBertForMultipleChoice_ANE(
+class DistilBertForMultipleChoice(
         modeling_distilbert.DistilBertForMultipleChoice):
-    config_class = DistilBertConfig_ANE
+    config_class = DistilBertConfig
 
     def __init__(self, config):
         super().__init__(config)
-        setattr(self, 'distilbert', DistilBertModel_ANE(config))
+        setattr(self, 'distilbert', DistilBertModel(config))
         setattr(self, 'pre_classifier', nn.Conv2d(config.dim, config.dim, 1))
         setattr(self, 'classifier', nn.Conv2d(config.dim, 1, 1))