475_信息检索与排序¶

# ---- 1. BM25 关键词检索 ----
# pip install rank-bm25 jieba
from rank_bm25 import BM25Okapi
import jieba

# 文档库（实际场景可能有百万文档）
corpus = [
    "深度学习在图像识别领域取得突破性进展",
    "自然语言处理BERT模型的预训练方法",
    "推荐系统协同过滤算法原理",
    "机器学习随机森林模型解析",
    "图神经网络在知识图谱中的应用",
]

# 中文分词（BM25需要分词后的token列表）
tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus]
bm25 = BM25Okapi(tokenized_corpus)

# 查询
query = "深度学习图像"
tokenized_query = list(jieba.cut(query))
scores = bm25.get_scores(tokenized_query)
top_n = bm25.get_top_n(tokenized_query, corpus, n=3)

print("BM25 Top-3结果:")
for i, doc in enumerate(top_n):
    print(f"  {i+1}. {doc}")

# ---- 2. 密集向量检索（DPR风格）----
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss  # pip install faiss-cpu

# 用Sentence-BERT编码文档库
model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
doc_embeddings = model.encode(corpus, show_progress_bar=True)
# doc_embeddings.shape: (5, 384)

# 构建FAISS索引（近似最近邻检索，百万级文档秒级返回）
dimension = doc_embeddings.shape[1]  # 384
index = faiss.IndexFlatIP(dimension)  # IP=内积（等效于L2归一化后的余弦相似度）

# L2归一化，使内积等于余弦相似度
faiss.normalize_L2(doc_embeddings)
index.add(doc_embeddings.astype(np.float32))  # 添加向量到索引

# 查询
query_emb = model.encode([query])
faiss.normalize_L2(query_emb)
distances, indices = index.search(query_emb.astype(np.float32), k=3)

print("\nFAISS 向量检索 Top-3:")
for i, (dist, idx) in enumerate(zip(distances[0], indices[0])):
    print(f"  {i+1}. (相似度={dist:.4f}) {corpus[idx]}")

# ---- 3. 混合检索（BM25 + 向量，取分数融合）----
def hybrid_search(query, corpus, bm25, model, alpha=0.5, top_k=3):
    """
    alpha: BM25权重；(1-alpha): 向量检索权重
    """
    # BM25分数
    tokenized_q = list(jieba.cut(query))
    bm25_scores = np.array(bm25.get_scores(tokenized_q))
    bm25_scores = bm25_scores / (bm25_scores.max() + 1e-9)  # 归一化

    # 向量相似度分数
    q_emb = model.encode([query])
    faiss.normalize_L2(q_emb)
    dense_scores, _ = index.search(q_emb.astype(np.float32), k=len(corpus))
    # FAISS返回的是排序后的，需要还原顺序
    vec_scores = np.zeros(len(corpus))
    for score, idx in zip(dense_scores[0], _[0]):
        if idx < len(corpus):
            vec_scores[idx] = score

    # 加权融合
    final_scores = alpha * bm25_scores + (1 - alpha) * vec_scores
    top_indices = np.argsort(final_scores)[::-1][:top_k]
    return [(corpus[i], final_scores[i]) for i in top_indices]

results = hybrid_search("深度学习图像识别", corpus, bm25, model)
print("\n混合检索结果:")
for doc, score in results:
    print(f"  {score:.4f}: {doc}")

# ---- 4. NDCG评估 ----
from sklearn.metrics import ndcg_score

# 真实相关性标签（0=不相关,1=相关,2=非常相关）
true_relevance = np.array([[2, 1, 0, 0, 0]])
predicted_scores = np.array([[0.9, 0.7, 0.3, 0.2, 0.1]])

ndcg = ndcg_score(true_relevance, predicted_scores, k=3)
print(f"\nNDCG@3: {ndcg:.4f}")