from transformers import PretrainedConfig


class DomMLConfig(PretrainedConfig):
    model_type = "dom_ml"

    def __init__(
        self,
        model_name="Synthetic_K3_Euclidean_Ensemble",
        model_file=None,
        feature_names=None,
        model_kind=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.model_name = model_name
        self.model_file = model_file
        self.model_kind = model_kind
        self.feature_names = feature_names or ["mass"]