Update custom_model_package/custom_model.py

custom_model_package/custom_model.py

@@ -47,3 +47,16 @@ class CustomModel(XLMRobertaForSequenceClassification):
             return {"loss": loss, "emotion_logits": emotion_logits, "sentiment_logits": sentiment_logits}
         return {"emotion_logits": emotion_logits, "sentiment_logits": sentiment_logits}
 
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop('config', None)
+        if config is None:
+            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        num_emotion_labels = kwargs.pop('num_emotion_labels', config.num_emotion_labels)
+        model = super(CustomModel, cls).from_pretrained(pretrained_model_name_or_path, config=config, *model_args, **kwargs)
+        model.num_emotion_labels = num_emotion_labels
+        return model
+
+# Register the custom configuration and model
+CONFIG_MAPPING.update({"custom_model": CustomConfig})
+MODEL_MAPPING.update({"custom_model": CustomModel})