Upload model_v2.py with huggingface_hub

model_v2.py

@@ -0,0 +1,138 @@
+"""
+Cozet: Native SYNAXIM Base Model (Enhanced)
+Architecture: SymbioticGate (M-matrix) + NaN Firewall + Dual-Track M
+(c) 2026 GRRN Research.
+"""
+import torch, torch.nn as nn, torch.nn.functional as F, math
+from dataclasses import dataclass
+
+@dataclass
+class CozConfig:
+    hidden_size: int = 1024
+    num_layers: int = 12
+    num_attention_heads: int = 16
+    num_kv_heads: int = 4
+    intermediate_size: int = 4096
+    vocab_size: int = 50257
+    max_seq_len: int = 4096
+    rope_theta: float = 10000.0
+    rms_norm_eps: float = 1e-6
+    memory_decay: float = 0.995
+    unrotated_decay: float = 0.995
+    tie_word_embeddings: bool = True
+    @property
+    def head_dim(self):
+        return self.hidden_size // self.num_attention_heads
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+    def forward(self, x):
+        return (x.float() * x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()).type_as(x) * self.weight
+
+class NaNFirewall(nn.Module):
+    def __init__(self, clamp_value=1e4, norm_ratio=8.0):
+        super().__init__()
+        self.clamp_value = clamp_value
+        self.norm_ratio = norm_ratio
+    def forward(self, delta, h_ref):
+        delta = torch.nan_to_num(delta, nan=0.0, posinf=self.clamp_value, neginf=-self.clamp_value)
+        d_norm = delta.norm()
+        max_norm = self.norm_ratio * (h_ref.norm() + 1e-8)
+        if d_norm > max_norm:
+            delta = delta * (max_norm / (d_norm + 1e-8))
+        if torch.isnan(delta).all() or torch.isinf(delta).all():
+            delta = torch.zeros_like(delta)
+        return delta
+
+class SymbioticGate(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        D, nh, nk, hd = config.hidden_size, config.num_attention_heads, config.num_kv_heads, config.head_dim
+        self.D, self.n_heads, self.n_kv, self.head_dim = D, nh, nk, hd
+        self.decay = config.memory_decay
+        self.unrot_decay = config.unrotated_decay
+        self.q_proj = nn.Linear(D, nh * hd, bias=False)
+        self.k_proj = nn.Linear(D, nk * hd, bias=False)
+        self.v_proj = nn.Linear(D, nk * hd, bias=False)
+        self.o_proj = nn.Linear(nh * hd, D, bias=False)
+        self.gate_scale = nn.Parameter(torch.ones(1))
+        self.gate_bias = nn.Parameter(torch.zeros(1))
+        self._rc, self._rs = None, None
+    def _rope(self, max_pos, dev):
+        if self._rc is not None and max_pos <= self._rc.shape[0]: return
+        half = self.head_dim // 2
+        f = 1.0 / (10000.0 ** (torch.arange(0, half, device=dev).float() / half))
+        a = torch.outer(torch.arange(max_pos, device=dev).float(), f)
+        self._rc, self._rs = a.cos(), a.sin()
+    def _apply_rope(self, x, pos):
+        half = self.head_dim // 2
+        c, s = self._rc[pos, :half], self._rs[pos, :half]
+        x1, x2 = x[..., :half], x[..., half:]
+        return torch.cat([x1*c - x2*s, x1*s + x2*c], dim=-1)
+    def forward(self, h, M, M_unrot, pos):
+        self._rope(pos + 1, h.device)
+        q, k, v = self.q_proj(h), self.k_proj(h), self.v_proj(h)
+        q_h = self._apply_rope(q.view(self.n_heads, self.head_dim), pos)
+        k_h = self._apply_rope(k.view(self.n_kv, self.head_dim), pos)
+        v_h = v.view(self.n_kv, self.head_dim)
+        if self.n_kv < self.n_heads:
+            r = self.n_heads // self.n_kv
+            k_h, v_h = k_h.repeat_interleave(r, 0), v_h.repeat_interleave(r, 0)
+        g = torch.sigmoid((q_h * k_h).sum(-1).mean() / math.sqrt(self.head_dim) * self.gate_scale + self.gate_bias)
+        kf, vf = k_h.reshape(-1), v_h.reshape(-1)
+        kn = kf / (kf.norm() + 1e-8)
+        vn = vf / (vf.norm() + 1e-8) * h.norm()
+        outer = torch.outer(kn, vn)
+        M2 = g * self.decay * M + (1.0 - g) * outer
+        M_unrot2 = self.unrot_decay * M_unrot + (1.0 - self.unrot_decay) * outer
+        return self.o_proj(q_h.reshape(-1) @ M2), M2, M_unrot2
+
+class SynaxBlock(nn.Module):
+    def __init__(self, config, i):
+        super().__init__()
+        self.norm_attn = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attn = SymbioticGate(config, i)
+        self.norm_mlp = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.firewall = NaNFirewall()
+    def forward(self, h, M, M_unrot, pos):
+        a, M2, M_unrot2 = self.attn(self.norm_attn(h), M, M_unrot, pos)
+        a = self.firewall(a, h)
+        h = h + a
+        n = self.norm_mlp(h)
+        m = self.down_proj(F.silu(self.gate_proj(n)) * self.up_proj(n))
+        m = self.firewall(m, h)
+        h = h + m
+        return h, M2, M_unrot2
+
+class CozModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config, self.D, self.n_layers = config, config.hidden_size, config.num_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([SynaxBlock(config, i) for i in range(config.num_layers)])
+        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        if not config.tie_word_embeddings:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        ds = 1.0 / math.sqrt(2 * self.n_layers)
+        for name, p in self.named_parameters():
+            if p.dim() < 2: continue
+            if "embed" in name: nn.init.normal_(p, std=0.02)
+            elif "down_proj" in name or "o_proj" in name: nn.init.normal_(p, std=0.02 * ds)
+            elif p.dim() == 2: nn.init.normal_(p, std=0.02)
+    def init_m(self, dev):
+        M = [torch.zeros(self.D, self.D, device=dev) for _ in range(self.n_layers)]
+        M_unrot = [torch.zeros(self.D, self.D, device=dev) for _ in range(self.n_layers)]
+        return M, M_unrot
+    def forward_token(self, tid, M, M_unrot, pos):
+        h = self.embed_tokens.weight[tid]
+        for i, layer in enumerate(self.layers):
+            h, M[i], M_unrot[i] = layer(h, M[i], M_unrot[i], pos)
+        h = self.final_norm(h)
+        logits = h @ self.embed_tokens.weight.T if self.config.tie_word_embeddings else self.lm_head(h)
+        return logits, M, M_unrot