model_type on top

configuration_bert.py

@@ -40,6 +40,8 @@ class BertConfig(TransformersBertConfig):
 
 
 class FlexBertConfig(TransformersBertConfig):
+    model_type = "flex_bert"
+
     def __init__(
         self,
         attention_layer: str = "base",
@@ -97,7 +99,6 @@ class FlexBertConfig(TransformersBertConfig):
         pad_logits: bool = False,
         compile_model: bool = False,
         masked_prediction: bool = False,
-        model_type: str = "flex_bert",
         **kwargs,
     ):
         """
@@ -214,7 +215,6 @@ class FlexBertConfig(TransformersBertConfig):
         self.pad_logits = pad_logits
         self.compile_model = compile_model
         self.masked_prediction = masked_prediction
-        self.model_type = model_type
 
         if loss_kwargs.get("return_z_loss", False):
             if loss_function != "fa_cross_entropy":