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