TinyLM — 세상에서 가장 작은 언어모델

import torch
import torch.nn as nn

# === 세상에서 가장 작은 언어모델 ===

# --- 1. 데이터 준비 (Ch2) ---
text = "나는 고양이를 좋아한다 나는 강아지를 좋아한다"
words = text.split()

# --- 2. 토큰화 (Ch3) ---
vocab = {word: i for i, word in enumerate(sorted(set(words)))}
vocab_size = len(vocab)
print(f"어휘: {vocab}")
print(f"어휘 크기: {vocab_size}")

X_train, Y_train = [], []
for i in range(len(words) - 1):
    X_train.append(vocab[words[i]])
    Y_train.append(vocab[words[i+1]])

X = torch.tensor(X_train)
Y = torch.tensor(Y_train)

# --- 3. 모델 정의 (Ch4 + Ch8) ---
class TinyLM(nn.Module):
    def __init__(self, vocab_size, embed_dim=16):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.output = nn.Linear(embed_dim, vocab_size)

    def forward(self, x):
        embeds = self.embedding(x)
        logits = self.output(embeds)
        return logits

model = TinyLM(vocab_size)
params = sum(p.numel() for p in model.parameters())
print(f"파라미터 수: {params:,}\n")

# --- 4. 학습 (Ch5 + Ch6) ---
criterion = nn.CrossEntropyLoss()              # Ch5
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)  # Ch6

print("=== 학습 시작 ===")
for epoch in range(100):
    logits = model(X)
    loss = criterion(logits, Y)
    optimizer.zero_grad()
    loss.backward()                             # Ch6 역전파
    optimizer.step()
    if epoch % 20 == 0:
        print(f"  Epoch {epoch:3d}: 손실 = {loss.item():.4f}")

# --- 5. 추론 (Ch5 Softmax) ---
print("\n=== 추론 (다음 단어 예측) ===")
idx_to_word = {i: w for w, i in vocab.items()}

for word in ["나는", "고양이를"]:
    if word in vocab:
        input_id = torch.tensor([vocab[word]])
        with torch.no_grad():
            logits = model(input_id)
            probs = torch.softmax(logits, dim=-1)

        print(f"\n'{word}' 다음에 올 단어:")
        for idx in range(vocab_size):
            w = idx_to_word[idx]
            p = probs[0][idx].item()
            bar = "█" * int(p * 30)
            print(f"  {w:8s}: {p:.3f} {bar}")

import math
import random
random.seed(42)

# === TinyLM: 순수 파이썬 미니 언어모델 ===

# 어휘 사전 (토큰화)
vocab = {"나는": 0, "AI를": 1, "공부": 2, "한다": 3}
idx2word = {v: k for k, v in vocab.items()}
V = len(vocab)  # 어휘 크기

# 학습 데이터: "나는" → "AI를", "AI를" → "공부", "공부" → "한다"
data = [(0, 1), (1, 2), (2, 3)]

# 임베딩 + 출력 가중치 (랜덤 초기화)
embed = [[random.gauss(0, 0.5) for _ in range(8)] for _ in range(V)]
W_out = [[random.gauss(0, 0.3) for _ in range(V)] for _ in range(8)]

def softmax(logits):
    mx = max(logits)
    exps = [math.exp(x - mx) for x in logits]
    s = sum(exps)
    return [e / s for e in exps]

def cross_entropy(probs, target):
    return -math.log(probs[target] + 1e-9)

# 학습 루프
lr = 0.5
for epoch in range(30):
    total_loss = 0
    for x_idx, y_idx in data:
        # 순전파: 임베딩 → 선형 → softmax
        h = embed[x_idx]
        logits = [sum(h[k] * W_out[k][j] for k in range(8)) for j in range(V)]
        probs = softmax(logits)
        loss = cross_entropy(probs, y_idx)
        total_loss += loss

        # 역전파 (수동): dL/d_logits
        d_logits = list(probs)
        d_logits[y_idx] -= 1  # softmax + CE 기울기

        # 가중치 업데이트
        for k in range(8):
            for j in range(V):
                W_out[k][j] -= lr * h[k] * d_logits[j]
            embed[x_idx][k] -= lr * sum(W_out[k][j] * d_logits[j] for j in range(V))

    if epoch % 5 == 0:
        print(f"에폭 {epoch:2d} | 손실: {total_loss:.4f}")

# 추론: 다음 단어 예측
print("\n--- 추론 결과 ---")
for word, idx in vocab.items():
    h = embed[idx]
    logits = [sum(h[k] * W_out[k][j] for k in range(8)) for j in range(V)]
    probs = softmax(logits)
    pred = max(range(V), key=lambda j: probs[j])
    print(f"  {word} → {idx2word[pred]} (확률 {probs[pred]:.1%})")