8장: 하이브리드 검색과 리랭킹 파이프라인

def rrf_fusion(
    results_list: list[list[tuple[str, float]]],
    k: int = 60,
    top_n: int = 10,
) -> list[tuple[str, float]]:
    """
    RRF 점수 퓨전
    results_list: 여러 검색 결과 [(doc_id, score), ...]의 리스트
    """
    rrf_scores: dict[str, float] = {}
 
    for results in results_list:
        for rank, (doc_id, _) in enumerate(results, start=1):
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + 1.0 / (k + rank)
 
    sorted_results = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_results[:top_n]

def min_max_normalize(scores: list[float]) -> list[float]:
    """Min-Max 정규화: 0-1 범위로 변환"""
    min_s = min(scores) if scores else 0
    max_s = max(scores) if scores else 1
    if max_s == min_s:
        return [0.5] * len(scores)
    return [(s - min_s) / (max_s - min_s) for s in scores]
 
def linear_interpolation_fusion(
    bm25_results: list[tuple[str, float]],
    semantic_results: list[tuple[str, float]],
    alpha: float = 0.5,
    top_n: int = 10,
) -> list[tuple[str, float]]:
    """
    선형 보간 퓨전
    alpha: 시맨틱 검색 가중치 (1-alpha가 BM25 가중치)
    """
    # 점수 정규화
    bm25_ids = [doc_id for doc_id, _ in bm25_results]
    bm25_scores = min_max_normalize([s for _, s in bm25_results])
    bm25_map = dict(zip(bm25_ids, bm25_scores))
 
    sem_ids = [doc_id for doc_id, _ in semantic_results]
    sem_scores = min_max_normalize([s for _, s in semantic_results])
    sem_map = dict(zip(sem_ids, sem_scores))
 
    # 모든 문서에 대해 결합 점수 계산
    all_docs = set(bm25_ids) | set(sem_ids)
    combined = {}
    for doc_id in all_docs:
        bm25_s = bm25_map.get(doc_id, 0.0)
        sem_s = sem_map.get(doc_id, 0.0)
        combined[doc_id] = (1 - alpha) * bm25_s + alpha * sem_s
 
    sorted_results = sorted(combined.items(), key=lambda x: x[1], reverse=True)
    return sorted_results[:top_n]

from sentence_transformers import CrossEncoder
from dataclasses import dataclass
 
@dataclass
class CascadeConfig:
    stage_name: str
    model_name: str
    input_size: int
    output_size: int
 
class RankingCascade:
    """다단계 리랭킹 캐스케이드"""
 
    def __init__(self, stages: list[CascadeConfig]):
        self.stages = stages
        self.models = {}
        for stage in stages:
            self.models[stage.stage_name] = CrossEncoder(stage.model_name)
 
    def rerank(self, query: str, candidates: list[dict]) -> list[dict]:
        """캐스케이드 리랭킹 실행"""
        current_candidates = candidates
 
        for stage in self.stages:
            model = self.models[stage.stage_name]
            top_candidates = current_candidates[:stage.input_size]
 
            pairs = [[query, c["text"]] for c in top_candidates]
            scores = model.predict(pairs)
 
            for candidate, score in zip(top_candidates, scores):
                candidate[f"{stage.stage_name}_score"] = float(score)
 
            top_candidates.sort(
                key=lambda x: x[f"{stage.stage_name}_score"],
                reverse=True,
            )
            current_candidates = top_candidates[:stage.output_size]
 
        return current_candidates
 
# 사용 예시
cascade = RankingCascade([
    CascadeConfig(
        stage_name="light_rerank",
        model_name="cross-encoder/ms-marco-MiniLM-L-6-v2",
        input_size=200,
        output_size=50,
    ),
    CascadeConfig(
        stage_name="heavy_rerank",
        model_name="cross-encoder/ms-marco-MiniLM-L-12-v2",
        input_size=50,
        output_size=10,
    ),
])

import time
from contextlib import contextmanager
from dataclasses import dataclass, field
 
@dataclass
class LatencyProfile:
    stages: dict[str, float] = field(default_factory=dict)
 
    @contextmanager
    def measure(self, stage_name: str):
        start = time.perf_counter()
        yield
        elapsed = time.perf_counter() - start
        self.stages[stage_name] = elapsed * 1000  # ms
 
    def total_ms(self) -> float:
        return sum(self.stages.values())
 
    def report(self) -> str:
        lines = [f"Total: {self.total_ms():.1f}ms"]
        for stage, ms in self.stages.items():
            pct = ms / self.total_ms() * 100
            lines.append(f"  {stage}: {ms:.1f}ms ({pct:.0f}%)")
        return "\n".join(lines)
 
# 사용 예시
profile = LatencyProfile()
 
with profile.measure("bm25_retrieval"):
    bm25_results = bm25_search(query, k=500)
 
with profile.measure("knn_retrieval"):
    knn_results = knn_search(query, k=500)
 
with profile.measure("rrf_fusion"):
    fused = rrf_fusion([bm25_results, knn_results], top_n=100)
 
with profile.measure("cross_encoder_rerank"):
    reranked = cross_encoder_rerank(query, fused, top_n=10)
 
print(profile.report())
# Total: 285.3ms
#   bm25_retrieval: 12.5ms (4%)
#   knn_retrieval: 18.2ms (6%)
#   rrf_fusion: 1.6ms (1%)
#   cross_encoder_rerank: 253.0ms (89%)

import torch
 
def batch_rerank(model, query: str, documents: list[str], batch_size: int = 32):
    """배치 단위 Cross-encoder 리랭킹"""
    all_scores = []
    pairs = [[query, doc] for doc in documents]
 
    for i in range(0, len(pairs), batch_size):
        batch = pairs[i:i + batch_size]
        with torch.no_grad():
            scores = model.predict(batch, show_progress_bar=False)
        all_scores.extend(scores)
 
    return all_scores

import hashlib
from functools import lru_cache
 
def cache_key(query: str, doc_ids: list[str]) -> str:
    content = f"{query}|{'|'.join(doc_ids)}"
    return hashlib.md5(content.encode()).hexdigest()
 
# 간단한 인메모리 캐시 (프로덕션에서는 Redis 등 사용)
rerank_cache: dict[str, list[float]] = {}
 
def cached_rerank(model, query: str, documents: list[dict]) -> list[float]:
    doc_ids = [d["id"] for d in documents]
    key = cache_key(query, doc_ids)
 
    if key in rerank_cache:
        return rerank_cache[key]
 
    scores = model.predict([[query, d["text"]] for d in documents])
    rerank_cache[key] = scores.tolist()
    return scores.tolist()