from transformers import PretrainedConfig
from typing import Optional

class LlamaEdgeConfig(PretrainedConfig):
    model_type = "llama_edge"

    def __init__(
        self,
        dim: int = 4096,
        n_layers: int = 32,
        n_heads: int = 32,
        n_kv_heads: int = 8,
        vocab_size: int = 9942,
        multiple_of: int = 256,
        ffn_dim_multiplier: Optional[float] = 1.3,
        norm_eps: float = 1e-5,
        rope_theta: float = 500000.0,
        max_seq_len: int = 8192,
        intermediate_size: int = 14336,
        **kwargs,
    ):
        self.dim = dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.vocab_size = vocab_size
        self.multiple_of = multiple_of
        self.ffn_dim_multiplier = ffn_dim_multiplier
        self.norm_eps = norm_eps
        self.rope_theta = rope_theta
        self.max_seq_len = max_seq_len
        self.intermediate_size = intermediate_size

        super().__init__(**kwargs)