리랭킹

# RAG 파이프라인에서 리랭킹 적용
from sentence_transformers import CrossEncoder
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings

# 1단계: 벡터 검색으로 후보 추출
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(embedding_function=embeddings)

query = "LLM의 할루시네이션을 줄이는 방법"
initial_docs = vectorstore.similarity_search(query, k=20)

# 2단계: Cross-Encoder로 리랭킹
reranker = CrossEncoder('BAAI/bge-reranker-large')

# 쿼리-문서 쌍 생성
pairs = [[query, doc.page_content] for doc in initial_docs]

# 관련성 점수 계산
scores = reranker.predict(pairs)

# 점수로 재정렬
reranked = sorted(
    zip(initial_docs, scores),
    key=lambda x: x[1],
    reverse=True
)

# 상위 5개 선택
top_docs = [doc for doc, score in reranked[:5]]

print("🔝 리랭킹 후 상위 문서:")
for i, (doc, score) in enumerate(reranked[:5]):
    print(f"{i+1}. [점수: {score:.4f}] {doc.page_content[:80]}...")

# Cohere API 사용 예시
import cohere
co = cohere.Client("your-api-key")

results = co.rerank(
    query=query,
    documents=[doc.page_content for doc in initial_docs],
    top_n=5,
    model="rerank-v3.0"
)

for r in results.results:
    print(f"순위 {r.index}: 점수 {r.relevance_score:.4f}")