Update modeling_esm_plusplus.py

modeling_esm_plusplus.py

@@ -49,6 +49,7 @@ class ESMplusplusConfig(PretrainedConfig):
         num_labels: int = 2,
         problem_type: str | None = None,
         dropout: float = 0.0,
+        initializer_range: float = 0.02,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -59,6 +60,7 @@ class ESMplusplusConfig(PretrainedConfig):
         self.num_labels = num_labels
         self.problem_type = problem_type
         self.dropout = dropout
+        self.initializer_range = initializer_range
 
 
 ### Rotary Embeddings
@@ -792,7 +794,7 @@ class ESMplusplusForMaskedLM(PreTrainedESMplusplusModel):
         )
 
 
-class ESMplusplusForSequenceClassification(PreTrainedESMplusplusModel):
+class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM):
     """
     ESM++ model for sequence classification.
     Extends the base ESM++ model with a classification head.
@@ -873,7 +875,7 @@ class ESMplusplusForSequenceClassification(PreTrainedESMplusplusModel):
         )
 
 
-class ESMplusplusForTokenClassification(PreTrainedESMplusplusModel):
+class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM):
     """
     ESM++ model for token classification.
     Extends the base ESM++ model with a token classification head.