실습 5 — 캐시로 같은 일 줄이기

import json, re, time, urllib.request, urllib.error
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer, CrossEncoder

HOST = "http://127.0.0.1:11434"
MODEL = "qwen2.5:3b"
BROKEN = re.compile(r"(.)\1{9,}")
retried = {"수": 0}


def post(path, body):
    req = urllib.request.Request(HOST + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    return json.loads(urllib.request.urlopen(req, timeout=300).read())


def unload_all():
    for m in json.loads(urllib.request.urlopen(HOST + "/api/ps").read())["models"]:
        post("/api/generate", {"model": m["name"], "keep_alive": 0})


def chat(messages, num_predict):
    """실습 4와 같은 함수"""
    for _ in range(3):
        try:
            r = post("/api/chat", {"model": MODEL, "stream": False, "messages": messages,
                                   "options": {"temperature": 0, "seed": 42, "num_predict": num_predict}})
            if not BROKEN.search(r["message"]["content"]):
                return r
        except urllib.error.HTTPError:
            pass
        retried["수"] += 1
        post("/api/generate", {"model": MODEL, "keep_alive": 0})
    return r


articles = [l.strip() for l in open("company_rules.txt", encoding="utf-8") if re.match(r"^제\d+조\(", l)]
SYSTEM = "당신은 솔내기술(가상 회사)의 규정 도우미입니다. [조항]에 있는 내용만으로 답하세요."
emb = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")
ce = CrossEncoder("cross-encoder/mmarco-mMiniLMv2-L12-H384-v1", device="cpu")
# 색인은 여기서 한 번만 만든다(실습 4는 질문마다 다시 만들었다)
A = emb.encode(["passage: " + a for a in articles], normalize_embeddings=True)


def rag(q):
    """색인을 다시 쓰고, 재정렬은 상위 10개만, 조각 3개로 한 문장 답을 만든다"""
    qv = emb.encode(["query: " + q], normalize_embeddings=True)[0]
    cand = np.argsort(-(A @ qv))[:10]
    cand = cand[np.argsort(-ce.predict([(q, articles[i]) for i in cand]))]
    ctx = "\n".join(articles[i] for i in cand[:3])
    r = chat([{"role": "system", "content": SYSTEM},
              {"role": "user", "content": f"[조항]\n{ctx}\n\n[질문]\n{q}\n한 문장으로 답하세요."}], num_predict=80)
    return r["message"]["content"].strip()


unload_all()
rag("안녕하세요")                                       # 몸풀기
log = pd.read_csv("rag_query_log.csv")
# 결과: 질문 기록의 앞부분
print(log.head(5).to_string(index=False))

def normalize(q):
    """빈칸을 하나로 줄이고 앞뒤 빈칸과 끝의 물음표·마침표를 지운다"""
    return re.sub(r"\s+", " ", q).strip().rstrip("?.").strip()


raw_unique = log["질문"].nunique()
norm_unique = log["질문"].map(normalize).nunique()
# 결과: 글자 그대로 다른 질문 수와, 다듬은 뒤 다른 질문 수
print(f"질문 {len(log)}줄 · 글자 그대로 {raw_unique}가지 · 다듬으면 {norm_unique}가지")
# 가장 많이 물은 질문 셋
print(log["질문"].map(normalize).value_counts().head(3).to_string())

def run_log(key_fn=None):
    """질문 기록을 차례로 처리한다. key_fn 이 있으면 그 열쇠로 답 캐시를 쓴다"""
    cache, hits = {}, 0
    start = time.perf_counter()
    for q in log["질문"]:
        if key_fn is None:
            rag(q)
            continue
        key = key_fn(q)
        if key in cache:
            hits += 1                                  # 캐시에 있으면 검색도 생성도 하지 않는다
        else:
            cache[key] = rag(q)
    return round(time.perf_counter() - start, 1), hits


rows = []
for label, fn in [("캐시 없음", None), ("글자 그대로", lambda q: q), ("다듬은 질문", normalize)]:
    sec, hits = run_log(fn)
    rows.append({"방법": label, "걸린 시간(초)": sec, "캐시 맞힘": hits, "질문 하나 평균(초)": round(sec / len(log), 2)})
# 결과: 80줄을 처리한 시간과 캐시에서 바로 답한 수
print(pd.DataFrame(rows).to_string(index=False))
print("다시 물음", retried["수"])

# 조항 30개 전체를 시스템 프롬프트에 넣는다. 질문마다 바뀌는 것은 맨 뒤의 질문뿐이다
long_system = SYSTEM + "\n[조항]\n" + "\n".join(articles)
unload_all()                                            # 모델을 내려 Ollama 의 프롬프트 캐시도 비운다
rows = []
for q in ["재택 근무는 며칠까지 되나요?", "반차를 쓰면 몇 시에 출근하나요?", "사원증을 잃어버렸어요."]:
    r = chat([{"role": "system", "content": long_system},
              {"role": "user", "content": q + " 한 문장으로 답하세요."}], num_predict=60)
    # prompt_eval_cached_count: 지난 요청과 앞부분이 같아 다시 계산하지 않은 입력 토큰 수
    rows.append({"질문": q, "입력토큰": r["prompt_eval_count"], "캐시토큰": r.get("prompt_eval_cached_count", 0),
                 "프리필(초)": round(r["prompt_eval_duration"] / 1e9, 2), "전체(초)": round(r["total_duration"] / 1e9, 2)})
# 결과: 첫 질문은 긴 앞부분을 모두 계산하고, 다음 질문부터는 앞부분을 캐시에서 읽는다
print(pd.DataFrame(rows).to_string(index=False))