add model

modeling_charmen.py

@@ -220,7 +220,7 @@ class CharmenElectraForSequenceClassification(ElectraForSequenceClassification):
 
         self.num_labels = config.num_labels
         self.config = config
-        self.model = CharmenElectraModel(config, compatibility_with_transformers=True)
+        self.electra = CharmenElectraModel(config, compatibility_with_transformers=True)
         self.classifier = CharmenElectraClassificationHead(config)
         self.cls_loss_fct = torch.nn.CrossEntropyLoss(weight=class_weight, label_smoothing=label_smoothing)
 
@@ -239,7 +239,7 @@ class CharmenElectraForSequenceClassification(ElectraForSequenceClassification):
             output_hidden_states=None,
             return_dict=None,
     ):
-        output_discriminator: CharmenElectraModelOutput = self.model(input_ids, attention_mask, token_type_ids)
+        output_discriminator: CharmenElectraModelOutput = self.electra(input_ids, attention_mask, token_type_ids)
 
         if self.carmen_config.upsample_output:
             cls = self.classifier(output_discriminator.upsampled_hidden_states)
@@ -256,7 +256,7 @@ class CharmenElectraForSequenceClassification(ElectraForSequenceClassification):
 
     def load_state_dict(self, state_dict: OrderDictType[str, Tensor], strict: bool = True):
         model = OrderedDict()
-        prefix = "discriminator.model"
+        prefix = "discriminator."
 
         for key, value in state_dict.items():
             if key.startswith('generator'):
@@ -270,8 +270,7 @@ class CharmenElectraForSequenceClassification(ElectraForSequenceClassification):
                     continue
                 model[key] = value
 
-        self.model.load_state_dict(state_dict=model, strict=False)
-        self.classifier.load_state_dict(state_dict=model, strict=False)
+        super(CharmenElectraForSequenceClassification, self).load_state_dict(state_dict=model, strict=False)
 
 
 class CharmenElectraForTokenClassification(ElectraForTokenClassification):