[Build GPT-2] Model

Model

이전 장에서 Dataset을 어떻게 처리하는지 살펴보았다.
이제 모델을 보자.
이 예시에서는 가장 간단한 예시인 bigram language model을 사용한다.
bigram language model이 뭐냐 ?
Bigram 언어 모델은 Markov 가정에 기반을 두고 있으며, 이는 특정 단어의 출현 확률이 그 이전 단어에만 의존한다고 가정한다. 이러한 모델은 다음과 같은 형식으로 표현할 수 있다:

$P(w_n|w_{n-1})$ 여기서 character $w_{n-1}$가 주어졌을 때 다음 $w_n$ 이 나올 확률을 의미한다.

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)

class BigramLangaugeModel(nn.Module):

	def __init__(self, vocab_size):
    	super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
        
    def forward(self, idx, targets):
    
    	# idx and targets are both (B,T) tensor of integers
    	logits = self.token_embeding_table(idx) # (B,T,C)
        
        return logits

m = BigramLangaugeModel(vocab_size)
out = m(xb, yb)
print(out.shape)

torch.Size([4, 8, 65])

먼저, pytorch를 임포트 해준다.

여기서 nn.Embedding(num_embeddings, embedding_dim)은

• num_embeddings(int) - size of the dictionary of embeddings
• embedding_dim(int) - the size of each embedding vector

여기선 간단하게 바로 65개의 vocab size로 예측할 수 있도록 만들었다.
따라서, idx # shape (4,8) batch 4개의 8개의 sequence(block size)를 가지는 데이터, 를 입력으로 nn.Embeddings(vocab_size(==65), vocab_size) 에 넣는 경우 (4,8,65) $ B, T, C 에 해당하는 예측 logit이 등장하고, 이때 Cross entropy loss를 target(다음 t에 해당되는 characters)를 넣는 것으로 학습을 진행.

Next token prediction

즉, 위 학습 방법은 이전 token을 보고 다음 token을 예측하는 모델을 학습하는 과정을 의미한다.
이는 다음 토큰을 예측한다고 해서 Next token prediction이라고 불린다.

이 때 근데 저 형태 그대로 그냥 넣으면 안되고, (B*T, C) 로 넣어야 하는데 이유는 pytorch Cross Entropy loss 가 작동하는 것 때문에 그럼.

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)

class BigramLanguageModel(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):

        # idx and targets are both (B,T) tensor of integers
        logits = self.token_embedding_table(idx) # (B,T,C)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):
            # get the predictions
            logits, loss = self(idx)
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx

m = BigramLanguageModel(vocab_size)
logits, loss = m(xb, yb)
print(logits.shape)
print(loss)

print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))

torch.Size([32, 65])
tensor(4.8786, grad_fn=<NllLossBackward0>)

SKIcLT;AcELMoTbvZv C?nq-QE33:CJqkOKH-q;:la!oiywkHjgChzbQ?u!3bLIgwevmyFJGUGp
wnYWmnxKWWev-tDqXErVKLgJ

batch_size = 32
for steps in range(100): # increase number of steps for good results...

    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = m(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

print(loss.item())

4.65630578994751

print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))

oTo.JUZ!!zqe!
xBP qbs$Gy'AcOmrLwwt
p$x;Seh-onQbfM?OjKbn'NwUAW -Np3fkz$FVwAUEa-wzWC -wQo-R!v -Mj?,SPiTyZ;o-opr$mOiPJEYD-CfigkzD3p3?zvS;ADz;.y?o,ivCuC'zqHxcVT cHA
rT'Fd,SBMZyOslg!NXeF$sBe,juUzLq?w-wzP-h
ERjjxlgJzPbHxf$ q,q,KCDCU fqBOQT
SV&CW:xSVwZv'DG'NSPypDhKStKzC -$hslxIVzoivnp ,ethA:NCCGoi
tN!ljjP3fwJMwNelgUzzPGJlgihJ!d?q.d
pSPYgCuCJrIFtb
jQXg
pA.P LP,SPJi
DBcuBM:CixjJ$Jzkq,OLf3KLQLMGph$O 3DfiPHnXKuHMlyjxEiyZib3FaHV-oJa!zoc'XSP :CKGUhd?lgCOF$;;DTHZMlvvcmZAm;:iv'MMgO&Ywbc;BLCUd&vZINLIzkuTGZa
D.?

처음에는 완전 막무가내로 나옴 (훈련 안됨)

generate (inference)시에는 그냥 다음 토큰을 내뱉는게 아니라 torch.multinomal에서 샘플링을 진행함 (몰랐네.)