| import math, torch, torch.nn as nn, torch.nn.functional as F | |
| def sinusoidal_position_embedding(seq_len, dim, device): | |
| pe = torch.zeros(seq_len, dim, device=device) | |
| position = torch.arange(0, seq_len, device=device).unsqueeze(1) | |
| div_term = torch.exp(torch.arange(0, dim, 2, device=device) * (-math.log(10000.0) / dim)) | |
| pe[:, 0::2] = torch.sin(position * div_term); pe[:, 1::2] = torch.cos(position * div_term) | |
| return pe | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, d_model, n_heads, dropout): | |
| super().__init__() | |
| assert d_model % n_heads == 0 | |
| self.n_heads = n_heads; self.head_dim = d_model // n_heads | |
| self.qkv = nn.Linear(d_model, 3*d_model, bias=False) | |
| self.proj = nn.Linear(d_model, d_model, bias=False) | |
| self.attn_drop = nn.Dropout(dropout); self.resid_drop = nn.Dropout(dropout) | |
| self.register_buffer("mask", None, persistent=False) | |
| def forward(self, x): | |
| B, T, C = x.size() | |
| if self.mask is None or self.mask.size(0) < T: | |
| self.mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0).unsqueeze(1) | |
| qkv = self.qkv(x); q, k, v = qkv.split(C, dim=2) | |
| def split(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1,2) | |
| q, k, v = map(split, (q,k,v)) | |
| att = (q @ k.transpose(-2,-1)) / math.sqrt(self.head_dim) | |
| att = att.masked_fill(self.mask[:,:,:T,:T]==0, float('-inf')) | |
| att = F.softmax(att, dim=-1); att = self.attn_drop(att) | |
| y = att @ v; y = y.transpose(1,2).contiguous().view(B, T, C) | |
| return self.resid_drop(self.proj(y)) | |
| class MLP(nn.Module): | |
| def __init__(self, d_model, dropout): | |
| super().__init__() | |
| self.fc1 = nn.Linear(d_model, 4*d_model); self.fc2 = nn.Linear(4*d_model, d_model) | |
| self.drop = nn.Dropout(dropout) | |
| def forward(self, x): | |
| return self.drop(self.fc2(self.drop(F.gelu(self.fc1(x))))) | |
| class Block(nn.Module): | |
| def __init__(self, d_model, n_heads, dropout): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(d_model); self.attn = CausalSelfAttention(d_model, n_heads, dropout) | |
| self.ln2 = nn.LayerNorm(d_model); self.mlp = MLP(d_model, 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, vocab_size, d_model=512, n_heads=8, n_layers=8, n_positions=512, dropout=0.1): | |
| super().__init__() | |
| self.vocab_size = vocab_size; self.n_positions = n_positions | |
| self.tok_emb = nn.Embedding(vocab_size, d_model) | |
| self.pos_emb = nn.Parameter(torch.zeros(n_positions, d_model), requires_grad=False) | |
| self.drop = nn.Dropout(dropout) | |
| self.blocks = nn.ModuleList([Block(d_model, n_heads, dropout) for _ in range(n_layers)]) | |
| self.ln_f = nn.LayerNorm(d_model); self.head = nn.Linear(d_model, vocab_size, bias=False) | |
| with torch.no_grad(): | |
| self.pos_emb.copy_(sinusoidal_position_embedding(n_positions, d_model, device='cpu')) | |
| def forward(self, idx): | |
| B, T = idx.shape; tok = self.tok_emb(idx); pos = self.pos_emb[:T, :].unsqueeze(0).expand(B, -1, -1) | |
| x = self.drop(tok + pos) | |
| for blk in self.blocks: x = blk(x) | |
| x = self.ln_f(x); return self.head(x) | |