실습 4 — 근거를 붙여 답하게 하기

import json
from collections import Counter
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 실습 1과 같은 준비: 모델과 chat(), grade()
name = "Qwen/Qwen2.5-0.5B-Instruct"
device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16 if device == "cuda" else torch.float32).to(device)

def chat(messages, max_new_tokens=100):
    enc = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(device)
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
    return tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()

REFUSE = ["찾지 못", "찾을 수 없", "나와 있지 않", "제공되지 않"]
def grade(key, answer):
    refused = any(w in answer for w in REFUSE)
    if key == "":
        return "모름 맞음" if refused else "지어냄"
    if any(k in answer for k in key.split("|")):
        return "정답"
    return "거절" if refused else "오답"

# 실습 3과 같은 준비: 조항 나누기와 검색
text = open("shop_policy.txt", encoding="utf-8").read()
chunks = [c.strip() for c in text.split("\n\n")[1:]]
qa = [json.loads(line) for line in open("shop_policy_qa.jsonl", encoding="utf-8")]
vec = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 3))
X = vec.fit_transform(chunks)

def search(question, k=3):
    scores = cosine_similarity(vec.transform([question]), X)[0]
    return [(int(i), float(scores[i])) for i in scores.argsort()[::-1][:k]]

print("조항", len(chunks), "개, 질문", len(qa), "개 준비")

SYSTEM = ("당신은 쇼핑몰 상담원입니다. 아래 [문서]에 적힌 내용만 근거로 짧게 답하세요. "
          "문서에 답이 없으면 '문서에서 찾지 못했습니다.'라고만 답하세요.")

def build_messages(question, k=3):
    """검색한 조항 k 개를 [문서] 로 붙인 메시지 목록. 본문 "근거를 붙인 프롬프트"의 함수에 검색을 더했다"""
    found = [chunks[i] for i, _ in search(question, k)]
    user = "[문서]\n" + "\n".join(found) + "\n\n[질문]\n" + question
    return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}]

msgs = build_messages("휴대폰으로 결제하면 한도가 있나요?")
print(msgs[1]["content"])
# 모델에 실제로 들어가는 글(채팅 템플릿을 씌운 것)의 길이를 잰다. 문서를 붙인 만큼 길어진다
prompt_text = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
print("프롬프트 길이:", len(prompt_text), "글자 /", len(tok(prompt_text)["input_ids"]), "토큰")

NO_DOC = [{"role": "system", "content": "당신은 쇼핑몰 상담원입니다. 짧게 답하세요."}]
for q in ["휴대폰으로 결제하면 한도가 있나요?", "골드 등급이 되려면 얼마나 사야 해요?", "새벽 배송도 해 주나요?"]:
    print("질문    :", q)
    print("문서 없이:", chat(NO_DOC + [{"role": "user", "content": q}]).replace("\n", " ")[:90])
    print("문서 붙여:", chat(build_messages(q)).replace("\n", " ")[:90])

results = [grade(r["정답열쇠"], chat(build_messages(r["질문"]))) for r in qa]
print("문서에 있는 질문:", Counter(x for x, r in zip(results, qa) if r["근거조항"]))
print("문서에 없는 질문:", Counter(x for x, r in zip(results, qa) if not r["근거조항"]))