실습 3 — 미니 GPT 만들기

import torch
from torch import nn
from torch.nn import functional as F

# 실습 1과 똑같이 글을 읽고 토크나이저를 만든다
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 = len(chars)      # 어휘 크기
BLOCK = 64          # 한 번에 보는 글자 수
EMB = 64            # 임베딩 좌표 수
HEADS = 4           # 헤드 수
LAYERS = 2          # 블록 수
print(V, BLOCK, EMB, HEADS, LAYERS)

class Block(nn.Module):
    """Transformer 블록 하나: LayerNorm → 마스크 셀프 어텐션 → 더하기 → LayerNorm → FFN → 더하기"""

    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(EMB)
        self.attn = nn.MultiheadAttention(EMB, HEADS, batch_first=True)
        self.ln2 = nn.LayerNorm(EMB)
        # 피드포워드: 좌표를 4배로 넓혔다가 다시 줄인다(2장의 신경망과 같은 모양)
        self.ff = nn.Sequential(nn.Linear(EMB, 4 * EMB), nn.ReLU(), nn.Linear(4 * EMB, EMB))

    def forward(self, x):
        T = x.shape[1]                                    # 지금 넣은 글자 수
        # True 인 칸(자기보다 뒤)을 가리는 마스크
        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]   # 잔차 연결 1
        x = x + self.ff(self.ln2(x))                                         # 잔차 연결 2
        return x

print("블록 정의 완료")

class MiniGPT(nn.Module):
    """임베딩 → 블록 LAYERS 개 → 다음 글자 점수"""

    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)             # 좌표 64개 → 어휘 크기만큼의 점수

    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 not None:
            # 분류 손실: 자리마다 "다음 글자"를 맞혔는지. 모양을 (건수, 어휘 크기) 로 펴서 넣는다
            loss = 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):
        """글 끝에 글자를 하나씩 n_new 번 붙여 나간다."""
        for _ in range(n_new):
            logits, _ = self(idx[:, -BLOCK:])                  # 마지막 64글자만 넣는다
            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)   # 마지막 자리의 점수 → 확률
            nxt = torch.multinomial(probs, 1, generator=generator)    # 확률대로 한 글자 뽑기
            idx = torch.cat([idx, nxt], dim=1)                  # 글 끝에 붙인다
        return idx

torch.manual_seed(42)
model = MiniGPT()
print("파라미터 수:", sum(p.numel() for p in model.parameters()))

# "고객: " 으로 시작하게 하고 80글자를 이어 쓰게 한다
start = torch.tensor([encode("고객: ")])
out = model.generate(start, 80, generator=torch.Generator().manual_seed(1))
print(decode(out[0].tolist()))