| 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} | |