| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| VOCAB = 16000 |
| CTX = 512 |
| N_LAYER = 16 |
| N_HEAD = 10 |
| N_EMBD = 640 |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self): |
| super().__init__() |
| 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.fc = nn.Linear(N_EMBD, 4 * N_EMBD) |
| self.fc2 = nn.Linear(4 * N_EMBD, N_EMBD) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| q, k, v = self.qkv(self.ln1(x)).split(C, dim=2) |
| q = q.view(B, T, N_HEAD, C // N_HEAD).transpose(1, 2) |
| k = k.view(B, T, N_HEAD, C // N_HEAD).transpose(1, 2) |
| v = v.view(B, T, N_HEAD, C // N_HEAD).transpose(1, 2) |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| x = x + self.proj(y) |
| x = x + self.fc2(F.gelu(self.fc(self.ln2(x)))) |
| return x |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.tok_emb = nn.Embedding(VOCAB, N_EMBD) |
| self.pos_emb = nn.Embedding(CTX, N_EMBD) |
| self.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)]) |
| self.ln_f = nn.LayerNorm(N_EMBD) |
| self.head = nn.Linear(N_EMBD, VOCAB, bias=False) |
| self.head.weight = self.tok_emb.weight |
| self.apply(self._init) |
|
|
| def _init(self, m): |
| if isinstance(m, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(m.weight, std=0.02) |
| if isinstance(m, nn.Linear) and m.bias is not None: |
| nn.init.zeros_(m.bias) |
|
|
| def forward(self, idx, targets=None): |
| B, T = idx.shape |
| pos = torch.arange(T, device=idx.device) |
| x = self.tok_emb(idx) + self.pos_emb(pos) |
| for b in self.blocks: |
| x = b(x) |
| logits = self.head(self.ln_f(x)) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.float().view(-1, VOCAB), targets.view(-1)) |
| return logits, loss |
|
|
|
|
| @torch.no_grad() |
| def sample(model, tok, prompt, device, max_new=120, temp=0.8, top_k=40): |
| was_training = model.training |
| model.eval() |
| idx = torch.tensor([tok.encode(prompt).ids], device=device) |
| for _ in range(max_new): |
| with torch.autocast("cuda", dtype=torch.float16): |
| logits, _ = model(idx[:, -CTX:]) |
| logits = logits[:, -1, :].float() / temp |
| v, _ = torch.topk(logits, top_k) |
| logits[logits < v[:, [-1]]] = -float("inf") |
| nxt = torch.multinomial(F.softmax(logits, dim=-1), 1) |
| idx = torch.cat([idx, nxt], dim=1) |
| model.train(was_training) |
| return tok.decode(idx[0].tolist()) |
|
|
