import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutputWithPooling from .configuration_rne_tiny_gpt import RNETinyGPTConfig class CausalSelfAttention(nn.Module): def __init__(self, config: RNETinyGPTConfig): super().__init__() if config.n_embd % config.n_head != 0: raise ValueError("n_embd must be divisible by n_head") self.n_head = config.n_head self.head_dim = config.n_embd // config.n_head self.attention_backend = getattr(config, "attention_backend", "sage") self.torch_fallback = bool(getattr(config, "torch_fallback", False)) if self.attention_backend not in ("sage", "torch"): raise ValueError("attention_backend must be 'sage' or 'torch'") if self.attention_backend == "sage" and self.head_dim not in (64, 96, 128): raise ValueError(f"SageAttention requires head_dim in [64, 96, 128], got {self.head_dim}.") if self.attention_backend == "sage" and config.dropout != 0.0: raise ValueError("SageAttention strict mode requires dropout=0.0") self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False) self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False) self.dropout = nn.Dropout(config.dropout) mask = torch.tril(torch.ones(config.ctx_len, config.ctx_len, dtype=torch.bool)) self.register_buffer("mask", mask.view(1, 1, config.ctx_len, config.ctx_len), persistent=False) self.sageattn = None if self.attention_backend == "sage": try: from sageattention import sageattn self.sageattn = sageattn except Exception as exc: if self.torch_fallback: self.attention_backend = "torch" self.sageattn = None else: raise RuntimeError( "Ce modèle a été entraîné avec SageAttention. " "Installe sageattention: pip install sageattention" ) from exc def _torch_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, t: int) -> torch.Tensor: scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) scores = scores.masked_fill(self.mask[:, :, :t, :t] == 0, float("-inf")) att = F.softmax(scores.float(), dim=-1).to(q.dtype) att = self.dropout(att) y = att @ v return y def _sage_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): if self.sageattn is None: raise RuntimeError("SageAttention demandé mais sageattn est None") if not q.is_cuda: if self.torch_fallback: return None raise RuntimeError( "SageAttention exige CUDA. Passe le modèle sur CUDA avec model.cuda(), " "ou active torch_fallback dans config.json." ) q = q.contiguous() k = k.contiguous() v = v.contiguous() return self.sageattn(q, k, v, tensor_layout="HND", is_causal=True) def forward(self, x: torch.Tensor) -> torch.Tensor: b, t, c = x.shape qkv = self.qkv(x) q, k, v = qkv.chunk(3, dim=-1) q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous() k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous() v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous() if self.attention_backend == "sage": y = self._sage_attention(q, k, v) if y is None: y = self._torch_attention(q, k, v, t) else: y = self._torch_attention(q, k, v, t) y = y.transpose(1, 2).contiguous().view(b, t, c) y = self.proj(y) return y class MLP(nn.Module): def __init__(self, config: RNETinyGPTConfig): super().__init__() self.fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False) self.proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False) self.dropout = nn.Dropout(config.dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.fc(x) x = F.gelu(x) x = self.proj(x) x = self.dropout(x) return x class Block(nn.Module): def __init__(self, config: RNETinyGPTConfig): super().__init__() self.ln1 = nn.LayerNorm(config.n_embd) self.attn = CausalSelfAttention(config) self.ln2 = nn.LayerNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class RNETinyGPTPreTrainedModel(PreTrainedModel): config_class = RNETinyGPTConfig base_model_prefix = "rne_tiny_gpt" supports_gradient_checkpointing = False 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) if isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) class RNETinyGPTModel(RNETinyGPTPreTrainedModel): def __init__(self, config: RNETinyGPTConfig): super().__init__(config) self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd) self.pos_emb = nn.Embedding(config.ctx_len, config.n_embd) self.drop = nn.Dropout(config.dropout) self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)]) self.ln_f = nn.LayerNorm(config.n_embd) self.post_init() def _mean_pool(self, hidden: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor) -> torch.Tensor: if attention_mask is None: mask = input_ids.ne(self.config.pad_token_id) else: mask = attention_mask.bool() mask = mask.unsqueeze(-1).to(hidden.dtype) summed = (hidden * mask).sum(dim=1) denom = mask.sum(dim=1).clamp(min=1.0) return summed / denom def _last_pool(self, hidden: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor) -> torch.Tensor: if attention_mask is None: mask = input_ids.ne(self.config.pad_token_id) else: mask = attention_mask.bool() lengths = mask.sum(dim=1).clamp(min=1) last_pos = lengths - 1 batch_idx = torch.arange(input_ids.size(0), device=input_ids.device) return hidden[batch_idx, last_pos, :] def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, return_dict: Optional[bool] = True, **kwargs, ): b, t = input_ids.shape if t > self.config.ctx_len: raise ValueError(f"Input length {t} > ctx_len {self.config.ctx_len}. Truncate before calling the model.") pos = torch.arange(0, t, dtype=torch.long, device=input_ids.device).unsqueeze(0) x = self.tok_emb(input_ids) + self.pos_emb(pos) x = self.drop(x) for block in self.blocks: x = block(x) hidden = self.ln_f(x) if self.config.pooling == "last": pooled = self._last_pool(hidden, attention_mask, input_ids) else: pooled = self._mean_pool(hidden, attention_mask, input_ids) pooled = pooled.float() if self.config.normalize_embeddings: pooled = F.normalize(pooled, p=2, dim=-1) if not return_dict: return (hidden, pooled) return BaseModelOutputWithPooling( last_hidden_state=hidden, pooler_output=pooled, hidden_states=None, attentions=None, ) @torch.no_grad() def encode(self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: out = self.forward(input_ids=input_ids, attention_mask=attention_mask, return_dict=True) return out.pooler_output