| """Ethoric config. The short field names (n_layer, d_model, ...) are the real |
| ones; the aliases Transformers and lm-eval look for are set alongside them in |
| __init__. |
| """ |
| from __future__ import annotations |
|
|
| from transformers import PretrainedConfig |
|
|
| |
| |
| DEFAULT_CANON = {"enabled": True, "kernel": 4, "positions": ["A", "C"], "doc_masking": True} |
|
|
|
|
| class EthoricConfig(PretrainedConfig): |
| """Decoder-only LM: RMSNorm, RoPE, grouped-query attention with QK-norm, |
| SwiGLU, Canon convolutions, Peri-LN, tied embeddings. |
| |
| Two fields are easy to trip over. `attn_window_layers` lists the layers |
| that attend to a short local window instead of the full context -- it is |
| part of the trained architecture, and an empty list silently gives you a |
| different model. `canon.doc_masking` only matters when a batch packs several |
| documents behind EOS. |
| """ |
|
|
| model_type = "ethoric" |
| keys_to_ignore_at_inference = [] |
|
|
| def __init__( |
| self, |
| vocab_size: int = 5120, |
| d_model: int = 128, |
| n_layer: int = 12, |
| n_head: int = 4, |
| n_kv_head: int = 1, |
| head_dim: int = 32, |
| ffn_hidden: int = 352, |
| max_seq_len: int = 512, |
| rope_theta: float = 10000.0, |
| norm_eps: float = 1e-5, |
| qk_norm: str = "per_head", |
| peri_ln: bool = True, |
| attn_gate: str = "headwise", |
| attn_window: int = 16, |
| attn_window_layers: list | None = None, |
| canon: dict | None = None, |
| doc_masking: bool = True, |
| tie_embeddings: bool = True, |
| **kwargs, |
| ): |
| self.vocab_size = int(vocab_size) |
| self.d_model = int(d_model) |
| self.n_layer = int(n_layer) |
| self.n_head = int(n_head) |
| self.n_kv_head = int(n_kv_head) |
| self.head_dim = int(head_dim) |
| self.ffn_hidden = int(ffn_hidden) |
| self.max_seq_len = int(max_seq_len) |
| self.rope_theta = float(rope_theta) |
| self.norm_eps = float(norm_eps) |
| self.qk_norm = str(qk_norm) |
| self.peri_ln = bool(peri_ln) |
| self.attn_gate = str(attn_gate) |
| self.attn_window = int(attn_window) |
| self.attn_window_layers = list(attn_window_layers or []) |
| self.canon = {**DEFAULT_CANON, **(canon or {})} |
| self.doc_masking = bool(doc_masking) |
|
|
| |
| |
| |
| self.hidden_size = self.d_model |
| self.num_hidden_layers = self.n_layer |
| self.num_attention_heads = self.n_head |
| self.num_key_value_heads = self.n_kv_head |
| self.intermediate_size = self.ffn_hidden |
| self.max_position_embeddings = self.max_seq_len |
| self.rms_norm_eps = self.norm_eps |
|
|
| kwargs.setdefault("tie_word_embeddings", bool(tie_embeddings)) |
| kwargs.setdefault("bos_token_id", None) |
| self.n_special = int(kwargs.pop("n_special", 2)) |
| kwargs.setdefault("eos_token_id", 0) |
| kwargs.setdefault("pad_token_id", 1) |
| super().__init__(**kwargs) |
|
|
| @property |
| def tie_embeddings(self) -> bool: |
| return bool(self.tie_word_embeddings) |
|
|
| @classmethod |
| def from_pipeline(cls, cfg: dict) -> "EthoricConfig": |
| """Build from a `model:` config block.""" |
| m = dict(cfg["model"]) |
| |
| m.pop("budget", None) |
| |
| |
| tok = cfg.get("tokenizer", {}) or {} |
| eos = int(tok.get("eos_id", 0)) |
| n_special = len(tok.get("special_tokens", []) or []) or 2 |
| pad = next((i for i in range(n_special) if i != eos), None) |
| return cls( |
| tie_embeddings=m.pop("tie_embeddings", True), |
| eos_token_id=eos, |
| pad_token_id=pad, |
| n_special=n_special, |
| **m, |
| ) |
|
|
