from transformers import PretrainedConfig


class SuaveMultitaskConfig(PretrainedConfig):
    model_type = "suave_multitask"

    def __init__(
        self,
        base_model_name="roberta-base",
        num_ai_classes=2,
        classifier_dropout=0.1,
        **kwargs,
    ):
        self.base_model_name = base_model_name
        self.num_ai_classes = num_ai_classes
        self.classifier_dropout = classifier_dropout
        super().__init__(**kwargs)