| """A very small GPT-style language model (character level).""" |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class Block(nn.Module): |
| """One transformer block: self-attention, then a small feed-forward net.""" |
|
|
| def __init__(self, n_embd, n_head, dropout): |
| super().__init__() |
| self.n_head = n_head |
| self.ln1 = nn.LayerNorm(n_embd) |
| self.qkv = nn.Linear(n_embd, 3 * n_embd) |
| self.proj = nn.Linear(n_embd, n_embd) |
| self.ln2 = nn.LayerNorm(n_embd) |
| self.mlp = nn.Sequential( |
| nn.Linear(n_embd, 4 * n_embd), nn.GELU(), nn.Linear(4 * n_embd, n_embd) |
| ) |
| self.dropout = dropout |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| q, k, v = self.qkv(self.ln1(x)).split(C, dim=2) |
| q, k, v = (t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in (q, k, v)) |
| |
| a = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0 |
| ) |
| x = x + self.proj(a.transpose(1, 2).reshape(B, T, C)) |
| return x + self.mlp(self.ln2(x)) |
|
|
|
|
| class TinyGPT(nn.Module): |
| def __init__(self, vocab_size, block_size=128, n_layer=4, n_head=4, n_embd=128, dropout=0.1): |
| super().__init__() |
| self.block_size = block_size |
| self.tok_emb = nn.Embedding(vocab_size, n_embd) |
| self.pos_emb = nn.Embedding(block_size, n_embd) |
| self.blocks = nn.ModuleList(Block(n_embd, n_head, dropout) for _ in range(n_layer)) |
| self.ln_f = nn.LayerNorm(n_embd) |
| self.head = nn.Linear(n_embd, vocab_size) |
|
|
| def forward(self, idx, targets=None): |
| T = idx.shape[1] |
| x = self.tok_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device)) |
| for block in self.blocks: |
| x = block(x) |
| logits = self.head(self.ln_f(x)) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new_tokens, temperature=0.8, top_k=None): |
| """Yield one new token id at a time, feeding each back in as input.""" |
| for _ in range(max_new_tokens): |
| logits, _ = self(idx[:, -self.block_size:]) |
| logits = logits[:, -1, :] / max(temperature, 1e-5) |
| if top_k: |
| kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, [-1]] |
| logits[logits < kth] = -float("inf") |
| nxt = torch.multinomial(F.softmax(logits, dim=-1), 1) |
| idx = torch.cat([idx, nxt], dim=1) |
| yield nxt.item() |
|
|
