from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_path = "sooh098/kanana-ko-rag"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).eval()

instruction = """당신은 한국어 어문 규범(맞춤법, 띄어쓰기, 표준어, 문장부호, 외래어 표기법 등)에 따라 문장에서 올바른 표현을 선택하고 그 이유를 설명하는 AI입니다.

- 틀린 표현을 정확히 찾아 수정하되, 그 외 표현은 변경하지 마십시오.
- 답변은 반드시 다음 형식을 따라야 합니다:
  "{수정문}이 옳다. {이유}"

[예시]
문제: 다음 문장에서 어문 규범에 부합하지 않는 부분을 찾아 고치고, 그렇게 고친 이유를 설명하세요.
"어서 쾌차하시길 바래요."
정답: "어서 쾌차하시길 바라요."가 옳다. 동사 '바라다'에 어미 '-아요'가 결합한 형태이므로 '바라요'로 써야 한다. '바래요'는 비표준어다."""

question = "\"나는 그를 본 적이 있음을 {기억해냈다/기억해 냈다}.\" 가운데 올바른 것을 선택하고, 그 이유를 설명하세요."

# 참고 규범을 추가할 경우 프롬프트 형식은 다음과 같이 변경됩니다.
# input_text = f"[참고 규범]\n{retrieved_rules}\n=====\n문제: {question}\n정답:"

prompt = (
    "<|begin_of_text|>\n"
    f"[|system|]{instruction}<|eot_id|>\n"
    f"[|user|]문제: {question}\n정답:<|eot_id|>\n"
    "[|assistant|]"
)

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
inputs.pop("token_type_ids", None)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id
    )

result = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(result)

🔎 RAG 활용 방식

본 과제에서는 정답 생성 시 어문 규범 정보를 효과적으로 반영하기 위해 사전에 구축된 규범 문서 DB에서 관련 규범을 검색하고 이를 LLM 입력에 포함하는 RAG(Retrieval-Augmented Generation) 방식을 사용하였습니다.

검색에는 dense retrieval 방식을 적용하였고,
한국어 전용 임베딩 모델과 ChromaDB 기반의 규범 문서 벡터 DB를 결합하여 구현하였습니다.

✅ 규범 검색 방식

python

def retrieve_context(query_text, top_k=3):
    query_vec = embed_query(query_text)
    results = collection.query(query_embeddings=[query_vec], n_results=top_k * 4)

    seen_rules = set()
    contexts = []

    for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
        rule_text = meta.get("rule", "").strip()
        if rule_text not in seen_rules:
            seen_rules.add(rule_text)
            context = f"[{meta['title']}]\n{rule_text}"
            contexts.append(context)
        if len(contexts) == top_k:
            break

    return "\n\n".join(contexts)

임베딩 모델: dragonkue/snowflake-arctic-embed-l-v2.0-ko
문서 저장소: ChromaDB로 구축된 규범 벡터 DB
검색 대상: 띄어쓰기, 맞춤법, 표준어 등 규칙 본문

🧾 프롬프트 내 규범 삽입 예시

검색된 규범은 다음과 같이 모델 입력에 포함됩니다:

python

question = "가축을 기를 때에는 {먹이량/먹이양}을 조절해 주어야 한다."

input_text = f"[참고 규범]\n{retrieved_rules}\n=====\n문제: {question}\n정답:"

📑 인스트럭션 포맷

모델은 아래와 같은 프롬프트 구성을 사용하여 질의에 응답합니다. 이는 ‘AI 말평 경진대회’ 과제에서 요구한 형식을 반영한 구조입니다.

교정형

당신은 한국어 어문 규범(맞춤법, 띄어쓰기, 표준어, 문장부호, 외래어 표기법 등)에 따라 문장을 교정하고 그 이유를 설명하는 AI입니다.

[문제 유형: 교정형]
- 문제와 함께 관련 규범이 주어질 수 있으나, 규범의 타당성을 스스로 판단하여 정답을 도출해야 합니다.
- 제시된 문장은 반드시 어문 규범에 어긋난 표현을 포함하고 있습니다.
- 틀린 표현을 정확히 찾아 수정하되, 그 외 표현은 변경하지 마십시오.
- 답변은 반드시 다음 형식을 따라야 합니다:
  "{수정문}이 옳다. {이유}"

[예시]
문제: 다음 문장에서 어문 규범에 부합하지 않는 부분을 찾아 고치고, 그렇게 고친 이유를 설명하세요.
"어서 쾌차하시길 바래요."
정답: "어서 쾌차하시길 바라요."가 옳다. 동사 '바라다'에 어미 '-아요'가 결합한 형태이므로 '바라요'로 써야 한다. '바래요'는 비표준어다."

선택형

당신은 한국어 어문 규범(맞춤법, 띄어쓰기, 표준어, 문장부호, 외래어 표기법 등)에 따라 문장에서 올바른 표현을 선택하고 그 이유를 설명하는 AI입니다.

[문제 유형: 선택형]
- 문제와 함께 관련 규범이 주어질 수 있으나, 규범의 타당성을 스스로 판단하여 정답을 도출해야 합니다.
- 보기 중 단 하나의 정답이 있으며, 반드시 어문 규범에 따라 하나의 표현만 선택해야 합니다.
- 선택한 표현 이외의 문장 구성은 수정하지 마십시오.
- 답변은 반드시 다음 형식을 따라야 합니다:
  "{정답문}이 옳다. {이유}"

[예시]
문제: "가축을 기를 때에는 {먹이량/먹이양}을 조절해 주어야 한다." 가운데 올바른 것을 선택하고, 그 이유를 설명하세요.
정답: "가축을 기를 때에는 먹이양을 조절해 주어야 한다."가 옳다. '먹이'는 고유어이므로, 한자어 '量'은 두음 법칙에 따라 '양'으로 표기해야 한다.

사용자

[참고 규범]
{retrieved_rules}
=====
문제: {question}
정답:

📎 참고

데이터 출처: 국립국어원 문화·규범 QA 말뭉치 V1.0