# 迷你語言模型：在本課程網站的文字（corpus.txt，約 5 萬字）上，訓練一個字元級的 Transformer 預測「下一個字」
import json
import math
from pathlib import Path

import torch
from torch import nn

HERE = Path(__file__).parent
text = (HERE / "corpus.txt").read_text(encoding="utf-8")
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
chunks = data.split(500)                              # 每 500 字一段，每 10 段取 1 段當驗證文字（模型沒看過）
train = torch.cat([c for i, c in enumerate(chunks) if i % 10 != 9])
val = torch.cat([c for i, c in enumerate(chunks) if i % 10 == 9])
cut = len(train)
V, T, D = len(chars), 64, 64                          # 字元種類、上下文長度、向量維度

# 基準：unigram（只看字頻）與 bigram（只看前一個字，加一平滑）
uni = torch.bincount(train, minlength=V).float() + 1
bi = torch.ones(V, V)
bi.index_put_((train[:-1], train[1:]), torch.ones(cut - 1), accumulate=True)
ppl_uni = math.exp(-(uni / uni.sum()).log()[val].mean().item())
ppl_bi = math.exp(-(bi / bi.sum(1, keepdim=True)).log()[val[:-1], val[1:]].mean().item())

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(D), nn.LayerNorm(D)
        self.att = nn.MultiheadAttention(D, 4, dropout=0.1, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(D, 4 * D), nn.GELU(), nn.Linear(4 * D, D), nn.Dropout(0.1))
    def forward(self, h, mask):
        a = self.n1(h)
        h = h + self.att(a, a, a, attn_mask=mask, need_weights=False)[0]   # 因果遮罩：看不到後面的字
        return h + self.ff(self.n2(h))

class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, D), nn.Embedding(T, D)      # 位置用可學習的向量
        self.blocks = nn.ModuleList([Block(), Block()])
        self.norm, self.head = nn.LayerNorm(D), nn.Linear(D, V)
    def forward(self, x):
        n = x.shape[1]
        mask = torch.ones(n, n, dtype=torch.bool).triu(1)
        h = self.tok(x) + self.pos(torch.arange(n))
        for b in self.blocks:
            h = b(h, mask)
        return self.head(self.norm(h))

def get_batch(src, n, g):
    i = torch.randint(0, len(src) - T - 1, (n,), generator=g)
    x = torch.stack([src[j:j + T] for j in i])
    return x, torch.stack([src[j + 1:j + T + 1] for j in i])          # 目標：每個位置的「下一個字」

torch.manual_seed(0)
g = torch.Generator().manual_seed(0)
net = GPT()
opt = torch.optim.AdamW(net.parameters(), lr=2e-3)
curve, best = [], (1e9, 0, None)
for step in range(1, 2001):
    net.train()
    x, y = get_batch(train, 32, g)
    loss = nn.functional.cross_entropy(net(x).reshape(-1, V), y.reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 100 == 0:
        net.eval()                                                    # 評估時關掉 dropout
        with torch.no_grad():
            xv, yv = get_batch(val, 256, torch.Generator().manual_seed(1))
            vl = nn.functional.cross_entropy(net(xv).reshape(-1, V), yv.reshape(-1)).item()
        curve.append([step, loss.item(), vl])
        if vl < best[0]:                                              # 提早停止：記下驗證誤差最低的模型
            best = (vl, step, {k: v.clone() for k, v in net.state_dict().items()})
net.load_state_dict(best[2]); net.eval()
print(f"字元 {len(text)} 個、{V} 種；模型參數 {sum(p.numel() for p in net.parameters())} 個")
print(f"驗證誤差最低在第 {best[1]} 步（之後開始過擬合），採用這時的模型")
print(f"驗證文字的困惑度（越小越好）：unigram {ppl_uni:.0f}、bigram {ppl_bi:.0f}、Transformer {math.exp(best[0]):.0f}")

def generate(prompt, n, temp, seed):
    g = torch.Generator().manual_seed(seed)
    x = torch.tensor([[stoi[c] for c in prompt]])
    with torch.no_grad():
        for _ in range(n):
            p = torch.softmax(net(x[:, -T:])[0, -1] / temp, -1)       # 溫度：越低越保守
            x = torch.cat([x, torch.multinomial(p, 1, generator=g).view(1, 1)], 1)
    return "".join(chars[i] for i in x[0].tolist())

for temp in (0.5, 1.0):
    print(f"\n溫度 {temp}：", generate("類神經網路", 60, temp, 0).replace("\n", " "))
nxt = {}                                                              # 給網頁動畫：幾個開頭之後，機率最高的 12 個字與其 logit
for prompt in ("類神經網", "倒單擺的", "梯度"):
    with torch.no_grad():
        top = net(torch.tensor([[stoi[c] for c in prompt]]))[0, -1].topk(12)
    nxt[prompt] = [[chars[i], round(v, 4)] for v, i in zip(top.values.tolist(), top.indices.tolist())]
(HERE / "char_lm.json").write_text(json.dumps({"curve": curve, "ppl": [ppl_uni, ppl_bi], "best": best[1], "next": nxt}, ensure_ascii=False), encoding="utf-8")