| """GPT compact : embeddings, attention causale, MLP, connexions residuelles.""" |
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import asdict, dataclass |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| @dataclass |
| class ModelConfig: |
| vocab_size: int |
| context_length: int = 512 |
| n_layers: int = 6 |
| n_heads: int = 6 |
| d_model: int = 384 |
| dropout: float = 0.1 |
|
|
| def __post_init__(self): |
| if min(self.vocab_size, self.context_length, self.n_layers, self.n_heads, self.d_model) <= 0: |
| raise ValueError("Les dimensions du modele doivent etre positives.") |
| if self.d_model % self.n_heads: |
| raise ValueError("d_model doit etre divisible par n_heads.") |
|
|
|
|
| class CausalAttention(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.n_heads = config.n_heads |
| self.dropout = config.dropout |
| self.qkv = nn.Linear(config.d_model, 3 * config.d_model) |
| self.proj = nn.Linear(config.d_model, config.d_model) |
| self.resid_dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, x): |
| batch, length, width = x.shape |
| q, k, v = self.qkv(x).chunk(3, dim=-1) |
| def split_heads(t): |
| return t.view(batch, length, self.n_heads, width // self.n_heads).transpose(1, 2) |
| q, k, v = map(split_heads, (q, k, v)) |
| attended = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0, |
| ) |
| attended = attended.transpose(1, 2).contiguous().view(batch, length, width) |
| return self.resid_dropout(self.proj(attended)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(config.d_model) |
| self.attn = CausalAttention(config) |
| self.ln2 = nn.LayerNorm(config.d_model) |
| self.mlp = nn.Sequential( |
| nn.Linear(config.d_model, 4 * config.d_model), nn.GELU(), |
| nn.Linear(4 * config.d_model, config.d_model), nn.Dropout(config.dropout), |
| ) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| return x + self.mlp(self.ln2(x)) |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.config = config |
| self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) |
| self.position_embedding = nn.Embedding(config.context_length, config.d_model) |
| self.dropout = nn.Dropout(config.dropout) |
| self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)]) |
| self.ln_final = nn.LayerNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| self.lm_head.weight = self.token_embedding.weight |
| self.apply(self._init_weights) |
| for name, param in self.named_parameters(): |
| if name.endswith("attn.proj.weight") or name.endswith("mlp.2.weight"): |
| nn.init.normal_(param, std=0.02 / math.sqrt(2 * config.n_layers)) |
|
|
| @staticmethod |
| def _init_weights(module): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, std=0.02) |
| if isinstance(module, nn.Linear) and module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def hidden_states(self, tokens): |
| length = tokens.size(1) |
| if length > self.config.context_length: |
| raise ValueError("Sequence plus longue que la fenetre de contexte.") |
| positions = torch.arange(length, device=tokens.device) |
| x = self.dropout(self.token_embedding(tokens) + self.position_embedding(positions)) |
| for block in self.blocks: |
| x = block(x) |
| return self.ln_final(x) |
|
|
| def forward(self, tokens, targets=None): |
| logits = self.lm_head(self.hidden_states(tokens)) |
| loss = None |
| if targets is not None: |
| |
| loss = F.cross_entropy(logits.float().reshape(-1, self.config.vocab_size), targets.reshape(-1), ignore_index=-100) |
| return logits, loss |
|
|
| def parameter_count(self): |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def optimizer(self, learning_rate, weight_decay, device): |
| decay, no_decay = [], [] |
| for param in self.parameters(): |
| (decay if param.ndim >= 2 else no_decay).append(param) |
| return torch.optim.AdamW( |
| [{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}], |
| lr=learning_rate, betas=(0.9, 0.95), fused=(device.type == "cuda"), |
| ) |
|
|
| @torch.inference_mode() |
| def generate(self, tokens, max_new_tokens=100, temperature=0.8, top_k=50, top_p=0.95, eos_id=2): |
| if max_new_tokens < 1 or not math.isfinite(temperature) or temperature < 0: |
| raise ValueError("max_new_tokens positif et temperature >= 0 attendus.") |
| if top_k < 0 or not 0 < top_p <= 1: |
| raise ValueError("top_k >= 0 et 0 < top_p <= 1 attendus.") |
| self.eval() |
| for _ in range(max_new_tokens): |
| logits, _ = self(tokens[:, -self.config.context_length:]) |
| next_logits = logits[:, -1, :].float() |
| next_logits[:, :2] = -float("inf") |
| if temperature == 0: |
| next_token = next_logits.argmax(dim=-1, keepdim=True) |
| else: |
| next_logits /= temperature |
| if top_k: |
| cutoff = torch.topk(next_logits, min(top_k, next_logits.size(-1))).values[:, -1:] |
| next_logits.masked_fill_(next_logits < cutoff, -float("inf")) |
| if top_p < 1: |
| sorted_logits, indices = next_logits.sort(descending=True) |
| remove = sorted_logits.softmax(-1).cumsum(-1) > top_p |
| remove[:, 1:] = remove[:, :-1].clone() |
| remove[:, 0] = False |
| next_logits.scatter_(1, indices, sorted_logits.masked_fill(remove, -float("inf"))) |
| next_token = torch.multinomial(next_logits.softmax(-1), num_samples=1) |
| tokens = torch.cat((tokens, next_token), dim=1) |
| if (next_token == eos_id).all(): |
| break |
| return tokens |
|
|
