"""Lesson 10: LoRA - freeze the trained model and fine-tune only small low-rank adapters on new text."""
import torch, torch.nn as nn
from tinygpt import TinyGPT, CharTokenizer, get_batch

class LoRALinear(nn.Module):
    def __init__(self, base, r=4, alpha=8):
        super().__init__()
        self.base = base
        for p in self.base.parameters(): p.requires_grad_(False)
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))       # B = 0, so training starts from the original model
        self.scale = alpha / r
    def forward(self, x):
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale

ck = torch.load("tinygpt.pt"); tok = CharTokenizer("".join(ck["chars"]))
model = TinyGPT(len(ck["chars"]), **ck["cfg"]); model.load_state_dict(ck["model"])
for p in model.parameters(): p.requires_grad_(False)
for blk in model.blocks:                                                # adapt the attention layers only
    blk.attn.qkv = LoRALinear(blk.attn.qkv); blk.attn.proj = LoRALinear(blk.attn.proj)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad); total = sum(p.numel() for p in model.parameters())
print(f"trainable {trainable:,} of {total:,} parameters ({100 * trainable / total:.1f}%)")
new_text = "Power BI measures use DAX. CALCULATE changes the filter context. " * 200
data = torch.tensor(tok.encode(new_text))
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-3)
torch.manual_seed(0)
for step in range(151):
    model.train(); _, loss = model(*get_batch(data, 64, 16)); opt.zero_grad(); loss.backward(); opt.step()
    if step % 50 == 0: print(f"step {step}  loss {loss.item():.3f}")