Update modeling_esm_plusplus.py

modeling_esm_plusplus.py

@@ -467,8 +467,8 @@ class ESMplusplusForMaskedLM(PreTrainedModel):
     Implements the base ESM++ architecture with a masked language modeling head.
     """
     config_class = ESMplusplusConfig
-    def __init__(self, config: ESMplusplusConfig):
-        super().__init__(config)
+    def __init__(self, config: ESMplusplusConfig, **kwargs):
+        super().__init__(config, **kwargs)
         self.config = config
         self.vocab_size = config.vocab_size
         self.embed = nn.Embedding(self.vocab_size, config.hidden_size)
@@ -642,10 +642,10 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM):
     
     Extends the base ESM++ model with a classification head.
     """
-    def __init__(self, config: ESMplusplusConfig, num_labels: int=2):
-        config.num_labels = num_labels
-        super().__init__(config)
+    def __init__(self, config: ESMplusplusConfig, **kwargs):
+        super().__init__(config, **kwargs)
         self.config = config
+        self.num_labels = config.num_labels
         self.classifier = RegressionHead(config.hidden_size * 2, config.num_labels, config.hidden_size * 4)
         # Large intermediate projections help with sequence classification tasks (*4)
         self.mse = nn.MSELoss()
@@ -715,8 +715,7 @@ class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM):
     
     Extends the base ESM++ model with a token classification head.
     """
-    def __init__(self, config: ESMplusplusConfig, num_labels: int=2):
-        config.num_labels = num_labels
+    def __init__(self, config: ESMplusplusConfig):
         super().__init__(config)
         self.config = config
         self.num_labels = config.num_labels