| from transformers import LlamaConfig | |
| class BlazeConfig(LlamaConfig): | |
| model_type = "blaze" | |
| def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False, | |
| prelude_layers=1, recurrent_layers=12, coda_layers=1, | |
| recurrent_passes=2, | |
| gradient_checkpointing=False, 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.recurrent_passes = recurrent_passes | |
| 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} | |