Reranking 전략

[쿼리] → Retriever (Top-K 후보) → Reranker (재정렬) → 최종 Top-N

from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank

# Cohere Reranker 설정
reranker = CohereRerank(
    model="rerank-v3.5",
    top_n=5  # 최종 반환 문서 수
)

# 기존 Retriever + Reranker 결합
compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=ensemble_retriever  # 앙상블 Retriever에서 후보 추출
)

results = compression_retriever.invoke("Unity Catalog 권한 설정 방법")

from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.retrievers.document_compressors import CrossEncoderReranker

# BGE Reranker (오픈소스, 로컬 실행 가능)
cross_encoder = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
reranker = CrossEncoderReranker(model=cross_encoder, top_n=5)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=ensemble_retriever
)

# ColBERT는 토큰 단위 벡터를 미리 계산 (인덱싱 시)
# 검색 시에는 MaxSim (Maximum Similarity) 연산으로 빠르게 비교

# RAGatouille 라이브러리를 통한 ColBERT 활용
from ragatouille import RAGPretrainedModel

# ColBERT 모델 로드
colbert = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

# 문서 인덱싱 (토큰 단위 벡터 사전 계산)
colbert.index(
    collection=[doc.page_content for doc in documents],
    index_name="korean_docs_colbert",
)

# 검색 (MaxSim 연산으로 빠른 비교)
results = colbert.search(query="Unity Catalog 권한 설정", k=5)

def llm_rerank(query: str, documents: list, top_n: int = 5) -> list:
    """LLM이 직접 문서 순위를 정렬 (RankGPT 방식)"""
    doc_list = "\n".join([
        f"[{i+1}] {doc.page_content[:300]}"
        for i, doc in enumerate(documents)
    ])

    prompt = f"""다음 질문에 가장 관련성이 높은 문서 순서대로 번호를 나열하세요.
상위 {top_n}개만 선택하세요.

질문: {query}

문서 목록:
{doc_list}

관련성 높은 순서 (번호만): """

    response = llm.invoke(prompt)
    # 응답에서 번호 추출하여 재정렬
    ranked_indices = [int(x.strip()) - 1 for x in response.content.split(",")[:top_n]]
    return [documents[i] for i in ranked_indices if i < len(documents)]

# 실전 Reranking 파이프라인: 50개 후보 → 5개 최종 선택
from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.retrievers.document_compressors import CrossEncoderReranker

# 1단계: 넓은 범위로 후보 검색
base_retriever = ensemble_retriever  # 또는 vs_retriever
base_retriever.search_kwargs = {"k": 50}

# 2단계: Cross-encoder로 정밀 재정렬
cross_encoder = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
reranker = CrossEncoderReranker(model=cross_encoder, top_n=5)

# 결합
final_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=base_retriever
)