不止于相似度：混合搜索如何重塑 RAG 的未来

# 安装必要的库
# pip install rank-bm25 sentence-transformers scikit-learn

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# 1. 准备数据
documents = [
    "Apple Inc. announced the new M3 chip, focusing on performance and efficiency.",
    "The latest financial report from Apple Inc. shows strong growth in the services sector.",
    "A detailed review of the MacBook Pro with M3 chip highlights its impressive speed.",
    "Google's new Pixel phone features an advanced AI-powered camera.",
    "How to bake the perfect apple pie from scratch.",
    "Microsoft's Surface Laptop competes directly with Apple's MacBook Air."
]

# 2. 创建双索引
# 稀疏索引 (BM25)
tokenized_docs = [doc.lower().split() for doc in documents]
bm25 = BM25Okapi(tokenized_docs)

# 稠密索引 (Sentence Transformer)
model = SentenceTransformer('all-MiniLM-L6-v2')
doc_embeddings = model.encode(documents)

# 3. 执行搜索
query = "latest review of Apple's M3 chip"
tokenized_query = query.lower().split()

# BM25 搜索
bm25_scores = bm25.get_scores(tokenized_query)

# 向量搜索
query_embedding = model.encode(query)
cosine_scores = cosine_similarity([query_embedding], doc_embeddings)[0]

# 4. 实现 RRF 融合
def reciprocal_rank_fusion(search_results_list, k=60):
    fused_scores = {}
    for doc_scores in search_results_list:
        # 对每个搜索结果列表按分数降序排序，获取排名
        sorted_docs = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)
        for rank, (doc_index, score) in enumerate(sorted_docs):
            if doc_index not in fused_scores:
                fused_scores[doc_index] = 0
            fused_scores[doc_index] += 1 / (k + rank + 1) # rank 从 0 开始，所以 +1
    
    # 按 RRF 分数重新排序
    reranked_results = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True)
    return reranked_results

# 准备 RRF 输入
bm25_results = {i: score for i, score in enumerate(bm25_scores)}
vector_results = {i: score for i, score in enumerate(cosine_scores)}

fused_results = reciprocal_rank_fusion([bm25_results, vector_results])

# 5. 展示结果
print("--- Query ---")
print(query)
print("
--- BM25 (Keyword) Search Results ---")
for i, score in sorted(bm25_results.items(), key=lambda item: item[1], reverse=True):
    print(f"Score: {score:.4f} Doc: {documents[i]}")

print("
--- Vector (Semantic) Search Results ---")
for i, score in sorted(vector_results.items(), key=lambda item: item[1], reverse=True):
    print(f"Score: {score:.4f} Doc: {documents[i]}")

print("
--- Hybrid Search (RRF Fused) Results ---")
for doc_index, score in fused_results:
    print(f"Score: {score:.4f} Doc: {documents[doc_index]}")