| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import tiktoken |
| import math |
|
|
| print(f"TORCH OK: {torch.__version__}") |
|
|
| enc = tiktoken.get_encoding("gpt2") |
| vocab_size = enc.n_vocab |
| print(f"Tokenizer: GPT2 {vocab_size} tokens") |
|
|
| class Config: |
| vocab_size = vocab_size |
| block_size = 128 |
| n_layer = 4 |
| n_head = 4 |
| n_embd = 256 |
| dropout = 0.1 |
|
|
| class CausalSelfAttention(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.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) |
| self.c_proj = nn.Linear(config.n_embd, config.n_embd) |
| self.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.size() |
| qkv = self.c_attn(x) |
| q, k, v = qkv.split(self.n_embd, dim=2) |
| q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
| k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
| v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
| att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) |
| att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf')) |
| att = F.softmax(att, dim=-1) |
| att = F.dropout(att, p=self.dropout, training=self.training) |
| y = att @ v |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| y = self.c_proj(y) |
| return y |
|
|
| class Block(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(config.n_embd) |
| self.attn = CausalSelfAttention(config) |
| self.ln2 = nn.LayerNorm(config.n_embd) |
| self.mlp = nn.Sequential( |
| nn.Linear(config.n_embd, 4 * config.n_embd), |
| nn.GELU(), |
| nn.Linear(4 * config.n_embd, config.n_embd), |
| nn.Dropout(config.dropout) |
| ) |
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
| class TinyLLM(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd) |
| self.pos_emb = nn.Embedding(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) |
|
|
| def forward(self, idx, targets=None): |
| B, T = idx.size() |
| tok = self.tok_emb(idx) |
| pos = self.pos_emb(torch.arange(T, device=idx.device)) |
| x = tok + pos |
| for block in self.blocks: |
| x = block(x) |
| x = self.ln_f(x) |
| logits = self.head(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=50, top_k=50): |
| for _ in range(max_new_tokens): |
| idx_cond = idx[:, -self.config.block_size:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1, :] |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float('Inf') |
| probs = F.softmax(logits, dim=-1) |
| idx_next = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat((idx, idx_next), dim=1) |
| return idx |
|
|
| config = Config() |
| model = TinyLLM(config) |
| print(f"Model params: {sum(p.numel() for p in model.parameters())/1e6:.2f}M") |
|
|
| x = torch.randint(0, vocab_size, (2, config.block_size)) |
| y = torch.randint(0, vocab_size, (2, config.block_size)) |
|
|
| logits, loss = model(x, y) |
| print(f"Initial loss: {loss.item():.4f}") |
|
|
| opt = torch.optim.AdamW(model.parameters(), lr=3e-4) |
| for i in range(20): |
| logits, loss = model(x, y) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| if i % 5 == 0: |
| print(f"step {i} loss {loss.item():.4f}") |
|
|
| prompt = "Hello, I am a" |
| ids = torch.tensor([enc.encode(prompt)], dtype=torch.long) |
| out = model.generate(ids, max_new_tokens=30) |
| print("\nPROMPT:", prompt) |
| print("GEN:", enc.decode(out[0].tolist())) |
|
|
