import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers import initialization as init from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast, SequenceClassifierOutput from transformers.cache_utils import Cache, DynamicCache from .configuration_neolm import NeoLMConfig # ───────────────────────────────────────────── # RMSNorm # ───────────────────────────────────────────── class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.weight * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) # ───────────────────────────────────────────── # RoPE # ───────────────────────────────────────────── def precompute_rope_freqs(head_dim: int, seq_len: int, base: int = 10_000): assert head_dim % 2 == 0 # float64 عشان نتجنب overflow inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).double() / head_dim)) t = torch.arange(seq_len).double() freqs = torch.outer(t, inv_freq) return freqs.cos().float(), freqs.sin().float() def apply_rope( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, offset: int = 0, base: int = 10_000, ): B, H, T, D = x.shape needed = offset + T if needed > cos.shape[0]: cos, sin = precompute_rope_freqs( head_dim=D, seq_len=needed + 64, base=base, ) cos = cos[offset:offset + T].to( device=x.device, dtype=x.dtype )[None, None] sin = sin[offset:offset + T].to( device=x.device, dtype=x.dtype )[None, None] x1 = x[..., ::2] x2 = x[..., 1::2] out = torch.empty_like(x) out[..., ::2] = x1 * cos - x2 * sin out[..., 1::2] = x1 * sin + x2 * cos return out # ───────────────────────────────────────────── # Causal Self-Attention (مع KV Cache) # ───────────────────────────────────────────── class CausalSelfAttention(nn.Module): def __init__(self, config: NeoLMConfig): super().__init__() self.num_heads = config.num_attention_heads self.hidden_size = config.hidden_size self.head_dim = config.hidden_size // config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = ( self.num_heads // self.num_key_value_heads ) self.kv_dim = self.num_key_value_heads * self.head_dim self.qk_norm = config.qk_norm self.q_norm = ( RMSNorm(self.head_dim, config.rms_norm_eps) if self.qk_norm else None ) self.k_norm = ( RMSNorm(self.head_dim, config.rms_norm_eps) if self.qk_norm else None ) self.rope_base = config.rope_base self.dropout = config.dropout self.qkv = nn.Linear(config.hidden_size, config.hidden_size + 2 * self.kv_dim, bias=False) self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) self.proj._residual_projection = config.scale_residual_init cos, sin = precompute_rope_freqs( self.head_dim, config.max_position_embeddings, config.rope_base ) # persistent=False → مش بتتحفظ في الـ checkpoint، بتتحسب من جديد عند الـ init self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) def forward( self, x: torch.Tensor, past_key_value: Cache | None = None, use_cache: bool = False, layer_idx: int = 0, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, Cache | None]: B, T, C = x.shape q, k, v = self.qkv(x).split( [ self.hidden_size, self.kv_dim, self.kv_dim, ], dim=-1, ) def reshape(t, num_heads): return t.view( B, T, num_heads, self.head_dim ).transpose(1, 2) q = reshape(q, self.num_heads) k = reshape(k, self.num_key_value_heads) v = reshape(v, self.num_key_value_heads) # QK-Norm: normalize queries and keys per attention head if self.qk_norm: q = self.q_norm(q) k = self.k_norm(k) # offset = عدد الـ tokens المحفوظة في الـ cache offset = 0 if use_cache and past_key_value is not None: offset = past_key_value.get_seq_length(layer_idx) q = apply_rope( q, self.rope_cos, self.rope_sin, offset=offset, base=self.rope_base, ) k = apply_rope( k, self.rope_cos, self.rope_sin, offset=offset, base=self.rope_base, ) # KV Cache: update بيضيف الـ current k,v وبيرجع الكل if use_cache and past_key_value is not None: k, v = past_key_value.update(k, v, layer_idx) # Causal + Padding Mask key_length = k.size(-2) if attention_mask is not None: if ( attention_mask.ndim != 2 or attention_mask.shape[0] != B or attention_mask.shape[1] < key_length ): raise ValueError( "Invalid attention_mask shape" ) valid_keys = attention_mask[ :, -key_length: ].to(device=q.device, dtype=torch.bool) q_pos = offset + torch.arange( T, device=q.device ) k_pos = torch.arange( key_length, device=q.device ) causal = k_pos[None, :] <= q_pos[:, None] attn_mask = ( valid_keys[:, None, None, :] & causal[None, None, :, :] ) is_causal = False elif offset == 0: attn_mask = None is_causal = True elif T == 1: attn_mask = None is_causal = False else: q_pos = offset + torch.arange(T, device=q.device) k_pos = torch.arange(key_length, device=q.device) attn_mask = k_pos[None, :] <= q_pos[:, None] is_causal = False y = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal, enable_gqa=( self.num_heads != self.num_key_value_heads ), ) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y), past_key_value # ───────────────────────────────────────────── # SwiGLU MLP # ───────────────────────────────────────────── class SwiGLU(nn.Module): def __init__(self, config: NeoLMConfig): super().__init__() # intermediate_size محسوب في الـ config مرة واحدة hidden = config.intermediate_size self.w1 = nn.Linear(config.hidden_size, hidden, bias=False) self.w2 = nn.Linear(hidden, config.hidden_size, bias=False) self.w3 = nn.Linear(config.hidden_size, hidden, bias=False) self.w2._residual_projection = config.scale_residual_init self.drop = nn.Dropout(config.dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x))) # ───────────────────────────────────────────── # Transformer Block # ───────────────────────────────────────────── class NeoLMBlock(nn.Module): def __init__(self, config: NeoLMConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.ln1 = RMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = CausalSelfAttention(config) self.ln2 = RMSNorm(config.hidden_size, config.rms_norm_eps) self.mlp = SwiGLU(config) def forward( self, x: torch.Tensor, past_key_value: Cache | None = None, use_cache: bool = False, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, Cache | None]: attn_out, past_key_value = self.attn( self.ln1(x), past_key_value=past_key_value, use_cache=use_cache, layer_idx=self.layer_idx, attention_mask=attention_mask, ) x = x + attn_out x = x + self.mlp(self.ln2(x)) return x, past_key_value # ───────────────────────────────────────────── # Base PreTrainedModel # ───────────────────────────────────────────── class NeoLMPreTrainedModel(PreTrainedModel): config_class = NeoLMConfig base_model_prefix = "model" supports_gradient_checkpointing = True def _init_weights(self, module: nn.Module): if isinstance(module, nn.Linear): std = 0.02 if ( self.config.scale_residual_init and getattr(module, "_residual_projection", False) ): std /= math.sqrt( 2 * self.config.num_hidden_layers ) init.normal_( module.weight, mean=0.0, std=std, ) if module.bias is not None: init.zeros_(module.bias) elif isinstance(module, nn.Embedding): init.normal_( module.weight, mean=0.0, std=0.02, ) elif isinstance(module, RMSNorm): init.ones_(module.weight) # ───────────────────────────────────────────── # NeoLMModel — Backbone # ───────────────────────────────────────────── class NeoLMModel(NeoLMPreTrainedModel): def __init__(self, config: NeoLMConfig): super().__init__(config) self.token_emb = nn.Embedding(config.vocab_size, config.hidden_size) self.blocks = nn.ModuleList([ NeoLMBlock(config, layer_idx=i) for i in range(config.num_hidden_layers) ]) self.ln_f = RMSNorm(config.hidden_size, config.rms_norm_eps) self.post_init() def get_input_embeddings(self): return self.token_emb def set_input_embeddings(self, value): self.token_emb = value def forward( self, input_ids: torch.Tensor, past_key_values: Cache | None = None, use_cache: bool = False, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, Cache | None]: B, T = input_ids.shape x = self.token_emb(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache() for block in self.blocks: x, past_key_values = block( x, past_key_value=past_key_values, use_cache=use_cache, attention_mask=attention_mask, ) x = self.ln_f(x) return x, past_key_values # ───────────────────────────────────────────── # _TiedWeightsKeys — متوافق مع كل إصدارات HF # ───────────────────────────────────────────── class _TiedWeightsKeys(dict): """ الإصدارات الجديدة تعمل .keys() → dict ✅ الإصدارات القديمة تعمل iteration → dict iteration = keys ✅ """ pass # ───────────────────────────────────────────── # NeoLMForCausalLM — الموديل الكامل # ───────────────────────────────────────────── class NeoLMForCausalLM(NeoLMPreTrainedModel, GenerationMixin): _tied_weights_keys = _TiedWeightsKeys({"lm_head.weight": "model.token_emb.weight"}) _keys_to_ignore_on_load_missing = [r"lm_head\.weight"] def __init__(self, config: NeoLMConfig): super().__init__(config) self.model = NeoLMModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() self._recompute_rope_buffers() # ── RoPE recompute (يمنع garbage values من القديم checkpoints) ── def _recompute_rope_buffers(self): device = next(self.parameters()).device for block in self.model.blocks: attn = block.attn cos, sin = precompute_rope_freqs( attn.head_dim, self.config.max_position_embeddings, self.config.rope_base, ) attn.rope_cos = cos.to(device) attn.rope_sin = sin.to(device) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) model._recompute_rope_buffers() return model # ── Embeddings helpers ─────────────────────────────────────── def get_input_embeddings(self): return self.model.token_emb def set_input_embeddings(self, value): self.model.token_emb = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, value): self.lm_head = value def tie_weights(self, **kwargs): if self.config.tie_word_embeddings: self.lm_head.weight = self.model.token_emb.weight super().tie_weights(**kwargs) # ── Forward ────────────────────────────────────────────────── def forward( self, input_ids: torch.Tensor, labels: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, # للتوافق مع Trainer past_key_values: Cache | None = None, use_cache: bool | None = None, **kwargs, ) -> CausalLMOutputWithPast: use_cache = use_cache if use_cache is not None else self.config.use_cache # لا نستخدم cache وقت التدريب if self.training: use_cache = False hidden, past_key_values = self.model( input_ids, past_key_values=past_key_values, use_cache=use_cache, attention_mask=attention_mask, ) logits = self.lm_head(hidden) loss = None if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_key_values if use_cache else None, hidden_states=(hidden,) if kwargs.get("output_hidden_states", False) else None ) # return CausalLMOutputWithPast( # loss=loss, # logits=logits, # past_key_values=past_key_values if use_cache else None, # ) # ── Generate helpers ───────────────────────────────────────── def prepare_inputs_for_generation( self, input_ids: torch.Tensor, past_key_values: Cache | None = None, **kwargs, ) -> dict: past_len = past_key_values.get_seq_length() if past_key_values is not None else 0 if past_len > 0: # في cache → نبعت آخر token بس input_ids = input_ids[:, -1:] else: # أول call → نقتص لو الـ prompt أطول من max_position_embeddings max_len = self.config.max_position_embeddings if input_ids.shape[1] > max_len: input_ids = input_ids[:, -max_len:] return { "input_ids": input_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache", True), "attention_mask": kwargs.get("attention_mask", None), } class NeoLMForSequenceClassification(NeoLMPreTrainedModel): def __init__(self, config: NeoLMConfig): super().__init__(config) self.num_labels = config.num_labels # backbone self.model = NeoLMModel(config) # classifier head self.score = nn.Linear( config.hidden_size, config.num_labels, bias=True, ) self.dropout = nn.Dropout(config.dropout) self.post_init() self._recompute_rope_buffers() def _recompute_rope_buffers(self): device = next(self.parameters()).device for block in self.model.blocks: attn = block.attn cos, sin = precompute_rope_freqs( attn.head_dim, self.config.max_position_embeddings, self.config.rope_base, ) attn.rope_cos = cos.to(device) attn.rope_sin = sin.to(device) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): model = super().from_pretrained( pretrained_model_name_or_path, *args, **kwargs ) model._recompute_rope_buffers() return model # ───────────────────────────────────────── # Embeddings Helpers # ───────────────────────────────────────── def get_input_embeddings(self): return self.model.token_emb def set_input_embeddings(self, value): self.model.token_emb = value # ───────────────────────────────────────── # Forward # ───────────────────────────────────────── def forward( self, input_ids: torch.Tensor, labels: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> SequenceClassifierOutput: # ===================================== # Backbone # ===================================== hidden, _ = self.model( input_ids=input_ids, attention_mask=attention_mask, use_cache=False, ) # ===================================== # Pooling # ===================================== if attention_mask is not None: mask = attention_mask.unsqueeze(-1).float() pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) else: pooled = hidden.mean(dim=1) pooled = self.dropout(pooled) # ===================================== # Classification Head # ===================================== logits = self.score(pooled) # ===================================== # Loss # ===================================== loss = None if labels is not None: loss = F.cross_entropy( logits.view(-1, self.num_labels), labels.view(-1), ) return SequenceClassifierOutput( loss=loss, logits=logits, hidden_states=(hidden,) if kwargs.get("output_hidden_states", False) else None, )