add model

modeling_charmen.py

@@ -26,7 +26,7 @@ class CharmenElectraModelOutput(ModelOutput):
 class CharmenElectraModel(ElectraPreTrainedModel):
     config_class = CharmenElectraConfig
 
-    def __init__(self, config: CharmenElectraConfig, compatibility_with_transformers=False):
+    def __init__(self, config: CharmenElectraConfig, compatibility_with_transformers=False, **kwargs):
         super().__init__(config)
         self.embeddings: GBST = GBST(
             num_tokens=config.vocab_size,
@@ -178,16 +178,20 @@ class CharmenElectraModel(ElectraPreTrainedModel):
         prefix = "discriminator.electra."
 
         for key, value in state_dict.items():
+            if key.startswith('generator'):
+                continue
             if key.startswith(prefix):
                 model[key[len(prefix):]] = value
+            else:
+                continue
 
-        super(CharmenElectraModel, self).load_state_dict(model, strict)
+        super(CharmenElectraModel, self).load_state_dict(state_dict=model, strict=strict)
 
 
 class CharmenElectraClassificationHead(nn.Module):
     """Head for sentence-level classification tasks."""
 
-    def __init__(self, config: CharmenElectraConfig):
+    def __init__(self, config: CharmenElectraConfig, **kwargs):
         super().__init__()
         self.config = config
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
@@ -211,7 +215,7 @@ class CharmenElectraClassificationHead(nn.Module):
 class CharmenElectraForSequenceClassification(ElectraForSequenceClassification):
     config_class = CharmenElectraConfig
 
-    def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0):
+    def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0, **kwargs):
         super().__init__(config)
 
         self.num_labels = config.num_labels
@@ -252,17 +256,26 @@ class CharmenElectraForSequenceClassification(ElectraForSequenceClassification):
 
     def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True):
         model = OrderedDict()
-        prefix = "discriminator.electra."
+        prefix = "discriminator.model"
 
         for key, value in state_dict.items():
+            if key.startswith('generator'):
+                continue
             if key.startswith(prefix):
+                if 'discriminator_predictions' in key:
+                    continue
                 model[key[len(prefix):]] = value
+            else:
+                if key.startswith('sop'):
+                    continue
+                model[key] = value
 
-        self.model.load_state_dict(state_dict=model, strict=strict)
+        self.model.load_state_dict(state_dict=model, strict=False)
+        self.classifier.load_state_dict(state_dict=model, strict=False)
 
 
 class CharmenElectraForTokenClassification(ElectraForTokenClassification):
-    def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0):
+    def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0, **kwargs):
         super().__init__(config)
 
         self.num_labels = config.num_labels
@@ -317,13 +330,17 @@ class CharmenElectraForTokenClassification(ElectraForTokenClassification):
 
     def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True):
         model = OrderedDict()
-        prefix = "discriminator.electra."
+        prefix = "discriminator."
 
         for key, value in state_dict.items():
+            if key.startswith('generator'):
+                continue
             if key.startswith(prefix):
-                model[key[len(prefix):]] = value
+                model[key[len(prefix):].replace('electra', 'model')] = value
+            else:
+                model[key] = value
 
-        self.model.load_state_dict(state_dict=model, strict=strict)
+        super(CharmenElectraForTokenClassification, self).load_state_dict(state_dict=model, strict=strict)
 
 
 class Pooler(nn.Module):
@@ -342,7 +359,7 @@ class Pooler(nn.Module):
 
 
 class CharmenElectraForMultipleChoice(ElectraForMultipleChoice):
-    def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0):
+    def __init__(self, config: CharmenElectraConfig, class_weight=None, label_smoothing=0.0, **kwargs):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.config = config
@@ -401,10 +418,14 @@ class CharmenElectraForMultipleChoice(ElectraForMultipleChoice):
 
     def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True):
         model = OrderedDict()
-        prefix = "discriminator.electra."
+        prefix = "discriminator."
 
         for key, value in state_dict.items():
+            if key.startswith('generator'):
+                continue
             if key.startswith(prefix):
-                model[key[len(prefix):]] = value
+                model[key[len(prefix):].replace('electra', 'model')] = value
+            else:
+                model[key] = value
 
-        self.model.load_state_dict(state_dict=model, strict=strict)
+        super(CharmenElectraForMultipleChoice, self).load_state_dict(state_dict=model, strict=strict)