확률적 디코딩

!pip install -q "transformers==4.55.4" accelerate jinja2

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "LGAI-EXAONE/EXAONE-4.0-1.2B"  # 허브 저장소 식별자
tokenizer = AutoTokenizer.from_pretrained(model_name)  # 토크나이저 로드
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # T4 GPU용 16비트 자료형
    device_map="auto",         # 사용 가능한 GPU에 자동 배치
)

logits = torch.tensor([2.0, 1.0, 0.0])  # 세 토큰의 가상 점수

for temperature in [0.7, 1.0, 1.3]:
    probabilities = torch.softmax(
        logits / temperature,
        dim=-1,  # 토큰 차원의 점수를 확률로 변환
    )
    print(temperature, probabilities)

prompt = "인공지능은 데이터를 학습하여 패턴을 찾는다. 이를 활용하면"  # 공통 입력
inputs = tokenizer(
    prompt,
    return_tensors="pt",  # 토큰 ID를 PyTorch 텐서로 반환
).to(model.device)         # 입력을 모형과 같은 장치로 이동

input_length = inputs["input_ids"].shape[1]  # 입력 토큰 수
inputs["input_ids"].shape  # 배치 크기와 입력 토큰 수

def sample_text(seed, **decoding_args):  # 디코딩 설정을 키워드 인자로 받음
    torch.manual_seed(seed)  # 반복 비교를 위해 난수 시드 고정
    with torch.inference_mode():  # 기울기 계산을 꺼 메모리와 연산 절약
        output_ids = model.generate(
            **inputs,           # input_ids와 attention_mask 전달
            max_new_tokens=24,  # 입력 뒤에 최대 24개 토큰 생성
            do_sample=True,     # 확률분포에서 다음 토큰을 추출
            use_cache=True,     # 이전 토큰의 Key·Value를 KV 캐시에 저장해 재사용
            **decoding_args,    # temperature, top_k, top_p 등을 전달
        )
    return tokenizer.decode(
        output_ids[0, input_length:],  # 첫 결과에서 입력을 제외한 토큰만 선택
        skip_special_tokens=True,     # 제어용 특수 토큰은 문자열에서 제외
    )

for temperature in [0.7, 1.0, 1.3]:
    text = sample_text(
        seed=5,                  # 동일한 난수 순서로 비교
        temperature=temperature,  # 현재 반복의 온도
        top_k=50,                # 상위 50개 토큰만 후보로 유지
        top_p=1.0,               # 누적 확률에 따른 후보 제한은 끄기
    )
    print(temperature, text)

for top_k in [5, 20, 50]:
    text = sample_text(
        seed=1,            # 동일한 난수 순서로 비교
        temperature=1.0,  # 원래 확률분포 유지
        top_k=top_k,       # 현재 반복의 후보 개수
        top_p=1.0,         # 누적 확률에 따른 후보 제한은 끄기
    )
    print(top_k, text)

for top_p in [0.7, 0.9, 0.95]:
    text = sample_text(
        seed=5,            # 동일한 난수 순서로 비교
        temperature=1.0,  # 원래 확률분포 유지
        top_k=0,           # top-k 후보 제한은 끄기
        top_p=top_p,       # 현재 반복의 누적 확률 기준
    )
    print(top_p, text)

for seed in [3, 6, 9]:
    text = sample_text(
        seed=seed,        # 현재 반복의 난수 시드
        temperature=0.8,  # 상위 확률 토큰에 조금 더 집중
        top_k=50,         # 상위 50개 토큰만 후보로 유지
        top_p=0.9,        # 누적 확률 0.9에 도달하는 최소 후보만 유지
    )
    print(seed, text)