실습 6 — 상담 말투로 SFT 하기

import json
import torch
from transformers import PreTrainedTokenizerFast, GPT2LMHeadModel, set_seed
from datasets import Dataset

tok = PreTrainedTokenizerFast.from_pretrained("out_pretrained")      # 실습 5에서 함께 저장한 토크나이저

# jsonl: 한 줄이 JSON 하나. 의도, 질문, 답변, 정답열쇠, 나눔(train/test)
with open("sft_shop.jsonl", encoding="utf-8") as f:
    rows = [json.loads(line) for line in f]
train_rows = [r for r in rows if r["나눔"] == "train"]
test_rows = [r for r in rows if r["나눔"] == "test"]
print("학습용:", len(train_rows), "/ 시험용:", len(test_rows))
print(train_rows[0])

def prompt(q):
    """학습할 때와 물어볼 때 똑같이 쓰는 템플릿."""
    return f"### 질문:\n{q}\n### 답변:\n"

def encode(batch):
    input_ids, labels = [], []
    for q, a in zip(batch["질문"], batch["답변"]):
        p = tok(prompt(q))["input_ids"]            # 질문 부분의 토큰
        r = tok(a + "<eos>")["input_ids"]          # 답 부분의 토큰. 끝 표시를 붙인다
        input_ids.append(p + r)
        labels.append([-100] * len(p) + r)         # 질문 자리는 채점하지 않는다
    return {"input_ids": input_ids, "labels": labels}

sft_data = Dataset.from_list(train_rows).map(encode, batched=True, remove_columns=list(train_rows[0].keys()))

def collate(features):
    """길이가 다른 줄을 가장 긴 줄에 맞춰 채운다. 채운 자리의 labels 는 -100."""
    L = max(len(f["input_ids"]) for f in features)
    ids = torch.full((len(features), L), tok.pad_token_id)
    lab = torch.full((len(features), L), -100)
    att = torch.zeros((len(features), L), dtype=torch.long)     # 1 인 자리만 진짜 토큰
    for i, f in enumerate(features):
        n = len(f["input_ids"])
        ids[i, :n] = torch.tensor(f["input_ids"])
        lab[i, :n] = torch.tensor(f["labels"])
        att[i, :n] = 1
    return {"input_ids": ids, "labels": lab, "attention_mask": att}

ex = sft_data[0]
print("토큰 수:", len(ex["input_ids"]), "/ 채점하는 자리:", sum(x != -100 for x in ex["labels"]))

def answer(model, q):
    """템플릿으로 묻고, 답 부분만 돌려준다(탐욕적 선택)."""
    enc = tok(prompt(q), return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=60, do_sample=False,
                             eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id)
    return tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()

pre = GPT2LMHeadModel.from_pretrained("out_pretrained")
if torch.cuda.is_available():
    pre = pre.to("cuda")
for q in ["배송비 알려 주세요.", "환불 처리 기간을 알려 주세요."]:
    print(f"[SFT 전] {q} → {answer(pre.eval(), q)}")

from transformers import Trainer, TrainingArguments

set_seed(42)
args = TrainingArguments("out_sft_ckpt", per_device_train_batch_size=16, num_train_epochs=30,
                         learning_rate=5e-4, logging_steps=50, save_strategy="no", report_to="none",
                         seed=42, fp16=torch.cuda.is_available(), disable_tqdm=True)
trainer = Trainer(model=GPT2LMHeadModel.from_pretrained("out_pretrained"), args=args,
                  train_dataset=sft_data, data_collator=collate)
trainer.train()
sft = trainer.model.eval()

# 학습에 쓰지 않은 시험 질문 24개: 답에 정답 열쇠가 들어 있는지 센다
hits = 0
for r in test_rows:
    a = answer(sft, r["질문"])
    ok = r["정답열쇠"] in a
    hits += ok
    print("O" if ok else "X", r["질문"], "→", a)
print(f"맞힌 수: {hits} / {len(test_rows)}")

sft.save_pretrained("out_sft_shop")
tok.save_pretrained("out_sft_shop")