File size: 4,834 Bytes
37d8122 886cf7c 37d8122 886cf7c 37d8122 886cf7c 37d8122 886cf7c 827713a 37d8122 886cf7c 37d8122 886cf7c 827713a 886cf7c 827713a 886cf7c 827713a 37d8122 | 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 | """Transformer texte decoder-only de Cortex-0 — écrit à la main.
Assemblage : embeddings de tokens + N blocs pre-norm (attention causale RoPE +
FFN GELU) + LayerNorm finale + tête de langage (`lm_head`).
Le `lm_head` peut être **lié** aux embeddings (weight tying), ce qui réduit le
nombre de paramètres et aide l'entraînement sur petits corpus.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn as nn
from src.model.block import TransformerBlock
from src.model.rope import build_rope_cache
@dataclass
class TextConfig:
vocab_size: int = 32000
n_layer: int = 6
n_head: int = 6
n_embd: int = 384
ctx_len: int = 512
dropout: float = 0.0
bias: bool = False
ffn_mult: int = 4
rope_theta: float = 10000.0
use_rope: bool = True
tie_embeddings: bool = True
init_std: float = 0.02
class CortexTextLM(nn.Module):
"""Modèle de langage autorégressif (decoder-only) de Cortex-0."""
def __init__(self, config: TextConfig):
super().__init__()
self.config = config
self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
self.dropout = nn.Dropout(config.dropout)
self.blocks = nn.ModuleList(
[
TransformerBlock(
config.n_embd,
config.n_head,
dropout=config.dropout,
bias=config.bias,
ffn_mult=config.ffn_mult,
rope_theta=config.rope_theta,
use_rope=config.use_rope,
)
for _ in range(config.n_layer)
]
)
self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
if config.tie_embeddings:
self.lm_head.weight = self.token_embedding.weight
self.apply(self._init_weights)
# Initialisation plus fine sur les projections résiduelles (comme GPT-2).
for name, param in self.named_parameters():
if name.endswith("out_proj.weight") or name.endswith("fc2.weight"):
nn.init.normal_(param, mean=0.0, std=config.init_std / (2 * config.n_layer) ** 0.5)
@staticmethod
def _init_weights(module: nn.Module) -> None:
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)
elif isinstance(module, nn.Embedding):
nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(
self,
input_ids: torch.Tensor | None = None,
targets: torch.Tensor | None = None,
inputs_embeds: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Retourne ``(logits, loss)``.
``input_ids`` : ``(B, T)`` ou ``inputs_embeds`` : ``(B, T, n_embd)``
(permet d'injecter des embeddings multimodaux). ``targets`` : ``(B, T)`` ;
les positions à ignorer (``-100``) ne comptent pas dans la perte.
"""
hidden = self.encode(input_ids=input_ids, inputs_embeds=inputs_embeds)
logits = self.lm_head(hidden)
loss = None
if targets is not None:
loss = nn.functional.cross_entropy(
logits.reshape(-1, logits.size(-1)), targets.reshape(-1),
ignore_index=-100,
)
return logits, loss
def encode(
self,
input_ids: torch.Tensor | None = None,
inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
"""Retourne les états cachés finaux ``(B, T, n_embd)`` (avant la tête)."""
if inputs_embeds is None:
if input_ids is None:
raise ValueError("fournir input_ids ou inputs_embeds")
x = self.dropout(self.token_embedding(input_ids))
else:
x = self.dropout(inputs_embeds)
T = x.size(1)
if T > self.config.ctx_len:
raise ValueError(f"longueur {T} > ctx_len {self.config.ctx_len}")
cos, sin = build_rope_cache(T, self.config.n_embd // self.config.n_head,
self.config.rope_theta, x.device, x.dtype)
for block in self.blocks:
x = block(x, cos, sin)
return self.ln_f(x)
def count_parameters(self, non_embedding: bool = False) -> int:
"""Compte les paramètres (option : hors embeddings de tokens et lm_head)."""
n = sum(p.numel() for p in self.parameters())
if non_embedding:
n -= self.token_embedding.weight.numel()
if not self.config.tie_embeddings:
n -= self.lm_head.weight.numel()
return n
|