from transformers import PretrainedConfig


class PulvisConfig(PretrainedConfig):
    model_type = "pulvis"

    def __init__(self, vocab_size=4096, hidden_size=192, num_attention_heads=6, num_key_value_heads=2, head_dim=32,
                 intermediate_size=368, prelude_layers=1, core_layers=3, core_loops=3, coda_layers=3,
                 max_position_embeddings=1024, rope_theta=10000.0, rms_norm_eps=1e-6, logit_cap=15.0,
                 bos_token_id=0, eos_token_id=0, pad_token_id=1, tie_word_embeddings=False, **kwargs):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.intermediate_size = intermediate_size
        self.prelude_layers = prelude_layers
        self.core_layers = core_layers
        self.core_loops = core_loops
        self.coda_layers = coda_layers
        self.num_hidden_layers = prelude_layers + core_layers * core_loops + coda_layers
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.rms_norm_eps = rms_norm_eps
        self.logit_cap = logit_cap
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id,
                         tie_word_embeddings=tie_word_embeddings, **kwargs)