"""Wren ASR configuration — transformers-compatible."""
from transformers import PretrainedConfig


class WrenASRConfig(PretrainedConfig):
    model_type = "wren_asr"

    def __init__(
        self,
        llm_name:        str = "Qwen/Qwen2.5-0.5B",
        mimi_model_name: str = "kyutai/mimi",
        k_codebooks:     int = 3,
        codebook_size:   int = 2048,
        vocab_size:      int = 151944,
        # Special-token IDs (in the resized text vocab)
        audio_start_id:  int = None,    # <|audio_start|> — opens audio prefix
        audio_end_id:    int = None,    # <|audio_end|>   — closes audio prefix; text begins after
        eos_token_id:    int = None,    # end of transcript (LLM's existing eos)
        sampling_rate:   int = 24000,
        **kwargs,
    ):
        self.llm_name        = llm_name
        self.mimi_model_name = mimi_model_name
        self.k_codebooks     = k_codebooks
        self.codebook_size   = codebook_size
        self.vocab_size      = vocab_size
        self.audio_start_id  = audio_start_id
        self.audio_end_id    = audio_end_id
        self.sampling_rate   = sampling_rate
        super().__init__(eos_token_id=eos_token_id, **kwargs)