131 lines
4.5 KiB
Python
131 lines
4.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
|
|
class GPTConfig:
|
|
"""Configuration for our GPT model."""
|
|
def __init__(
|
|
self,
|
|
vocab_size: int,
|
|
block_size: int,
|
|
n_layer: int = 8,
|
|
n_head: int = 8,
|
|
n_embd: int = 512,
|
|
):
|
|
self.vocab_size = vocab_size
|
|
self.block_size = block_size
|
|
self.n_layer = n_layer
|
|
self.n_head = n_head
|
|
self.n_embd = n_embd
|
|
|
|
|
|
class CausalSelfAttention(nn.Module):
|
|
"""A single multi-head causal self-attention layer."""
|
|
def __init__(self, config: GPTConfig):
|
|
super().__init__()
|
|
assert config.n_embd % config.n_head == 0
|
|
self.key = nn.Linear(config.n_embd, config.n_embd)
|
|
self.query = nn.Linear(config.n_embd, config.n_embd)
|
|
self.value = nn.Linear(config.n_embd, config.n_embd)
|
|
self.proj = nn.Linear(config.n_embd, config.n_embd)
|
|
self.n_head = config.n_head
|
|
self.head_dim = config.n_embd // config.n_head
|
|
# causal mask, buffer not a parameter
|
|
mask = torch.tril(torch.ones(config.block_size, config.block_size))
|
|
self.register_buffer("mask", mask)
|
|
|
|
def forward(self, x):
|
|
B, T, C = x.size()
|
|
k = self.key(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
|
|
q = self.query(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
|
|
# compute attention scores
|
|
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)
|
|
v = self.value(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
|
|
out = att @ v
|
|
out = out.transpose(1, 2).contiguous().view(B, T, C)
|
|
return self.proj(out)
|
|
|
|
|
|
class MLP(nn.Module):
|
|
"""Feed-forward layer."""
|
|
def __init__(self, config: GPTConfig):
|
|
super().__init__()
|
|
self.fc1 = nn.Linear(config.n_embd, 4 * config.n_embd)
|
|
self.fc2 = nn.Linear(4 * config.n_embd, config.n_embd)
|
|
|
|
def forward(self, x):
|
|
return self.fc2(F.gelu(self.fc1(x)))
|
|
|
|
|
|
class Block(nn.Module):
|
|
"""Transformer block: attention + feed-forward."""
|
|
def __init__(self, config: GPTConfig):
|
|
super().__init__()
|
|
self.ln1 = nn.LayerNorm(config.n_embd)
|
|
self.ln2 = nn.LayerNorm(config.n_embd)
|
|
self.attn = CausalSelfAttention(config)
|
|
self.mlp = MLP(config)
|
|
|
|
def forward(self, x):
|
|
x = x + self.attn(self.ln1(x))
|
|
x = x + self.mlp(self.ln2(x))
|
|
return x
|
|
|
|
|
|
class GPT(nn.Module):
|
|
"""GPT language model from scratch."""
|
|
def __init__(self, config: GPTConfig):
|
|
super().__init__()
|
|
self.token_emb = nn.Embedding(config.vocab_size, config.n_embd)
|
|
self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
|
|
self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
|
|
self.ln_f = nn.LayerNorm(config.n_embd)
|
|
self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
|
|
self.block_size = config.block_size
|
|
|
|
def forward(self, idx, targets=None):
|
|
B, T = idx.size()
|
|
tok_emb = self.token_emb(idx) # (B,T,C)
|
|
pos_emb = self.pos_emb[:, :T, :] # (1,T,C)
|
|
x = tok_emb + pos_emb
|
|
for block in self.blocks:
|
|
x = block(x)
|
|
x = self.ln_f(x)
|
|
logits = self.head(x) # (B,T,vocab)
|
|
loss = None
|
|
if targets is not None:
|
|
# shift logits and targets for next-token prediction
|
|
logits = logits.view(B * T, -1)
|
|
targets = targets.view(B * T)
|
|
loss = F.cross_entropy(logits, targets)
|
|
return logits, loss
|
|
|
|
@torch.no_grad()
|
|
def generate(
|
|
self,
|
|
idx,
|
|
max_new_tokens: int,
|
|
temperature: float = 1.0,
|
|
top_k: int = None
|
|
):
|
|
"""
|
|
Iteratively predict next token and append to sequence.
|
|
- idx is (B,T) starting context.
|
|
- Returns (B, T+max_new_tokens).
|
|
"""
|
|
for _ in range(max_new_tokens):
|
|
idx_cond = idx[:, -self.block_size :]
|
|
logits, _ = self(idx_cond)
|
|
logits = logits[:, -1, :] / temperature
|
|
if top_k is not None:
|
|
v, _ = torch.topk(logits, top_k)
|
|
logits[logits < v[:, [-1]]] = -float('Inf')
|
|
probs = F.softmax(logits, dim=-1)
|
|
next_token = torch.multinomial(probs, num_samples=1)
|
|
idx = torch.cat([idx, next_token], dim=1)
|
|
return idx
|