실습 5 — 선호 데이터와 로그 확률

import json
from collections import Counter

with open("dpo_shop.jsonl", encoding="utf-8") as f:
    pairs = [json.loads(line) for line in f]
print("쌍의 수:", len(pairs))
print(json.dumps(pairs[0], ensure_ascii=False, indent=1))
print(Counter(p["rejected"][:12] for p in pairs))

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float32).eval()
SYSTEM = "당신은 온라인 쇼핑몰 고객센터 상담원입니다."

def answer_logprob(q, a):
    """질문 q 다음에 답 a 가 나올 로그 확률의 평균(토큰 하나당)을 돌려준다."""
    msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": q}]
    p_ids = tok(tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False), add_special_tokens=False)["input_ids"]
    f_ids = tok(tok.apply_chat_template(msgs + [{"role": "assistant", "content": a}], tokenize=False), add_special_tokens=False)["input_ids"]
    ids = torch.tensor([f_ids])
    with torch.no_grad():
        logits = model(ids).logits[0, :-1]                    # 자리마다 "다음 토큰"의 점수
    logp = torch.log_softmax(logits, dim=-1)                  # 점수를 로그 확률로
    target = ids[0, 1:]                                       # 실제 다음 토큰
    token_logp = logp[torch.arange(len(target)), target]      # 실제 다음 토큰의 로그 확률
    return token_logp[len(p_ids) - 1:].mean().item()          # 답 부분만 평균

p = pairs[0]
print("고른 답:", round(answer_logprob(p["prompt"], p["chosen"]), 3))
print("반대 답:", round(answer_logprob(p["prompt"], p["rejected"]), 3))

prefer_chosen = 0
for p in pairs[::6]:                      # 72쌍 가운데 6개 건너 하나씩, 12쌍
    c = answer_logprob(p["prompt"], p["chosen"])
    r = answer_logprob(p["prompt"], p["rejected"])
    prefer_chosen += c > r
    print(f"{'고른 답' if c > r else '반대 답'} 쪽  고른 {c:6.3f}  반대 {r:6.3f}  {p['prompt']}")
print(f"고른 답을 더 낼 것 같은 쌍: {prefer_chosen} / 12")