两阶段检索为何有效

# Bi-encoder: compute query embedding ONCE, compare to all doc embeddings
# O(1) query encoding + O(n) dot products via ANN index = fast
query_vec = embed(query)  # done once
results = vector_index.search(query_vec, top_k=100)  # fast ANN search

# Cross-encoder: re-scores (query, doc) pairs jointly
# O(k) forward passes for k candidate documents = slow but accurate
for doc in results[:100]:
    score = cross_encoder.score(query, doc.text)  # joint scoring

from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

# Stage 1: retrieve 100 candidates (high recall, modest precision)
vectorstore = FAISS.from_documents(documents, OpenAIEmbeddings())
coarse_retriever = vectorstore.as_retriever(
    search_kwargs={'k': 100}  # large candidate set
)

candidates = coarse_retriever.invoke(query)
print(f'Stage 1: retrieved {len(candidates)} candidate documents')

from sentence_transformers import CrossEncoder

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank(query: str, candidates: list[str], top_k: int = 5) -> list[str]:
    # Score each (query, document) pair jointly
    pairs = [[query, doc] for doc in candidates]
    scores = reranker.predict(pairs)

    # Sort by score descending
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, _ in ranked[:top_k]]

candidate_texts = [doc.page_content for doc in candidates]
final_docs = rerank(query, candidate_texts, top_k=5)
print(f'Stage 2: selected top {len(final_docs)} documents after re-ranking')

import time

def two_stage_search(query, coarse_retriever, reranker, top_k=5):
    t0 = time.perf_counter()

    candidates = coarse_retriever.invoke(query)        # stage 1
    t1 = time.perf_counter()

    candidate_texts = [c.page_content for c in candidates]
    final_docs = rerank(query, candidate_texts, top_k)  # stage 2
    t2 = time.perf_counter()

    print(f'Stage 1 (retrieval): {(t1-t0)*1000:.1f}ms')
    print(f'Stage 2 (re-ranking): {(t2-t1)*1000:.1f}ms')
    print(f'Total: {(t2-t0)*1000:.1f}ms')
    return final_docs

def recall_at_n(coarse_retriever, test_queries, golden_relevant, n_values):
    for n in n_values:
        recalls = []
        for query, relevant in zip(test_queries, golden_relevant):
            # Temporarily set k to n
            coarse_retriever.search_kwargs['k'] = n
            results = coarse_retriever.invoke(query)
            retrieved_ids = {r.metadata.get('id') for r in results}
            relevant_found = len(set(relevant) & retrieved_ids)
            recalls.append(relevant_found / len(relevant))
        avg = sum(recalls) / len(recalls)
        print(f'N={n}: recall={avg:.3f}')

from langchain.retrievers import EnsembleRetriever

# Stage 1: hybrid retrieval for maximum recall
hybrid_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6],
)

# Stage 2: cross-encoder re-ranking for high precision
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

cross_encoder_model = HuggingFaceCrossEncoder(model_name='cross-encoder/ms-marco-MiniLM-L-6-v2')
compressor = CrossEncoderReranker(model=cross_encoder_model, top_n=5)

from langchain.retrievers import ContextualCompressionRetriever
two_stage = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=hybrid_retriever,
)

import cohere

co = cohere.Client('YOUR_API_KEY')

def cohere_rerank(query: str, documents: list[str], top_k: int = 5):
    response = co.rerank(
        model='rerank-english-v3.0',
        query=query,
        documents=documents,
        top_n=top_k,
    )
    return [
        {'text': documents[r.index], 'score': r.relevance_score}
        for r in response.results
    ]

final = cohere_rerank(query, candidate_texts, top_k=5)
for doc in final:
    print(f'Score {doc["score"]:.3f}: {doc["text"][:80]}')