File size: 4,852 Bytes
4acd7cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_sovythos import SovythosConfig
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)
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, persistent=False)
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)
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, is_causal=is_causal)
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.o_proj(out)
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
class SovythosModel(PreTrainedModel, GenerationMixin):
config_class = SovythosConfig
def __init__(self, config):
# حل مشكلة num_hidden_layers بربطها ديناميكياً بالمتغير الحالي
if not hasattr(config, "num_hidden_layers"):
config.num_hidden_layers = getattr(config, "layers", 12)
super().__init__(config)
self.emb = nn.Embedding(config.vocab_size, config.dim)
self.blocks = nn.ModuleList([Block(config.dim, config.heads) for _ in range(config.layers)])
self.norm = RMSNorm(config.dim)
self.fc = nn.Linear(config.dim, config.vocab_size, bias=False)
self.fc.weight = self.emb.weight
self.post_init()
def get_input_embeddings(self):
return self.emb
def set_input_embeddings(self, value):
self.emb = value
def get_output_embeddings(self):
return self.fc
def set_output_embeddings(self, new_embeddings):
self.fc = new_embeddings
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
x = self.emb(input_ids)
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
logits = self.fc(x)
loss = None
if labels is not None:
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits.view(-1, self.config.vocab_size), labels.view(-1))
return CausalLMOutput(loss=loss, logits=logits)
def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
return {
"input_ids": input_ids,
"attention_mask": attention_mask
} |