%pip install -q torch requests matplotlib

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

# 재현성을 위한 시드 고정
torch.manual_seed(1337)

# GPU가 있으면 GPU를, 없으면 CPU를 사용
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"사용 장치:{device}")
if device == 'cuda':
    print(f"GPU 이름:{torch.cuda.get_device_name(0)}")

# matplotlib 한글 폰트 설정 (Windows 기본 폰트인 맑은 고딕 사용)
# 폰트를 지정하지 않으면 그래프의 한글 제목/라벨이 네모(□)로 깨져서 표시됨
mpl.rcParams['font.family'] = 'Malgun Gothic'
mpl.rcParams['axes.unicode_minus'] = False  # 맑은 고딕에서 마이너스 기호가 깨지는 문제 방지

import os
import requests

data_path = './data/gpt/tinyshakespeare.txt'

if not os.path.exists(data_path):
    url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
    response = requests.get(url)
    response.raise_for_status()
    with open(data_path, 'w', encoding='utf-8') as f:
        f.write(response.text)
    print("데이터셋을 다운로드했습니다.")
else:
    print("이미 다운로드된 데이터셋을 사용합니다.")

with open(data_path, 'r', encoding='utf-8') as f:
    text = f.read()

print(f"전체 글자 수:{len(text):,}")
print("--- 앞부분 미리보기 ---")
print(text[:300])

# 텍스트에 등장하는 모든 고유 문자를 정렬된 리스트로 추출
chars = sorted(list(set(text)))
vocab_size = len(chars)

print(f"고유 문자 수 (vocab_size):{vocab_size}")
print(''.join(chars))

# 문자 -> 정수, 정수 -> 문자 매핑 테이블
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

def encode(s: str) -> list[int]:
    """문자열을 정수 리스트로 변환"""
    return [stoi[c] for c in s]

def decode(ids: list[int]) -> str:
    """정수 리스트를 문자열로 변환"""
    return ''.join(itos[i] for i in ids)

# 왕복 변환이 잘 되는지 확인
sample = "Hello, world!"
encoded = encode(sample)
print(encoded)
print(decode(encoded))

data = torch.tensor(encode(text), dtype=torch.long)
print(f"인코딩된 전체 데이터 shape:{data.shape}, dtype:{data.dtype}")

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

print(f"train:{len(train_data):,}자, val:{len(val_data):,}자")

block_size_demo = 8
x_demo = train_data[:block_size_demo]
y_demo = train_data[1:block_size_demo+1]  # y는 x보다 한 칸씩 뒤로 밀린 정답

for t in range(block_size_demo):
    context = x_demo[:t+1]
    target = y_demo[t]
    print(f"문맥{context.tolist()} -> 정답{target.item()} ({decode([target.item()])!r})")

block_size = 256  # 모델이 한 번에 보는 문맥 길이
batch_size = 64    # 한 스텝에 동시에 처리하는 시퀀스 개수

def get_batch(split: str):
    d = train_data if split == 'train' else val_data
    # 배치 크기만큼 랜덤한 시작 위치를 뽑는다
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i:i+block_size] for i in ix])
    y = torch.stack([d[i+1:i+block_size+1] for i in ix])
    return x.to(device), y.to(device)

xb, yb = get_batch('train')
print(f"입력 배치 shape:{xb.shape}")   # (batch_size, block_size)
print(f"정답 배치 shape:{yb.shape}")

eval_iters = 200

@torch.no_grad()
def estimate_loss(model):
    """train/val 각각에서 여러 배치를 뽑아 평균 loss를 계산 (노이즈 감소용)"""
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out

class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # (vocab_size, vocab_size) 테이블: 각 문자가 다음 문자에 대한 로짓을 직접 담고 있음
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        logits = self.token_embedding_table(idx)  # (B, T, vocab_size)

        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):
        for _ in range(max_new_tokens):
            logits, _ = self(idx)
            logits = logits[:, -1, :]  # 마지막 타임스텝의 로짓만 사용
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)  # 확률적으로 다음 토큰 샘플링
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

bigram_model = BigramLanguageModel(vocab_size).to(device)

# 학습 전 loss는 이론상 -ln(1/vocab_size)에 가까워야 함 (완전 랜덤 예측)
import math
print(f"이론적 초기 loss (완전 랜덤):{math.log(vocab_size):.4f}")
losses = estimate_loss(bigram_model)
print(f"실제 초기 loss: train{losses['train']:.4f}, val{losses['val']:.4f}")

# Bigram 모델을 짧게 학습
optimizer = torch.optim.AdamW(bigram_model.parameters(), lr=1e-3)

for step in range(3000):
    xb, yb = get_batch('train')
    logits, loss = bigram_model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

losses = estimate_loss(bigram_model)
print(f"학습 후 loss: train{losses['train']:.4f}, val{losses['val']:.4f}")

# Bigram 모델로 텍스트 생성
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = bigram_model.generate(context, max_new_tokens=300)[0].tolist()
print(decode(generated))

# 모델 하이퍼파라미터 (nanoGPT의 소형 설정을 따름 — GPU에서 수 분 내 학습 가능)
n_embd = 384    # 임베딩 차원
n_head = 6      # attention head 개수
n_layer = 6     # Transformer 블록(레이어) 개수
dropout = 0.2   # 과적합 방지용 dropout 비율

print(f"head 하나의 차원 (n_embd // n_head):{n_embd // n_head}")

class Head(nn.Module):
    """Self-Attention head 하나"""

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # 학습되지 않는 버퍼로 등록 (미래 위치를 가리는 하삼각 마스크)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)    # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)

        # Scaled Dot-Product Attention: q와 k의 내적을 head_size의 제곱근으로 나눠 스케일 조정
        wei = q @ k.transpose(-2, -1) * (C ** -0.5)  # (B, T, T)

        # Causal mask: 미래 위치(대각선 위쪽)는 -inf로 막아 softmax 이후 0이 되게 함
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)

        v = self.value(x)
        out = wei @ v  # (B, T, head_size)
        return out

class MultiHeadAttention(nn.Module):
    """여러 head를 병렬로 실행하고 결과를 이어붙인 뒤 다시 선형 변환"""

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)  # (B, T, n_embd)
        out = self.dropout(self.proj(out))
        return out

class FeedForward(nn.Module):
    """토큰별로 독립적으로 적용되는 간단한 MLP"""

    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

class Block(nn.Module):
    """Transformer 디코더 블록: 통신(Self-Attention) + 계산(FeedForward)"""

    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        # Pre-LN + Residual: 입력을 LayerNorm 통과시킨 뒤 서브레이어에 넣고, 원래 입력에 더해준다
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.token_embedding_table(idx)                          # (B, T, n_embd)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T, n_embd)
        x = tok_emb + pos_emb   # 브로드캐스팅으로 (B, T, n_embd)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        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

    @torch.no_grad()
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            # position embedding 테이블 크기(block_size)를 넘지 않도록 최근 문맥만 사용
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

model = GPTLanguageModel().to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"모델 파라미터 수:{n_params/1e6:.2f}M")

learning_rate = 3e-4
max_iters = 5000
eval_interval = 500

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

train_losses, val_losses, logged_steps = [], [], []

for step in range(max_iters):
    if step % eval_interval == 0 or step == max_iters - 1:
        losses = estimate_loss(model)
        print(f"step{step:5d}: train loss{losses['train']:.4f}, val loss{losses['val']:.4f}")
        train_losses.append(losses['train'])
        val_losses.append(losses['val'])
        logged_steps.append(step)

    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

plt.figure(figsize=(8, 4))
plt.plot(logged_steps, train_losses, label='train loss')
plt.plot(logged_steps, val_losses, label='val loss')
plt.xlabel('step')
plt.ylabel('loss')
plt.title('학습 손실 곡선')
plt.legend()
plt.grid(alpha=0.3)
plt.show()

context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=500)[0].tolist()
print(decode(generated))

prompt = "ROMEO:"
context = torch.tensor([encode(prompt)], dtype=torch.long, device=device)
generated = model.generate(context, max_new_tokens=300)[0].tolist()
print(decode(generated))

# 연습 1: 하이퍼파라미터별 파라미터 수 / val loss 비교
def build_and_train(n_layer_, n_head_, n_embd_, iters=1500):
    """주어진 하이퍼파라미터로 새 모델을 만들고 짧게 학습시킨 뒤 (파라미터 수, val loss)를 반환"""
    global n_layer, n_head, n_embd  # Head/Block/GPTLanguageModel의 __init__이 참조하는 전역 변수를 임시로 덮어씀
    n_layer, n_head, n_embd = n_layer_, n_head_, n_embd_

    m = GPTLanguageModel().to(device)
    n_params = sum(p.numel() for p in m.parameters())

    opt = torch.optim.AdamW(m.parameters(), lr=learning_rate)
    for step in range(iters):
        xb, yb = get_batch('train')
        _, loss = m(xb, yb)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()

    val_loss = estimate_loss(m)['val']
    return n_params, val_loss

# (n_layer, n_head, n_embd) 조합 3가지를 비교 — head_size(=n_embd//n_head)가 정수가 되도록 선택
configs = [
    (2, 2, 128),   # 작은 모델
    (6, 6, 384),   # 노트북 본문과 동일한 기본 설정
    (8, 8, 512),   # 더 큰 모델
]

for nl, nh, ne in configs:
    n_params, val_loss = build_and_train(nl, nh, ne)
    print(f"n_layer={nl}, n_head={nh}, n_embd={ne} -> 파라미터{n_params/1e6:.2f}M, val loss{val_loss:.4f}")

# 전역 하이퍼파라미터를 원래 값으로 복원 (이후 셀에서 model 변수를 다시 쓸 때 혼동을 피하기 위함)
n_layer, n_head, n_embd = 6, 6, 384

# 연습 2: temperature를 조절할 수 있는 generate 함수
@torch.no_grad()
def generate_with_temperature(m, idx, max_new_tokens, temperature=1.0):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -block_size:]
        logits, _ = m(idx_cond)
        logits = logits[:, -1, :] / temperature  # temperature로 나눠 분포의 뾰족함(sharpness)을 조절
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

context = torch.zeros((1, 1), dtype=torch.long, device=device)

for temp in [0.3, 1.0, 1.5]:
    print(f"--- temperature={temp} ---")
    out = generate_with_temperature(model, context, max_new_tokens=200, temperature=temp)
    print(decode(out[0].tolist()))
    print()

# 연습 3: dropout 값에 따른 train/val loss 격차 비교
def train_with_dropout(dropout_value, iters=1500):
    global dropout  # Head/MultiHeadAttention/FeedForward가 참조하는 전역 dropout 값을 변경
    dropout = dropout_value

    m = GPTLanguageModel().to(device)
    opt = torch.optim.AdamW(m.parameters(), lr=learning_rate)
    for step in range(iters):
        xb, yb = get_batch('train')
        _, loss = m(xb, yb)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()

    losses = estimate_loss(m)
    return losses['train'], losses['val']

for d in [0.2, 0.0]:
    train_loss, val_loss = train_with_dropout(d)
    gap = val_loss - train_loss
    print(f"dropout={d}: train{train_loss:.4f}, val{val_loss:.4f}, 격차(val-train){gap:.4f}")

dropout = 0.2  # 원래 값으로 복원

%pip install -q tiktoken

# 연습 4: 문자 단위 토크나이저 vs BPE 토크나이저 비교
import tiktoken

enc = tiktoken.get_encoding("gpt2")  # 실제 GPT-2가 사용하는 BPE 토크나이저

sample_text = text[:1000]  # 앞부분 1000자로 비교

char_tokens = encode(sample_text)     # 우리가 만든 문자 단위 토크나이저
bpe_tokens = enc.encode(sample_text)  # GPT-2의 BPE 토크나이저

print(f"원문 글자 수:{len(sample_text)}")
print(f"문자 단위 토큰 수:{len(char_tokens)} (vocab_size={vocab_size})")
print(f"BPE 토큰 수:{len(bpe_tokens)} (vocab_size={enc.n_vocab:,})")
print(f"같은 block_size={block_size} 안에 BPE는 문자 단위보다 약 "
      f"{len(char_tokens) / len(bpe_tokens):.1f}배 더 많은 원문을 담을 수 있음")

# 연습 5-1: 학습된 모델 저장
torch.save(model.state_dict(), './model/gpt_scratch.pt')
print("모델 저장 완료: gpt_scratch.pt")

%pip install peft tqdm

# 연습 5-2: peft로 LoRA 적용
from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,                                 # 저랭크(low-rank) 차원 — 클수록 표현력은 늘지만 파라미터도 늘어남
    lora_alpha=16,                       # LoRA 출력 스케일링 계수 (보통 r의 2배 정도로 설정)
    target_modules=["query", "value"],   # Head 클래스 안의 query/value Linear 레이어에만 적용
    lora_dropout=0.05,
)

lora_model = get_peft_model(model, lora_config)
lora_model.print_trainable_parameters()  # 전체 파라미터 대비 학습 대상(LoRA) 파라미터 비율 확인

# 이후 특정 스타일의 소량 텍스트로 섹션 7의 학습 루프를 다시 돌리면,
# 원본 파라미터는 얼어붙은(freeze) 채 LoRA 어댑터만 업데이트되어 훨씬 적은 자원으로 파인튜닝할 수 있다.