from transformers import Cohere2Config


class TinyAyaConfig(Cohere2Config):
    model_type = "tinyaya_mimi"

    def __init__(
        self,
        num_quantizers=8,
        codebook_size=2048,
        audio_start_token_id=None,
        audio_end_token_id=None,
        speakers=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.num_quantizers = int(num_quantizers)
        self.codebook_size = int(codebook_size)
        self.audio_start_token_id = audio_start_token_id
        self.audio_end_token_id = audio_end_token_id
        self.speakers = speakers or ["Ira", "Aisha", "Siya", "Zoya", "Silver"]