| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import math |
|
|
|
|
| class MultiHeadAttention(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| assert config.n_embd % config.n_head == 0 |
| self.n_head = config.n_head |
| self.n_embd = config.n_embd |
| self.head_dim = config.n_embd // config.n_head |
| self.dropout = config.dropout |
|
|
| self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) |
| self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) |
| self.attn_dropout = nn.Dropout(config.dropout) |
| self.resid_dropout = nn.Dropout(config.dropout) |
|
|
| self.register_buffer( |
| "mask", |
| torch.tril(torch.ones(config.block_size, config.block_size)).view( |
| 1, 1, config.block_size, config.block_size |
| ), |
| ) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| q, k, v = self.c_attn(x).split(self.n_embd, dim=2) |
|
|
| q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| scale = 1.0 / math.sqrt(self.head_dim) |
| attn = (q @ k.transpose(-2, -1)) * scale |
| attn = attn.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf")) |
| attn = F.softmax(attn, dim=-1) |
| attn = self.attn_dropout(attn) |
|
|
| out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) |
| return self.resid_dropout(self.c_proj(out)) |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias), |
| nn.GELU(), |
| nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias), |
| nn.Dropout(config.dropout), |
| ) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(config.n_embd, bias=config.bias) |
| self.attn = MultiHeadAttention(config) |
| self.ln2 = nn.LayerNorm(config.n_embd, bias=config.bias) |
| self.ff = FeedForward(config) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.ff(self.ln2(x)) |
| return x |
|
|
|
|
| class GPTConfig: |
| def __init__( |
| self, |
| vocab_size=65, |
| block_size=256, |
| n_layer=6, |
| n_head=6, |
| n_embd=384, |
| dropout=0.2, |
| bias=True, |
| ): |
| 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 |
| self.dropout = dropout |
| self.bias = bias |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
|
|
| self.transformer = nn.ModuleDict( |
| { |
| "wte": nn.Embedding(config.vocab_size, config.n_embd), |
| "wpe": nn.Embedding(config.block_size, config.n_embd), |
| "drop": nn.Dropout(config.dropout), |
| "h": nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)]), |
| "ln_f": nn.LayerNorm(config.n_embd, bias=config.bias), |
| } |
| ) |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
| self.transformer.wte.weight = self.lm_head.weight |
|
|
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| 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, idx, targets=None): |
| B, T = idx.shape |
| assert T <= self.config.block_size |
|
|
| pos = torch.arange(0, T, dtype=torch.long, device=idx.device) |
| tok_emb = self.transformer.wte(idx) |
| pos_emb = self.transformer.wpe(pos) |
| x = self.transformer.drop(tok_emb + pos_emb) |
|
|
| for block in self.transformer.h: |
| x = block(x) |
| x = self.transformer.ln_f(x) |
|
|
| if targets is not None: |
| logits = self.lm_head(x) |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) |
| return logits, loss |
|
|
| logits = self.lm_head(x[:, [-1], :]) |
| return logits, None |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): |
| for _ in range(max_new_tokens): |
| idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1, :] / temperature |
|
|
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| 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 |
|
|
| @torch.no_grad() |
| def stream(self, idx, max_new_tokens, temperature=1.0, top_k=None): |
| """Yield one token id at a time for real-time streaming.""" |
| for _ in range(max_new_tokens): |
| idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1, :] / temperature |
|
|
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| 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) |
| yield next_token.item() |
|
|
| def num_params(self): |
| return sum(p.numel() for p in self.parameters()) |
|
|
