hyungi
d53fcc2b36
맥미니 llama.cpp 리랭크는 후보 수 선형(실측 50=0.60s/200=1.89s) — NAS 배포에서 MAX_RERANK_INPUT=50 으로 tail 지연 축소. 기본 200 = 현행 무회귀. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
276 lines
9.7 KiB
Python
276 lines
9.7 KiB
Python
"""Reranker 서비스 — bge-reranker-v2-m3 통합 (Phase 1.3).
|
|
|
|
TEI 컨테이너 호출 + asyncio.Semaphore(2) + soft timeout fallback.
|
|
|
|
데이터 흐름 원칙:
|
|
- fusion = doc 기준 / reranker = chunk 기준 — 절대 섞지 말 것
|
|
- raw chunks를 끝까지 보존, fusion은 압축본만 사용
|
|
- reranker는 chunks_by_doc dict에서 raw chunks 회수해서 chunk 단위로 호출
|
|
- diversity는 reranker 직후 마지막 단계에서만 적용
|
|
|
|
snippet 생성:
|
|
- 200~400 토큰(800~1500자) 기준
|
|
- query keyword 위치 중심 ±target_chars/2 윈도우
|
|
- keyword 매치 없으면 첫 target_chars 문자 fallback (성능 손실 방지)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import re
|
|
from typing import TYPE_CHECKING
|
|
|
|
import httpx
|
|
|
|
from ai.client import AIClient
|
|
from core.utils import setup_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from api.search import SearchResult
|
|
|
|
logger = setup_logger("rerank")
|
|
|
|
# 동시 rerank 호출 제한 (GPU saturation 방지)
|
|
RERANK_SEMAPHORE = asyncio.Semaphore(2)
|
|
|
|
# rerank input 크기 제한 (latency / VRAM hard cap).
|
|
# 2노드 이관(2026-07-02): env MAX_RERANK_INPUT 로 조정 가능 — 맥미니 llama.cpp 리랭크는
|
|
# 후보 수에 선형(NAS발 실측 50=0.60s / 100=0.95s / 200=1.89s)이라 NAS 배포는 50 권장.
|
|
# 기본 200 = 현행(GPU TEI) 무회귀.
|
|
MAX_RERANK_INPUT = int(os.getenv("MAX_RERANK_INPUT", "200"))
|
|
MAX_CHUNKS_PER_DOC = 2
|
|
|
|
# Soft timeout (초)
|
|
RERANK_TIMEOUT = 5.0
|
|
|
|
# ─── Phase 2B Diagnose dispatcher (R2-B1 slug-based) ──────────────
|
|
# server-side allowlist map. query parameter 가 raw endpoint URL 받지 않음.
|
|
RERANKER_BACKEND_MAP: dict[str, dict[str, str] | None] = {
|
|
"baseline": None, # production reranker (config.yaml endpoint via AIClient.rerank)
|
|
"cand_gte_ml_base": {
|
|
"endpoint": "http://rerank-cand-gte-ml-base:80/rerank",
|
|
},
|
|
# mxbai_large 후보 (deberta-v2 → TEI 1.7 미지원) Phase 2B-Extended 이관
|
|
# bge_v2_gemma_2b 후보 (LLM-based reranker, 1_Pooling/config.json 부재) Phase 2B-Extended 이관
|
|
}
|
|
|
|
|
|
def _resolve_reranker(slug: str | None) -> str | None:
|
|
"""slug → endpoint URL or None (baseline = config.yaml via AIClient).
|
|
|
|
Raises ValueError on unknown slug (caller 가 HTTP 400 으로 translate).
|
|
"""
|
|
if slug is None or slug == "baseline":
|
|
return None
|
|
if slug not in RERANKER_BACKEND_MAP:
|
|
raise ValueError(f"unknown_reranker_backend: {slug!r}")
|
|
cfg = RERANKER_BACKEND_MAP[slug]
|
|
return cfg["endpoint"] if cfg else None
|
|
|
|
|
|
async def _rerank_via_candidate_endpoint(
|
|
endpoint: str, query: str, texts: list[str]
|
|
) -> list[dict]:
|
|
"""후보 TEI reranker endpoint 호출 (cache 미사용).
|
|
|
|
Returns:
|
|
[{"index": int, "score": float}, ...] sorted score desc.
|
|
Raises:
|
|
httpx errors — caller 가 timeout/fallback path 로.
|
|
"""
|
|
async with httpx.AsyncClient(timeout=RERANK_TIMEOUT) as c:
|
|
r = await c.post(endpoint, json={"query": query, "texts": texts})
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
if not isinstance(data, list):
|
|
raise ValueError(f"unexpected candidate TEI shape: {type(data).__name__}")
|
|
return data
|
|
|
|
|
|
def _extract_window(text: str, query: str, target_chars: int = 800) -> str:
|
|
"""query keyword 위치 중심으로 ±target_chars/2 윈도우 추출.
|
|
|
|
fallback: keyword 매치 없으면 첫 target_chars 문자 그대로.
|
|
이게 없으면 reranker가 무관한 텍스트만 보고 점수 매겨 성능 급락.
|
|
"""
|
|
keywords = [k for k in re.split(r"\s+", query) if len(k) >= 2]
|
|
best_pos = -1
|
|
for kw in keywords:
|
|
pos = text.lower().find(kw.lower())
|
|
if pos >= 0:
|
|
best_pos = pos
|
|
break
|
|
|
|
if best_pos < 0:
|
|
# Fallback: 첫 target_chars 문자
|
|
return text[:target_chars]
|
|
|
|
half = target_chars // 2
|
|
start = max(0, best_pos - half)
|
|
end = min(len(text), start + target_chars)
|
|
return text[start:end]
|
|
|
|
|
|
def _make_snippet(c: "SearchResult", query: str, max_chars: int = 1500) -> str:
|
|
"""Reranker input snippet — title + query 중심 본문 윈도우.
|
|
|
|
feedback_search_phase1_implementation.md 3번 항목 강제:
|
|
snippet 200~400 토큰(800~1500자), full document 절대 안 됨.
|
|
"""
|
|
title = c.title or ""
|
|
text = c.snippet or ""
|
|
|
|
# snippet은 chunk text 앞 200자 또는 doc text 앞 200자
|
|
# 더 긴 chunk text가 필요하면 호출자가 따로 채워서 넘김
|
|
if len(text) > max_chars:
|
|
text = _extract_window(text, query, target_chars=max_chars - 100)
|
|
|
|
return f"{title}\n\n{text}"
|
|
|
|
|
|
def _wrap_doc_as_chunk(doc: "SearchResult") -> "SearchResult":
|
|
"""text-only 매치 doc(chunks_by_doc에 없는 doc)을 ChunkResult 형태로 변환.
|
|
|
|
Phase 1.3 reranker 입력에 doc 자체가 들어가야 하는 경우.
|
|
snippet은 documents.extracted_text 앞 200자 (이미 SearchResult.snippet에 채워짐).
|
|
chunk_id 등은 None 그대로.
|
|
"""
|
|
return doc
|
|
|
|
|
|
async def rerank_chunks(
|
|
query: str,
|
|
candidates: list["SearchResult"],
|
|
limit: int,
|
|
*,
|
|
reranker_backend: str | None = None,
|
|
snapshot_doc_id_max: int | None = None,
|
|
snapshot_chunk_id_max: int | None = None,
|
|
) -> list["SearchResult"]:
|
|
"""RRF 결과 candidates를 bge-reranker로 재정렬.
|
|
|
|
Args:
|
|
query: 사용자 쿼리
|
|
candidates: chunk-level SearchResult 리스트 (이미 chunks_by_doc에서 회수)
|
|
limit: 반환할 결과 수
|
|
|
|
Returns:
|
|
reranked SearchResult 리스트 (rerank score로 score 필드 업데이트)
|
|
|
|
Fallback (timeout/HTTPError): RRF 순서 그대로 candidates[:limit] 반환.
|
|
"""
|
|
if not candidates:
|
|
return []
|
|
|
|
# input 크기 제한 (latency/VRAM hard cap)
|
|
if len(candidates) > MAX_RERANK_INPUT:
|
|
logger.warning(
|
|
f"rerank input {len(candidates)} > MAX {MAX_RERANK_INPUT}, 자름"
|
|
)
|
|
candidates = candidates[:MAX_RERANK_INPUT]
|
|
|
|
snippets = [_make_snippet(c, query) for c in candidates]
|
|
|
|
# Phase 2B dispatcher (R2-B1 + R2-B2): slug → endpoint resolve, snapshot id dispatch log
|
|
cand_endpoint = _resolve_reranker(reranker_backend)
|
|
logger.info(
|
|
"[reranker-dispatch] backend=%s endpoint=%s snapshot_doc_id_max=%s snapshot_chunk_id_max=%s",
|
|
reranker_backend or "baseline",
|
|
cand_endpoint or "production(config.yaml)",
|
|
snapshot_doc_id_max,
|
|
snapshot_chunk_id_max,
|
|
)
|
|
|
|
client: AIClient | None = AIClient() if cand_endpoint is None else None
|
|
|
|
try:
|
|
async with asyncio.timeout(RERANK_TIMEOUT):
|
|
async with RERANK_SEMAPHORE:
|
|
if cand_endpoint is None:
|
|
results = await client.rerank(query, snippets)
|
|
else:
|
|
results = await _rerank_via_candidate_endpoint(
|
|
cand_endpoint, query, snippets
|
|
)
|
|
# results: [{"index": int, "score": float}, ...] (이미 정렬됨)
|
|
reranked: list["SearchResult"] = []
|
|
for r in results:
|
|
idx = r.get("index")
|
|
sc = r.get("score")
|
|
if idx is None or sc is None or idx >= len(candidates):
|
|
continue
|
|
chunk = candidates[idx]
|
|
score = float(sc)
|
|
chunk.score = score
|
|
# Phase 3.1: reranker raw 점수를 별도 필드에 보존.
|
|
# normalize_display_scores가 나중에 .score를 랭크 기반으로 덮어써도
|
|
# fast-path 판단에 쓸 수 있는 원본 신호 유지.
|
|
chunk.rerank_score = score
|
|
chunk.match_reason = (chunk.match_reason or "") + "+rerank"
|
|
reranked.append(chunk)
|
|
return reranked[:limit]
|
|
except (asyncio.TimeoutError, httpx.HTTPError) as e:
|
|
logger.warning(f"rerank failed → RRF fallback: {type(e).__name__}: {e}")
|
|
return candidates[:limit]
|
|
except Exception as e:
|
|
logger.warning(f"rerank unexpected error → RRF fallback: {type(e).__name__}: {e}")
|
|
return candidates[:limit]
|
|
finally:
|
|
if client is not None:
|
|
try:
|
|
await client.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def warmup_reranker() -> bool:
|
|
"""TEI 부팅 후 모델 로딩 완료 대기 (10회 retry).
|
|
|
|
TEI는 health 200을 빠르게 반환하지만 첫 모델 로딩(10~30초) 전에는
|
|
rerank 요청이 실패하거나 매우 느림. FastAPI startup 또는 첫 요청 전 호출.
|
|
"""
|
|
client = AIClient()
|
|
try:
|
|
for attempt in range(10):
|
|
try:
|
|
await client.rerank("warmup", ["dummy text for model load"])
|
|
logger.info(f"reranker warmup OK (attempt {attempt + 1})")
|
|
return True
|
|
except Exception as e:
|
|
logger.info(f"reranker warmup retry {attempt + 1}: {e}")
|
|
await asyncio.sleep(3)
|
|
logger.error("reranker warmup failed after 10 attempts")
|
|
return False
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
def apply_diversity(
|
|
results: list["SearchResult"],
|
|
max_per_doc: int = MAX_CHUNKS_PER_DOC,
|
|
top_score_threshold: float = 0.90,
|
|
) -> list["SearchResult"]:
|
|
"""chunk-level 결과를 doc 기준으로 압축 (max_per_doc).
|
|
|
|
조건부 완화: 가장 상위 결과 score가 threshold 이상이면 unlimited
|
|
(high confidence relevance > diversity).
|
|
"""
|
|
if not results:
|
|
return []
|
|
|
|
# 가장 상위 score가 threshold 이상이면 diversity 제약 해제
|
|
top_score = results[0].score if results else 0.0
|
|
if top_score >= top_score_threshold:
|
|
return results
|
|
|
|
seen: dict[int, int] = {}
|
|
out: list["SearchResult"] = []
|
|
for r in results:
|
|
doc_id = r.id
|
|
if seen.get(doc_id, 0) >= max_per_doc:
|
|
continue
|
|
out.append(r)
|
|
seen[doc_id] = seen.get(doc_id, 0) + 1
|
|
return out
|