# GPT-2의 Zero-shot 개념
# Fine-tuning 없이 프롬프트만으로 태스크 수행

# 번역 (학습하지 않았지만 수행)
prompt = "Translate English to French: cheese =>"
# 모델 출력: "fromage"

# 요약 (학습하지 않았지만 수행)
prompt = "Article: [긴 기사 텍스트]\nTL;DR:"
# 모델 출력: 기사 요약

# GPT-2의 주요 아키텍처 변경점

# 1. Pre-Norm (Layer Normalization 위치 변경)
# GPT-1: Post-Norm (Attention → Add → LayerNorm)
# GPT-2: Pre-Norm (LayerNorm → Attention → Add) → 학습 안정성 향상

class GPT2Block:
    """GPT-2 Transformer 블록의 개념적 구조입니다."""
    def forward(self, x):
        # Pre-Norm 구조
        residual = x
        x = self.layer_norm_1(x)     # LayerNorm 먼저
        x = self.attention(x)         # Attention
        x = residual + x              # Residual 연결

        residual = x
        x = self.layer_norm_2(x)     # LayerNorm 먼저
        x = self.feed_forward(x)     # FFN
        x = residual + x              # Residual 연결
        return x

# 2. 어휘 크기 확대: 40,478 → 50,257 (BPE)
# 3. 최종 LayerNorm 추가 (마지막 블록 출력에)

from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# GPT-2 모델과 토크나이저 로드
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()

# 텍스트 생성
prompt = "Artificial intelligence will"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

# 다양한 생성 전략
with torch.no_grad():
    # Greedy 생성
    greedy = model.generate(input_ids, max_new_tokens=50)
    print("Greedy:", tokenizer.decode(greedy[0], skip_special_tokens=True))

    # Top-k 샘플링
    top_k = model.generate(
        input_ids,
        max_new_tokens=50,
        do_sample=True,
        top_k=50,
        temperature=0.7,
    )
    print("Top-k:", tokenizer.decode(top_k[0], skip_special_tokens=True))

    # Nucleus (Top-p) 샘플링
    top_p = model.generate(
        input_ids,
        max_new_tokens=50,
        do_sample=True,
        top_p=0.9,
        temperature=0.7,
    )
    print("Top-p:", tokenizer.decode(top_p[0], skip_special_tokens=True))