Create mini-transformers-ANN.py

mini-transformers-ANN.py

@@ -0,0 +1,242 @@
+# transformer_mini_addition.py
+# Train a tiny Transformer encoder-decoder on a toy task: "a + b =" -> "c"
+# Example: input tokens ["3","+","5","="] -> output tokens ["8"]
+# Run: python transformer_mini_addition.py
+
+import math
+import random
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+
+# -----------------------
+# Config
+# -----------------------
+VOCAB = ["0","1","2","3","4","5","6","7","8","9","+","=","<pad>","<s>","</s>"]
+PAD, BOS, EOS = VOCAB.index("<pad>"), VOCAB.index("<s>"), VOCAB.index("</s>")
+TOK2ID = {t:i for i,t in enumerate(VOCAB)}
+ID2TOK = {i:t for t,i in TOK2ID.items()}
+
+MAX_IN_LEN = 4         # "d + d ="  -> 4 tokens
+MAX_OUT_LEN = 3        # could be 1 or 2 digits + EOS
+EMB_DIM = 128
+FF_DIM = 256
+N_HEAD = 4
+N_LAYERS = 2
+DROPOUT = 0.1
+BATCH_SIZE = 128
+STEPS = 1500           # keep small for a mini training run
+LR = 3e-4
+WARMUP = 100
+SAVE_PATH = "mini_transformer_addition.pt"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.manual_seed(42)
+random.seed(42)
+
+# -----------------------
+# Data: generate on-the-fly mini dataset
+# -----------------------
+def encode(seq: List[str], max_len: int) -> List[int]:
+    ids = [TOK2ID[s] for s in seq]
+    if len(ids) < max_len:
+        ids += [PAD]*(max_len-len(ids))
+    return ids[:max_len]
+
+def sample_pair() -> Tuple[List[int], List[int]]:
+    a, b = random.randint(0,9), random.randint(0,9)
+    c = a + b
+    inp = [str(a), "+", str(b), "="]                         # length 4
+    out_tokens = list(str(c))                                # "0".."18"
+    tgt = [BOS] + [TOK2ID[t] for t in out_tokens] + [EOS]    # BOS ... EOS
+    # pad to MAX_OUT_LEN + 2 (BOS/EOS)
+    max_len = MAX_OUT_LEN + 2
+    if len(tgt) < max_len:
+        tgt += [PAD] * (max_len - len(tgt))
+    return encode(inp, MAX_IN_LEN), tgt
+
+class MiniAddDataset(Dataset):
+    def __init__(self, size=5000):
+        self.size = size
+    def __len__(self): return self.size
+    def __getitem__(self, idx):
+        src, tgt = sample_pair()
+        return torch.tensor(src), torch.tensor(tgt)
+
+train_ds = MiniAddDataset(size=8000)
+val_ds   = MiniAddDataset(size=500)
+train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
+val_dl   = DataLoader(val_ds, batch_size=BATCH_SIZE)
+
+# -----------------------
+# Model
+# -----------------------
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=64):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0)/d_model))
+        pe[:, 0::2] = torch.sin(pos * div)
+        pe[:, 1::2] = torch.cos(pos * div)
+        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1), :]
+
+class TinyTransformer(nn.Module):
+    def __init__(self, vocab_size: int):
+        super().__init__()
+        self.src_emb = nn.Embedding(vocab_size, EMB_DIM, padding_idx=PAD)
+        self.tgt_emb = nn.Embedding(vocab_size, EMB_DIM, padding_idx=PAD)
+        self.pos_enc_src = PositionalEncoding(EMB_DIM, max_len=MAX_IN_LEN+8)
+        self.pos_enc_tgt = PositionalEncoding(EMB_DIM, max_len=MAX_OUT_LEN+8)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=EMB_DIM, nhead=N_HEAD, dim_feedforward=FF_DIM, dropout=DROPOUT, batch_first=True
+        )
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=EMB_DIM, nhead=N_HEAD, dim_feedforward=FF_DIM, dropout=DROPOUT, batch_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=N_LAYERS)
+        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=N_LAYERS)
+        self.lm_head = nn.Linear(EMB_DIM, vocab_size)
+
+    def make_padding_mask(self, seq, pad_idx=PAD):
+        # returns (batch, 1, 1, seq_len) for nn.Transformer; but with batch_first=True we can use (batch, seq_len)
+        return (seq == pad_idx)
+
+    def generate_square_subsequent_mask(self, sz):
+        # causal mask for decoder (tgt): allow attention to previous positions only
+        return torch.triu(torch.ones(sz, sz, device=DEVICE), diagonal=1).bool()
+
+    def forward(self, src_ids, tgt_ids):
+        # src_ids: (B, S) ; tgt_ids: (B, T)
+        src_key_padding_mask = self.make_padding_mask(src_ids)    # (B,S)
+        tgt_key_padding_mask = self.make_padding_mask(tgt_ids)    # (B,T)
+        tgt_mask = self.generate_square_subsequent_mask(tgt_ids.size(1))
+
+        src = self.src_emb(src_ids)
+        src = self.pos_enc_src(src)
+        memory = self.encoder(src, src_key_padding_mask=src_key_padding_mask)
+
+        tgt = self.tgt_emb(tgt_ids)
+        tgt = self.pos_enc_tgt(tgt)
+        out = self.decoder(
+            tgt, memory,
+            tgt_mask=tgt_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=src_key_padding_mask
+        )
+        logits = self.lm_head(out)  # (B,T,V)
+        return logits
+
+# -----------------------
+# Training utils
+# -----------------------
+class WarmupAdam(torch.optim.Adam):
+    def __init__(self, params, lr, warmup_steps=1000):
+        super().__init__(params, lr=lr, betas=(0.9, 0.98), eps=1e-9)
+        self.warmup_steps = warmup_steps
+        self._step = 0
+        self._base_lr = lr
+    def step(self, closure=None):
+        self._step += 1
+        scale = min(self._step ** (-0.5), self._step * (self.warmup_steps ** (-1.5)))
+        for g in self.param_groups:
+            g['lr'] = self._base_lr * scale * (self.warmup_steps ** 0.5)
+        return super().step(closure=closure)
+
+def shift_right(tgt):
+    """
+    Teacher forcing: model sees BOS + y[:-1] and predicts y.
+    Here tgt is already [BOS, y..., EOS, PAD...]
+    We return inp=tgt[:, :-1], label=tgt[:, 1:]
+    """
+    return tgt[:, :-1], tgt[:, 1:]
+
+def accuracy_from_logits(logits, labels, pad=PAD):
+    # logits: (B,T,V), labels: (B,T)
+    preds = logits.argmax(-1)
+    mask = labels.ne(pad)
+    correct = (preds.eq(labels) & mask).sum().item()
+    total = mask.sum().item() + 1e-9
+    return correct/total
+
+# -----------------------
+# Train
+# -----------------------
+model = TinyTransformer(vocab_size=len(VOCAB)).to(DEVICE)
+criterion = nn.CrossEntropyLoss(ignore_index=PAD)
+optim = WarmupAdam(model.parameters(), lr=LR, warmup_steps=WARMUP)
+
+def run_epoch(dl, train=True):
+    model.train(train)
+    total_loss, total_acc, n = 0.0, 0.0, 0
+    for src, tgt in dl:
+        src, tgt = src.to(DEVICE), tgt.to(DEVICE)
+        dec_inp, labels = shift_right(tgt)
+        logits = model(src, dec_inp)
+        loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))
+        acc = accuracy_from_logits(logits, labels)
+
+        if train:
+            optim.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optim.step()
+
+        bs = src.size(0)
+        total_loss += loss.item() * bs
+        total_acc += acc * bs
+        n += bs
+    return total_loss/n, total_acc/n
+
+best_val = 0.0
+for step in range(1, STEPS+1):
+    tr_loss, tr_acc = run_epoch(train_dl, train=True)
+    if step % 50 == 0:
+        val_loss, val_acc = run_epoch(val_dl, train=False)
+        print(f"[step {step:4d}] train loss {tr_loss:.3f} acc {tr_acc:.3f} | val loss {val_loss:.3f} acc {val_acc:.3f}")
+        if val_acc > best_val:
+            best_val = val_acc
+            torch.save({"model": model.state_dict()}, SAVE_PATH)
+
+print(f"Saved best model to: {SAVE_PATH}")
+
+# -----------------------
+# Inference demo
+# -----------------------
+def encode_inp(a:int,b:int):
+    seq = [str(a), "+", str(b), "="]
+    return torch.tensor([encode(seq, MAX_IN_LEN)], device=DEVICE)
+
+def greedy_decode(src_ids, max_len=MAX_OUT_LEN+2):
+    model.eval()
+    with torch.no_grad():
+        # Start with BOS
+        ys = torch.tensor([[BOS]], device=DEVICE)
+        for _ in range(max_len-1):
+            logits = model(src_ids, ys)
+            next_tok = logits[:, -1, :].argmax(-1, keepdim=True)  # (B,1)
+            ys = torch.cat([ys, next_tok], dim=1)
+            if next_tok.item() == EOS:
+                break
+    return ys.squeeze(0).tolist()
+
+def detok(ids: List[int]) -> str:
+    toks = [ID2TOK[i] for i in ids if i not in (PAD, BOS)]
+    out = []
+    for t in toks:
+        if t == "</s>": break
+        out.append(t)
+    return "".join(out)
+
+# Load best (optional, already in memory)
+ckpt = torch.load(SAVE_PATH, map_location=DEVICE)
+model.load_state_dict(ckpt["model"])
+
+for (a,b) in [(3,5),(9,8),(0,0),(7,2),(4,6)]:
+    src = encode_inp(a,b)
+    out_ids = greedy_decode(src)
+    print(f"{a}+{b} => {detok(out_ids)}")