from transformers import PreTrainedConfig


class CheddarConfig(PreTrainedConfig):
    model_type = "gouda-cheddar"
    attribute_map = {
        "num_hidden_layers": "n_layer",
        "hidden_size": "n_embd",
        "num_attention_heads": "n_head",
        "num_key_value_heads": "n_kv_head",
        "intermediate_size": "mlp_hidden",
        "max_position_embeddings": "max_seq_len",
    }
    
    def __init__(
        self,
        vocab_size=32768,
        max_seq_len=4096,
        n_layer=44,
        n_embd=768,
        n_head=12,
        n_kv_head=4,
        mlp_hidden=2048,
        rope_base=10000.0,
        rms_eps=1e-6,
        qk_norm=True,
        z_loss=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.n_layer = n_layer
        self.n_embd = n_embd
        self.n_head = n_head
        self.n_kv_head = n_kv_head
        self.mlp_hidden = mlp_hidden
        self.rope_base = rope_base
        self.rms_eps = rms_eps
        self.qk_norm = qk_norm
        self.z_loss = z_loss

        if self.n_embd % self.n_head:
            raise ValueError("n_embd must be divisible by n_head")
        if self.n_head % self.n_kv_head:
            raise ValueError("n_head must be divisible by n_kv_head")

    @property
    def head_dim(self):
        return self.n_embd // self.n_head