실습 4 — 느린 RAG 의 병목 찾기

import json, re, time, urllib.request, urllib.error
import numpy as np
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):
    """답과 Ollama 가 잰 숫자를 돌려준다. 서버 오류나 깨진 답이면 모델을 내렸다 두 번까지 다시 묻는다"""
    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


def correct(answer, key):
    """정답열쇠(세로줄로 나눈 여러 낱말) 가운데 하나라도 답에 들어 있으면 맞힘"""
    return any(k in answer.replace(" ", "") for k in key.replace(" ", "").split("|"))


articles = [l.strip() for l in open("company_rules.txt", encoding="utf-8") if re.match(r"^제\d+조\(", l)]
qa = [json.loads(l) for l in open("company_rules_qa.jsonl", encoding="utf-8")]
SYSTEM = "당신은 솔내기술(가상 회사)의 규정 도우미입니다. [조항]에 있는 내용만으로 답하세요."

# 모델 불러오기는 처음 한 번만 드는 시간이라 따로 잰다
start = time.perf_counter()
emb = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")     # GPU 는 Ollama 에 비워 둔다
ce = CrossEncoder("cross-encoder/mmarco-mMiniLMv2-L12-H384-v1", device="cpu")
unload_all()
chat([{"role": "user", "content": "안녕"}], 5)                                 # 몸풀기: 생성 모델을 올려 둔다
print(f"모델 불러오기 {time.perf_counter() - start:.1f}초 · 조항 {len(articles)}개 · 질문 {len(qa)}개")

def slow_rag(q):
    """일부러 느리게 만든 RAG. 단계마다 걸린 시간을 t 에 적는다"""
    t = {}
    s = time.perf_counter()
    # 1 색인: 질문마다 조항 30개를 다시 임베딩한다(낭비)
    A = emb.encode(["passage: " + a for a in articles], normalize_embeddings=True)
    t["색인"] = time.perf_counter() - s; s = time.perf_counter()
    # 2 검색: 질문을 임베딩해 모든 조항과의 가까움(내적)으로 줄 세운다
    qv = emb.encode(["query: " + q], normalize_embeddings=True)[0]
    cand = np.argsort(-(A @ qv))
    t["검색"] = time.perf_counter() - s; s = time.perf_counter()
    # 3 재정렬: 조항 30개를 모두 크로스 인코더에 넣는다(낭비)
    scores = ce.predict([(q, articles[i]) for i in cand])
    cand = cand[np.argsort(-scores)]
    t["재정렬"] = time.perf_counter() - s; s = time.perf_counter()
    # 4 생성: 조각 8개를 붙이고 자세히 설명하게 한다(입력도 출력도 길다)
    ctx = "\n".join(articles[i] for i in cand[:8])
    r = chat([{"role": "system", "content": SYSTEM},
              {"role": "user", "content": f"[조항]\n{ctx}\n\n[질문]\n{q}\n자세히 설명해 주세요."}], num_predict=300)
    t["생성"] = time.perf_counter() - s
    return r["message"]["content"].strip(), t, r["prompt_eval_count"], r["eval_count"]


answer, t, n_in, n_out = slow_rag(qa[0]["질문"])
# 결과: 질문 하나의 단계별 시간(초)과 입력·출력 토큰 수
print({k: round(v, 2) for k, v in t.items()}, "· 입력", n_in, "· 출력", n_out)
print(answer[:120])

import pandas as pd

rows = []
for x in qa:
    a, t, n_in, n_out = slow_rag(x["질문"])
    rows.append({**t, "전체": sum(t.values()), "입력토큰": n_in, "출력토큰": n_out, "정답": correct(a, x["정답열쇠"])})
slow = pd.DataFrame(rows)
stages = ["색인", "검색", "재정렬", "생성"]
summary = pd.DataFrame({"평균(초)": slow[stages].mean().round(3), "몫": (slow[stages].sum() / slow["전체"].sum()).round(3)})
# 결과: 단계마다 평균 시간과 전체에서의 몫
print(summary.to_string())
print(f"질문 하나 평균 {slow['전체'].mean():.2f}초 · p95 {slow['전체'].quantile(0.95):.2f}초 · "
      f"평균 입력 {slow['입력토큰'].mean():.0f}토큰 · 평균 출력 {slow['출력토큰'].mean():.0f}토큰 · "
      f"정답 {slow['정답'].sum()}/{len(qa)} · 다시 물음 {retried['수']}")

import matplotlib.pyplot as plt
plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False

fig, ax = plt.subplots(figsize=(7, 3.5))
ax.barh(stages[::-1], summary["평균(초)"][::-1])     # 위에서부터 단계 순서대로 보이게 뒤집는다
for i, v in enumerate(summary["평균(초)"][::-1]):
    ax.text(v, i, f" {v:.2f}초", va="center")
ax.set_xlabel("질문 하나의 평균 시간(초)")
ax.set_title("느린 RAG 의 단계별 시간")
ax.set_xlim(0, summary["평균(초)"].max() * 1.25)
plt.tight_layout()
plt.show()