feat: 重排处理

This commit is contained in:
李如威 2025-07-29 09:49:53 +08:00
parent b2fcd4a69c
commit 5f3fc6ff07
1 changed files with 7 additions and 11 deletions

View File

@ -48,13 +48,9 @@ def main():
"type": "local",
"model_name": "sentence-transformers/all-MiniLM-L6-v2",
}
# 重排配置
rerank_config = {
"enabled": True,
"method": "similarity",
"top_k": 3
}
rerank_config = {"enabled": True, "method": "cross_encoder", "top_k": 3}
# 初始化RAG系统
rag = RerankRAG(
@ -111,21 +107,21 @@ def main():
for i, test_case in enumerate(test_queries, 1):
query = test_case["query"]
expected = test_case["expected_topic"]
print(f"=== 测试查询 {i}: {query} ===")
print(f"期望结果: {expected}\n")
# 获取两种搜索结果
docs_no_rerank, docs_with_rerank = rag.query_with_scores(query, k=5)
print("📍 不使用重排的结果:")
for j, doc in enumerate(docs_no_rerank[:3], 1):
print(f" {j}. {doc.page_content}")
print("\n🎯 使用重排的结果:")
for j, doc in enumerate(docs_with_rerank[:3], 1):
print(f" {j}. {doc.page_content}")
print("\n" + "="*80 + "\n")
print("重排功能测试完成!")