프로젝트 1 — 친절한 빵집 상담 모델 만들기

import json, time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, set_seed
from peft import LoraConfig, get_peft_model
from datasets import Dataset

name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
SYSTEM = "당신은 동네 빵집의 안내 직원입니다."

with open("bakery_qa.jsonl", encoding="utf-8") as f:
    rows = [json.loads(line) for line in f]
for r in rows:
    # 메모 말투: 인사와 끝인사를 떼고 사실만 남긴다
    r["메모"] = r["답변"].removeprefix("안녕하세요. ").removesuffix(" 맛있게 드세요.")
train_rows = [r for r in rows if r["나눔"] == "train"]
test_rows = [r for r in rows if r["나눔"] == "test"]
print("친절:", train_rows[0]["답변"])
print("메모:", train_rows[0]["메모"])

def messages(q, a=None):
    m = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": q}]
    return m + ([{"role": "assistant", "content": a}] if a is not None else [])

def encode(batch):
    input_ids, labels = [], []
    for q, a in zip(batch["질문"], batch["메모"]):
        p = tok(tok.apply_chat_template(messages(q), add_generation_prompt=True, tokenize=False), add_special_tokens=False)["input_ids"]
        f = tok(tok.apply_chat_template(messages(q, a), tokenize=False), add_special_tokens=False)["input_ids"]
        input_ids.append(f)
        labels.append([-100] * len(p) + f[len(p):])
    return {"input_ids": input_ids, "labels": labels}

def collate(features):
    L = max(len(x["input_ids"]) for x 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)
    for i, x in enumerate(features):
        n = len(x["input_ids"])
        ids[i, :n], lab[i, :n], att[i, :n] = torch.tensor(x["input_ids"]), torch.tensor(x["labels"]), 1
    return {"input_ids": ids, "labels": lab, "attention_mask": att}

def ask(m, q):
    enc = tok.apply_chat_template(messages(q), add_generation_prompt=True, return_tensors="pt", return_dict=True).to(m.device)
    with torch.no_grad():
        out = m.generate(**enc, max_new_tokens=60, do_sample=False)
    return tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip().replace("\n", " ")

def report(m, label):
    answers = [ask(m, r["질문"]) for r in test_rows]
    hit = sum(r["정답열쇠"] in a for r, a in zip(test_rows, answers))
    kind = sum(("안녕하세요" in a) and ("맛있게 드세요" in a) for a in answers)
    print(f"[{label}] 정답 {hit}/20, 친절한 말투 {kind}/20")
    for r, a in list(zip(test_rows, answers))[:3]:
        print("   ", r["질문"], "→", a[:80])

lora = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], task_type="CAUSAL_LM")
set_seed(42)
model = get_peft_model(AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16), lora)
sft_data = Dataset.from_list(train_rows).map(encode, batched=True, remove_columns=list(train_rows[0].keys()))
args = TrainingArguments("out_bakery_lora_ckpt", per_device_train_batch_size=8, num_train_epochs=8, learning_rate=5e-4,
                         save_strategy="no", report_to="none", seed=42, fp16=torch.cuda.is_available(),
                         disable_tqdm=True, logging_steps=100)
sft_trainer = Trainer(model=model, args=args, train_dataset=sft_data, data_collator=collate)
sft_trainer.train()
sft_model = sft_trainer.model.eval()
report(sft_model, "SFT 뒤")

from trl import DPOConfig, DPOTrainer

merged = sft_model.merge_and_unload()            # SFT 의 LoRA 를 합쳐 새 출발점으로
pairs = Dataset.from_list([{
    "prompt": messages(r["질문"]),
    "chosen": [{"role": "assistant", "content": r["답변"]}],      # 친절한 말투
    "rejected": [{"role": "assistant", "content": r["메모"]}],    # 메모 말투
} for r in train_rows])
config = DPOConfig("out_bakery_dpo_ckpt", beta=0.1, per_device_train_batch_size=4, num_train_epochs=2,
                   learning_rate=3e-4, max_length=256, save_strategy="no", report_to="none", seed=42,
                   fp16=torch.cuda.is_available(), disable_tqdm=True, logging_steps=100,
                   # DPO 손실(sigmoid)에 고른 답의 SFT 손실(sft)을 1:1 로 섞는다. 문장이 무너지는 것을 막는다
                   loss_type=["sigmoid", "sft"], loss_weights=[1.0, 1.0])
set_seed(42)
dpo_trainer = DPOTrainer(model=merged, args=config, train_dataset=pairs, processing_class=tok, peft_config=lora)
dpo_trainer.train()
dpo_model = dpo_trainer.model.eval()
report(dpo_model, "DPO 뒤")

# 빵집과 상관없는 질문: 원래 모델이 하던 일을 잊지 않았는지 본다
for q in ["대한민국의 수도는 어디인가요?", "1부터 10까지 더하면 얼마인가요?"]:
    print("Q:", q)
    print("A:", ask(dpo_model, q)[:100])