File size: 6,376 Bytes
67ac79e | 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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """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:
# Le decalage de +1 a deja ete fait dans le chargeur de donnees.
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") # PAD et BOS ne sont pas du texte.
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
|