model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as ckpt
import math

# ============================================================
# SOVYTHOS — Sovereign Egyptian Reasoning/Code Model
# نسخة مصرية مصغّرة، معمارية Titan (RMSNorm + RoPE + SwiGLU + GQA-ready)
# مبنية للاستدلال والبرمجة كأولوية، والعربي/المصري/الانجليزي كطبقة لغوية فوقها
# ============================================================
MODEL_IDENTITY = "SOVYTHOS"

# ========= RMSNorm =========
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.w = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True)
        return self.w * x * torch.rsqrt(rms + self.eps)

# ========= RoPE (Cached — أسرع في الـ inference) =========
class RoPE(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq)
        self._cos_cache = None
        self._sin_cache = None

    def _build_cache(self, seq_len, device):
        if self._cos_cache is not None and self._cos_cache.shape[0] >= seq_len:
            return
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self._cos_cache = emb.cos()
        self._sin_cache = emb.sin()

    def forward(self, x, seq_len):
        self._build_cache(seq_len, x.device)
        cos = self._cos_cache[:seq_len][None, None, :, :]
        sin = self._sin_cache[:seq_len][None, None, :, :]
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return (x * cos) + (torch.cat((-x2, x1), dim=-1) * sin)

# ========= TitanAttention =========
# الـ shapes متطابقة 100% مع الـ checkpoint:
#   k_proj: [512, 512]  v_proj: [512, 512]
class TitanAttention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.o_proj = nn.Linear(dim, dim, bias=False)

        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.rope = RoPE(self.head_dim)

    def forward(self, x, is_causal=True):
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)

        q = self.rope(self.q_norm(q), T)
        k = self.rope(self.k_norm(k), T)

        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=0.1 if self.training else 0.0,
            is_causal=is_causal
        )

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)

# ========= Block =========
# الـ shapes متطابقة 100% مع الـ checkpoint:
#   w1/w2: [2048, 512]   w3: [512, 2048]
class Block(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.n1 = RMSNorm(dim)
        self.attn = TitanAttention(dim, heads)
        self.n2 = RMSNorm(dim)

        self.w1 = nn.Linear(dim, 4 * dim, bias=False)
        self.w2 = nn.Linear(dim, 4 * dim, bias=False)
        self.w3 = nn.Linear(4 * dim, dim, bias=False)

    def forward(self, x):
        x = x + self.attn(self.n1(x))
        h = self.n2(x)
        x = x + self.w3(F.silu(self.w1(h)) * self.w2(h))
        return x

# ========= Sovereign V16 Titan =========
class Model(nn.Module):
    def __init__(self, vocab_size, dim=512, heads=16, layers=12, use_grad_checkpoint=False):
        super().__init__()
        self.identity = MODEL_IDENTITY
        self.vocab_size = vocab_size
        self.dim = dim
        self.use_grad_checkpoint = use_grad_checkpoint  # وفّر VRAM على T4 لو True (أبطأ ~20% بس بيسمح بـ batch أكبر)

        self.emb = nn.Embedding(vocab_size, dim)
        self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(layers)])
        self.norm = RMSNorm(dim)
        self.fc = nn.Linear(dim, vocab_size, bias=False)

        # Weight Tying
        self.fc.weight = self.emb.weight

        self.apply(self._init_weights)

        n_params = sum(p.numel() for p in self.parameters())
        print(f"🔱 {self.identity} | {layers}L/{heads}H | Dim:{dim} | Params: {n_params/1e6:.1f}M")

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, x):
        x = self.emb(x)
        for blk in self.blocks:
            if self.use_grad_checkpoint and self.training:
                x = ckpt.checkpoint(blk, x, use_reentrant=False)
            else:
                x = blk(x)
        x = self.norm(x)
        return self.fc(x)