실습 4 — 미니 GPT 학습하고 글 이어 쓰기

import torch
from torch import nn
from torch.nn import functional as F
import matplotlib.pyplot as plt

plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False

with open("cs_dialogs.txt", encoding="utf-8") as f:
    text = f.read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
V, BLOCK, EMB, HEADS, LAYERS = len(chars), 64, 64, 4, 2

data = torch.tensor(encode(text))
n = int(len(data) * 0.9)
train_data, val_data = data[:n], data[n:]


# 실습 3과 같은 모델 정의
class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(EMB)
        self.attn = nn.MultiheadAttention(EMB, HEADS, batch_first=True)
        self.ln2 = nn.LayerNorm(EMB)
        self.ff = nn.Sequential(nn.Linear(EMB, 4 * EMB), nn.ReLU(), nn.Linear(4 * EMB, EMB))

    def forward(self, x):
        T = x.shape[1]
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.ff(self.ln2(x))


class MiniGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok_emb = nn.Embedding(V, EMB)
        self.pos_emb = nn.Embedding(BLOCK, EMB)
        self.blocks = nn.Sequential(*[Block() for _ in range(LAYERS)])
        self.ln = nn.LayerNorm(EMB)
        self.head = nn.Linear(EMB, V)

    def forward(self, idx, targets=None):
        T = idx.shape[1]
        x = self.tok_emb(idx) + self.pos_emb(torch.arange(T))
        logits = self.head(self.ln(self.blocks(x)))
        loss = None if targets is None else F.cross_entropy(logits.view(-1, V), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, n_new, temperature=1.0, generator=None):
        for _ in range(n_new):
            logits, _ = self(idx[:, -BLOCK:])
            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1, generator=generator)], dim=1)
        return idx

print("준비 완료:", len(train_data), len(val_data))

g = torch.Generator().manual_seed(42)     # 배치를 뽑는 순서를 고정한다

def get_batch(split, batch_size=32):
    """글의 아무 곳에서나 64글자 묶음을 batch_size 개 뽑는다. y 는 x 를 한 칸 민 것(다음 글자)."""
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - BLOCK - 1, (batch_size,), generator=g)   # 시작 위치들
    x = torch.stack([d[i:i + BLOCK] for i in ix])
    y = torch.stack([d[i + 1:i + BLOCK + 1] for i in ix])
    return x, y

xb, yb = get_batch("train", batch_size=1)
print("입력 앞 12글자:", repr(decode(xb[0, :12].tolist())))
print("정답 앞 12글자:", repr(decode(yb[0, :12].tolist())))

torch.manual_seed(42)
model = MiniGPT()
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)   # Adam 에 가중치 감쇠를 더한 옵티마이저

steps, train_hist, val_hist = [], [], []
for step in range(1, 1501):
    xb, yb = get_batch("train")
    _, loss = model(xb, yb)          # 순전파 + 손실
    opt.zero_grad()                  # 지난 기울기 지우기
    loss.backward()                  # 역전파
    opt.step()                       # 가중치 고치기

    if step % 100 == 0:              # 100번마다 검증용 손실도 잰다
        with torch.no_grad():
            xv, yv = get_batch("val")
            _, val_loss = model(xv, yv)
        steps.append(step); train_hist.append(loss.item()); val_hist.append(val_loss.item())
        if step % 300 == 0:
            print(f"{step:5d}번  학습 손실 {loss.item():.3f}  검증 손실 {val_loss.item():.3f}")

# 체크포인트: 학습한 가중치를 파일로 저장해 둔다
torch.save(model.state_dict(), "out_minigpt.pt")
print("저장: out_minigpt.pt")

plt.figure(figsize=(8, 3.5))
plt.plot(steps, train_hist, label="학습 손실")
plt.plot(steps, val_hist, label="검증 손실")
plt.xlabel("학습 횟수")
plt.ylabel("손실(교차 엔트로피)")
plt.title("미니 GPT 의 손실")
plt.legend()
plt.show()

start = torch.tensor([encode("고객: ")])
for temp in [0.5, 1.0, 1.5]:
    out = model.generate(start, 110, temperature=temp, generator=torch.Generator().manual_seed(1))
    print(f"===== 온도 {temp}")
    print(decode(out[0].tolist()))