from transformers import PretrainedConfig


class HyperStylerConfig(PretrainedConfig):
    model_type = "hyperstyler"

    def __init__(
        self,
        base_model_name="google/t5-v1_1-large",
        styler_config=None,
        max_encoder_len=80,
        max_decoder_len=80,
        style_model_name="AnnaWegmann/Style-Embedding",
        **kwargs,
    ):
        self.base_model_name = base_model_name
        self.styler_config = styler_config
        self.max_encoder_len = max_encoder_len
        self.max_decoder_len = max_decoder_len
        self.style_model_name = style_model_name
        super().__init__(**kwargs)