from transformers import LlamaConfig

class EmberConfig(LlamaConfig):
    model_type = "ember"
    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
                 prelude_layers=1, recurrent_layers=3, coda_layers=4,
                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
        kwargs.setdefault("use_cache", False)
        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
        self.xsa_projection = xsa_projection
        self.rope_theta = rope_theta
        self.attention_bias = attention_bias
        self.prelude_layers = prelude_layers
        self.recurrent_layers = recurrent_layers
        self.coda_layers = coda_layers
        self.gradient_checkpointing = gradient_checkpointing
        self.use_flash_attn = use_flash_attn
        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}