Cohere와 BGE를 활용한 교차 인코더 재순위 지정

# Bi-encoder: separate encoding
query_vec = encoder.encode(query)          # [768] vector
doc_vec = encoder.encode(document)         # [768] vector
score = cosine_similarity(query_vec, doc_vec)  # compare independently

# Cross-encoder: joint encoding
combined_input = '[CLS] ' + query + ' [SEP] ' + document + ' [SEP]'
logits = cross_encoder.forward(combined_input)  # score from joint attention
# The model attends from every query token to every document token

import cohere

co = cohere.Client('YOUR_COHERE_API_KEY')

candidates = [
    'How to configure pgvector in PostgreSQL for vector search',
    'BM25 scoring formula and hyperparameters explained',
    'Installing and using the pgvector extension for embeddings',
    'Hybrid search combining BM25 and dense retrieval',
]

result = co.rerank(
    model='rerank-english-v3.0',
    query='pgvector setup guide',
    documents=candidates,
    top_n=3,
    return_documents=True,
)

for item in result.results:
    print(f'Score {item.relevance_score:.4f}: {item.document.text[:60]}')

import cohere
from openai import OpenAI

co = cohere.Client('COHERE_KEY')
client = OpenAI()

def rag_with_rerank(question: str, retriever, top_k: int = 5) -> str:
    # Stage 1: coarse retrieval
    candidates = retriever.invoke(question)  # returns 30 docs
    texts = [doc.page_content for doc in candidates]

    # Stage 2: Cohere re-ranking
    rerank_result = co.rerank(
        model='rerank-english-v3.0',
        query=question,
        documents=texts,
        top_n=top_k,
    )
    top_texts = [texts[r.index] for r in rerank_result.results]
    context = '\n\n'.join(top_texts)

    # Stage 3: generation
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[
            {'role': 'system', 'content': 'Answer using only the provided context.'},
            {'role': 'user', 'content': f'Context:\n{context}\n\nQuestion: {question}'},
        ],
    )
    return response.choices[0].message.content

from sentence_transformers import CrossEncoder

# Load BGE reranker model
reranker = CrossEncoder(
    'BAAI/bge-reranker-v2-m3',
    max_length=512,
    device='cpu',  # use 'cuda' if GPU is available
)

def bge_rerank(query: str, documents: list[str], top_k: int = 5) -> list[dict]:
    pairs = [[query, doc] for doc in documents]
    scores = reranker.predict(pairs, show_progress_bar=False)
    ranked = sorted(
        zip(documents, scores.tolist()),
        key=lambda x: x[1],
        reverse=True,
    )
    return [{'text': doc, 'score': score} for doc, score in ranked[:top_k]]

# Model size vs accuracy trade-offs
models = [
    ('BAAI/bge-reranker-base',   '278M params', 'fast,   English-focused'),
    ('BAAI/bge-reranker-large',  '560M params', 'better, English-focused'),
    ('BAAI/bge-reranker-v2-m3',  '570M params', 'best for multilingual use'),
    ('BAAI/bge-reranker-v2-gemma', '2B params', 'SOTA open-source accuracy'),
]

# For most RAG applications, bge-reranker-v2-m3 offers the best
# balance of accuracy, multilingual support, and inference speed on CPU

def batch_rerank(reranker, query: str, documents: list[str],
                 top_k: int = 5, batch_size: int = 32) -> list[dict]:
    pairs = [[query, doc] for doc in documents]

    # Batch prediction
    all_scores = reranker.predict(
        pairs,
        batch_size=batch_size,
        show_progress_bar=len(pairs) > 100,
    )

    ranked = sorted(
        zip(documents, all_scores.tolist()),
        key=lambda x: x[1],
        reverse=True,
    )
    return [{'text': doc, 'score': score} for doc, score in ranked[:top_k]]

from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

# Base retriever (coarse stage)
base_retriever = FAISS.from_documents(
    documents, OpenAIEmbeddings()
).as_retriever(search_kwargs={'k': 30})

# Re-ranker (fine stage)
model = HuggingFaceCrossEncoder(model_name='BAAI/bge-reranker-v2-m3')
compressor = CrossEncoderReranker(model=model, top_n=5)

# Combined two-stage retriever
two_stage_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever,
)

docs = two_stage_retriever.invoke('how does hybrid search work?')

def rerank_with_threshold(reranker, query, documents, top_k=5, min_score=-2.0):
    pairs = [[query, doc] for doc in documents]
    scores = reranker.predict(pairs)

    scored_docs = [
        {'text': doc, 'score': float(score)}
        for doc, score in zip(documents, scores)
        if float(score) >= min_score  # filter low-relevance docs
    ]

    scored_docs.sort(key=lambda x: x['score'], reverse=True)

    if not scored_docs:
        return []  # nothing passed the threshold
    return scored_docs[:top_k]

from sentence_transformers import CrossEncoder

# Multilingual BGE reranker
ml_reranker = CrossEncoder('BAAI/bge-reranker-v2-m3')

# Cross-lingual example: English query, French document
query_en = 'How to configure vector search?'
doc_fr = 'La configuration de la recherche vectorielle dans PostgreSQL'

score = ml_reranker.predict([[query_en, doc_fr]])
print(f'Cross-lingual relevance score: {float(score):.3f}')
# Positive score indicates the document is relevant to the query
# despite being in a different language

import redis
import hashlib
import json

r = redis.Redis(host='localhost', port=6379)

def cached_rerank_score(reranker, query: str, doc_text: str, doc_id: str) -> float:
    cache_key = 'rerank:' + hashlib.md5((query + doc_id).encode()).hexdigest()
    cached = r.get(cache_key)
    if cached:
        return float(cached)

    score = float(reranker.predict([[query, doc_text]]))
    r.setex(cache_key, 3600, str(score))  # cache 1 hour
    return score