"""
Configuration class for TaoNet model.
"""

from transformers import PretrainedConfig


class TaoNetConfig(PretrainedConfig):
    """Configuration for TaoNet model."""
    
    model_type = "taonet"
    
    def __init__(
        self,
        vocab_size: int = 25000,
        d_model: int = 512,
        d_embed_rank: int = 384,
        d_state: int = 512,
        d_ff: int = 512,
        n_heads: int = 4,
        d_kv_comp: int = 384,
        d_rope: int = 64,
        n_layers: int = 8,
        max_seq_len: int = 256,
        dropout: float = 0.02,
        block_arrangement: str = "layered",
        ssm_per_mla: int = 3,
        layered_mla_num: int = 0,
        pad_token_id: int = 3,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        unk_token_id: int = 0,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            unk_token_id=unk_token_id,
            **kwargs,
        )
        
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.d_embed_rank = d_embed_rank
        self.d_state = d_state
        self.d_ff = d_ff
        self.n_heads = n_heads
        self.d_kv_comp = d_kv_comp
        self.d_rope = d_rope
        self.n_layers = n_layers
        self.max_seq_len = max_seq_len
        self.dropout = dropout
        self.block_arrangement = block_arrangement
        self.ssm_per_mla = ssm_per_mla
        self.layered_mla_num = layered_mla_num