실습 6 — 작은 챗봇 모델 써 보기

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(name)
# CausalLM: 앞만 보고 다음 토큰을 맞히는(디코더만 쓰는) 모델. float32 로 CPU 에서 돌린다
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float32)
model.eval()
print("파라미터 수:", sum(p.numel() for p in model.parameters()))

# 대화는 역할(role)과 내용(content)을 가진 메시지의 목록이다
messages = [
    {"role": "system", "content": "당신은 쇼핑몰 고객센터 상담원입니다."},
    {"role": "user", "content": "배송이 늦어요."},
]
# 모델이 학습 때 본 꼴의 글로 바꾼다. add_generation_prompt=True 는 "이제 상담원 차례"라는 표시를 붙인다
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
print(prompt)

def ask(system, question, max_new_tokens=80):
    """시스템 메시지와 질문을 넣고 모델의 답만 돌려준다."""
    messages = [{"role": "system", "content": system}, {"role": "user", "content": question}]
    enc = tokenizer.apply_chat_template(messages, add_generation_prompt=True,
                                        return_tensors="pt", return_dict=True)
    with torch.no_grad():
        # do_sample=False: 확률대로 뽑지 않고 늘 1등 토큰을 고른다(탐욕적 선택). 다시 돌려도 같은 답
        out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
    new_tokens = out[0][enc["input_ids"].shape[1]:]     # 입력 뒤에 새로 생긴 토큰만
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

system = "당신은 온라인 쇼핑몰 고객센터 상담원입니다. 한국어로 두 문장 안에서 짧고 친절하게 답합니다."
for q in ["주문한 지 일주일이 지났는데 아직 안 왔어요.", "환불은 며칠 걸려요?"]:
    print("고객:", q)
    print("모델:", ask(system, q))

# 같은 질문에 말투만 바꾼 시스템 메시지를 준다
playful = "당신은 반말로 장난스럽게 말하는 쇼핑몰 마스코트입니다. 한 문장으로 답합니다."
q = "주문한 지 일주일이 지났는데 아직 안 왔어요."
print("시스템 메시지를 바꾼 답:", ask(playful, q))