5장: 랭킹 모델 — Bi-encoder와 Cross-encoder

from sentence_transformers import CrossEncoder
 
# Cross-encoder 모델 로드
model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
 
query = "파이썬 비동기 프로그래밍"
documents = [
    "Python asyncio를 활용한 비동기 웹 서버 구축 가이드",
    "JavaScript Promise와 async/await 패턴",
    "파이썬 동시성: threading vs multiprocessing 비교",
    "Node.js 이벤트 루프의 작동 원리",
]
 
# 쿼리-문서 쌍의 관련성 점수 계산
pairs = [[query, doc] for doc in documents]
scores = model.predict(pairs)
 
for doc, score in sorted(zip(documents, scores), key=lambda x: x[1], reverse=True):
    print(f"{score:.4f} | {doc}")

from sentence_transformers import SentenceTransformer, CrossEncoder
import numpy as np
 
class TwoStageSearchPipeline:
    """2단계 검색 파이프라인: Bi-encoder 리트리벌 + Cross-encoder 리랭킹"""
 
    def __init__(
        self,
        retriever_model: str = "intfloat/multilingual-e5-large",
        reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-12-v2",
    ):
        self.retriever = SentenceTransformer(retriever_model)
        self.reranker = CrossEncoder(reranker_model)
        self.doc_embeddings = None
        self.documents = []
 
    def index(self, documents: list[str]):
        """문서 인덱싱 (Bi-encoder 벡터 생성)"""
        self.documents = documents
        texts = [f"passage: {doc}" for doc in documents]
        self.doc_embeddings = self.retriever.encode(texts, normalize_embeddings=True)
 
    def search(self, query: str, top_k: int = 10, rerank_k: int = 100) -> list[dict]:
        """2단계 검색: 리트리벌 후 리랭킹"""
        # 1단계: Bi-encoder 리트리벌
        query_vec = self.retriever.encode(
            f"query: {query}", normalize_embeddings=True
        )
        similarities = self.doc_embeddings @ query_vec
        top_indices = np.argsort(similarities)[::-1][:rerank_k]
        candidates = [(i, self.documents[i]) for i in top_indices]
 
        # 2단계: Cross-encoder 리랭킹
        pairs = [[query, doc] for _, doc in candidates]
        rerank_scores = self.reranker.predict(pairs)
 
        # 최종 결과 정렬
        results = []
        for (idx, doc), score in zip(candidates, rerank_scores):
            results.append({
                "doc_id": idx,
                "text": doc,
                "retrieval_score": float(similarities[idx]),
                "rerank_score": float(score),
            })
 
        results.sort(key=lambda x: x["rerank_score"], reverse=True)
        return results[:top_k]

{
  "service": "elasticsearch",
  "service_settings": {
    "num_allocations": 1,
    "num_threads": 1,
    "model_id": ".rerank-v1"
  }
}

{
  "retriever": {
    "text_similarity_reranker": {
      "retriever": {
        "standard": {
          "query": {
            "match": {
              "content": "벡터 데이터베이스 성능 비교"
            }
          }
        }
      },
      "field": "content",
      "inference_id": "my-rerank-model",
      "inference_text": "벡터 데이터베이스 성능 비교",
      "rank_window_size": 100
    }
  },
  "size": 10
}

def reciprocal_rank_fusion(
    ranked_lists: list[list[str]],
    k: int = 60,
) -> list[tuple[str, float]]:
    """
    RRF 점수 퓨전
    ranked_lists: 여러 검색 결과 목록 (각 목록은 문서 ID 순서)
    k: 스무딩 상수 (기본값 60)
    """
    rrf_scores: dict[str, float] = {}
 
    for ranked_list in ranked_lists:
        for rank, doc_id in enumerate(ranked_list, start=1):
            if doc_id not in rrf_scores:
                rrf_scores[doc_id] = 0.0
            rrf_scores[doc_id] += 1.0 / (k + rank)
 
    # 점수 기준 내림차순 정렬
    sorted_results = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_results
 
# BM25 결과와 시맨틱 검색 결과 병합
bm25_results = ["doc_A", "doc_B", "doc_C", "doc_D", "doc_E"]
semantic_results = ["doc_C", "doc_A", "doc_F", "doc_B", "doc_G"]
 
fused = reciprocal_rank_fusion([bm25_results, semantic_results])
for doc_id, score in fused[:5]:
    print(f"{doc_id}: {score:.6f}")

from dataclasses import dataclass
 
@dataclass
class SearchFeatures:
    """LTR 모델의 입력 특성"""
    # 텍스트 매칭 특성
    bm25_score: float
    semantic_similarity: float
    title_match_score: float
 
    # 문서 품질 특성
    doc_length: int
    doc_freshness_days: int
    doc_pagerank: float
 
    # 사용자 행동 특성
    historical_ctr: float
    avg_dwell_time: float
 
    # 쿼리-문서 상호작용 특성
    query_doc_topic_overlap: float
    entity_match_count: int
 
def extract_features(query: str, doc: dict, user_profile: dict) -> SearchFeatures:
    """쿼리-문서 쌍의 LTR 특성 추출"""
    return SearchFeatures(
        bm25_score=compute_bm25(query, doc["text"]),
        semantic_similarity=compute_similarity(query, doc["embedding"]),
        title_match_score=compute_title_match(query, doc["title"]),
        doc_length=len(doc["text"]),
        doc_freshness_days=(today() - doc["published_date"]).days,
        doc_pagerank=doc.get("pagerank", 0.0),
        historical_ctr=doc.get("ctr", 0.0),
        avg_dwell_time=doc.get("avg_dwell_time", 0.0),
        query_doc_topic_overlap=compute_topic_overlap(query, doc),
        entity_match_count=count_entity_matches(query, doc),
    )