from transformers import PreTrainedConfig


class GenAIDConfig(PreTrainedConfig):
    model_type = "genaid"

    def __init__(
        self,
        encoder_config=None,
        bottleneck_dim=64,
        num_accents=9,
        num_speakers=336,
        accent_labels=None,
        sampling_rate=16000,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.encoder_config = encoder_config or {}
        self.bottleneck_dim = bottleneck_dim
        self.num_accents = num_accents
        self.num_speakers = num_speakers
        self.accent_labels = accent_labels or []
        self.sampling_rate = sampling_rate