Create train_once.py

app/train_once.py

@@ -0,0 +1,41 @@
+import os, random, torch, torch.nn.functional as F
+from slm_qa import TinyTransformer, encode, wrap_bos_eos, itos, PAD, DATA_QA
+
+def make_sequences():
+    return [wrap_bos_eos(encode("q: "+q) + encode("a: "+a)) for q,a in DATA_QA]
+
+def pad_batches(seqs, batch_size=4, device=torch.device("cpu")):
+    random.shuffle(seqs)
+    batches=[]
+    for i in range(0,len(seqs),batch_size):
+        chunk=seqs[i:i+batch_size]; T=max(len(s) for s in chunk)
+        x=torch.full((len(chunk),T-1),PAD); y=torch.full((len(chunk),T-1),PAD)
+        for bi,s in enumerate(chunk):
+            s_pad=s+[PAD]*(T-len(s))
+            x[bi]=torch.tensor(s_pad[:-1]); y[bi]=torch.tensor(s_pad[1:])
+        batches.append((x.to(device),y.to(device)))
+    return batches
+
+def main():
+    device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model=TinyTransformer(vocab_size=len(itos),max_len=128).to(device)
+    opt=torch.optim.AdamW(model.parameters(),lr=3e-4)
+    batches=pad_batches(make_sequences(),device=device)
+    best=1e9
+    for ep in range(60):  # more epochs
+        loss_sum=0
+        for x,y in batches:
+            logits=model(x)
+            B,T,V=logits.size()
+            loss=F.cross_entropy(logits.view(B*T,V),y.view(B*T),ignore_index=PAD)
+            opt.zero_grad(); loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)
+            opt.step(); loss_sum+=loss.item()
+        avg=loss_sum/len(batches)
+        print(f"Epoch {ep+1:02d} loss {avg:.4f}")
+        if avg<best:
+            best=avg
+            os.makedirs("app/models",exist_ok=True)
+            torch.save(model.state_dict(),"app/models/slm_qa_best.pt")
+
+if __name__=="__main__": main()