프로젝트 1 — 우리 동네 빵집 맞춤 모델 만들기

import json, math, os
import torch
from tokenizers import ByteLevelBPETokenizer
from transformers import (PreTrainedTokenizerFast, GPT2Config, GPT2LMHeadModel, Trainer,
                          TrainingArguments, DataCollatorForLanguageModeling, set_seed)
from datasets import Dataset

bpe = ByteLevelBPETokenizer()
bpe.train(["bakery_corpus.txt"], vocab_size=2000, min_frequency=2, special_tokens=["<pad>", "<eos>"])
os.makedirs("out_bakery_tok", exist_ok=True)
bpe.save("out_bakery_tok/tokenizer.json")
tok = PreTrainedTokenizerFast(tokenizer_file="out_bakery_tok/tokenizer.json", pad_token="<pad>", eos_token="<eos>")
print("어휘:", tok.vocab_size, "/ '통밀 깜파뉴' →", [tok.decode([i]) for i in tok("통밀 깜파뉴")["input_ids"]])

with open("bakery_corpus.txt", encoding="utf-8") as f:
    lines = [line for line in f.read().splitlines() if line]
ds = Dataset.from_dict({"text": lines}).train_test_split(test_size=0.05, seed=42)
BLOCK = 128

def group(batch):
    ids = sum(batch["input_ids"], [])
    n = len(ids) // BLOCK * BLOCK
    return {"input_ids": [ids[i:i + BLOCK] for i in range(0, n, BLOCK)]}

lm_data = ds.map(lambda b: tok([t + "<eos>" for t in b["text"]]), batched=True, remove_columns=["text"]) \
            .map(group, batched=True, remove_columns=["attention_mask"])

set_seed(42)
config = GPT2Config(vocab_size=len(tok), n_positions=BLOCK, n_embd=256, n_layer=4, n_head=4,
                    bos_token_id=tok.eos_token_id, eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id)
model = GPT2LMHeadModel(config)
args = TrainingArguments("out_bakery_ckpt", per_device_train_batch_size=32, num_train_epochs=8, learning_rate=1e-3,
                         warmup_steps=50, save_strategy="no", report_to="none", seed=42,
                         fp16=torch.cuda.is_available(), disable_tqdm=True, logging_steps=1000)
trainer = Trainer(model=model, args=args, train_dataset=lm_data["train"], eval_dataset=lm_data["test"],
                  data_collator=DataCollatorForLanguageModeling(tok, mlm=False))
trainer.train()
print("검증 퍼플렉서티:", round(math.exp(trainer.evaluate()["eval_loss"]), 2))
trainer.model.save_pretrained("out_bakery_pre")
tok.save_pretrained("out_bakery_pre")

with open("bakery_qa.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"]

def prompt(q):
    return f"### 질문:\n{q}\n### 답변:\n"

def encode(batch):
    input_ids, labels = [], []
    for q, a in zip(batch["질문"], batch["답변"]):
        p, r = tok(prompt(q))["input_ids"], tok(a + "<eos>")["input_ids"]
        input_ids.append(p + r)
        labels.append([-100] * len(p) + r)
    return {"input_ids": input_ids, "labels": labels}

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

sft_data = Dataset.from_list(train_rows).map(encode, batched=True, remove_columns=list(train_rows[0].keys()))
set_seed(42)
sft_args = TrainingArguments("out_bakery_ckpt", per_device_train_batch_size=16, num_train_epochs=30,
                             learning_rate=5e-4, save_strategy="no", report_to="none", seed=42,
                             fp16=torch.cuda.is_available(), disable_tqdm=True, logging_steps=1000)
sft_trainer = Trainer(model=GPT2LMHeadModel.from_pretrained("out_bakery_pre"), args=sft_args,
                      train_dataset=sft_data, data_collator=collate)
sft_trainer.train()
print("SFT 끝:", len(train_rows), "쌍")

model = sft_trainer.model.eval()

def answer(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()

hits = 0
for r in test_rows:
    a = answer(r["질문"])
    ok = r["정답열쇠"] in a
    hits += ok
    print("O" if ok else "X", f"[{r['의도']}]", r["질문"], "→", a)
print(f"맞힌 수: {hits} / {len(test_rows)}")