test_siliconflow_reranker.py•4.57 kB
#!/usr/bin/env python3
"""
测试SiliconFlow重排器
"""
import sys
from pathlib import Path
# 添加src路径
sys.path.insert(0, str(Path(__file__).parent / "src"))
from memos.configs.reranker import RerankerConfigFactory
from memos.rerankers.factory import RerankerFactory
from usage_examples import load_env_file
def test_siliconflow_reranker():
"""测试SiliconFlow重排器"""
print("🧪 测试SiliconFlow重排器")
print("=" * 50)
# 加载环境变量
load_env_file()
# 创建SiliconFlow重排器配置
config = RerankerConfigFactory.model_validate({
"backend": "siliconflow",
"config": {
"model_name_or_path": "Qwen/Qwen3-Reranker-0.6B",
"model_name": "Qwen/Qwen3-Reranker-0.6B",
"api_key": "sk-ygqlrgrxrypykiiskuspuahkwihhbhhjhazqokntwdzfwqdv",
"api_base": "https://api.siliconflow.cn/v1",
"top_k": 5,
"max_chunks_per_query": 50
}
})
print("✅ 配置创建成功")
# 创建重排器
reranker = RerankerFactory.from_config(config)
print("✅ 重排器创建成功")
# 测试数据
query = "数据分析平台的技术架构设计"
documents = [
"数据分析平台需要考虑数据采集、存储、计算、可视化四个核心层次",
"推荐使用ClickHouse作为OLAP数据库，Spark作为计算引擎",
"MemOS是一个智能记忆管理系统，支持向量搜索和语义检索",
"Claude可以通过MCP调用MemOS获取记忆上下文",
"SiliconFlow提供高质量的AI API服务，包括嵌入和重排功能",
"机器学习模型的训练需要大量的计算资源和数据",
"云原生架构采用容器化部署和微服务设计模式"
]
print(f"🔄 正在重排 {len(documents)} 个文档...")
print(f"查询: {query}")
try:
# 测试重排功能
reranked_results = reranker.rerank(query, documents, top_k=5)
print("✅ 重排成功！")
print(f"📊 重排统计:")
print(f" 原始文档数量: {len(documents)}")
print(f" 重排后返回数量: {len(reranked_results)}")
# 显示重排结果
print(f"\n📝 重排结果 (按相关度排序):")
for i, (doc, score) in enumerate(reranked_results, 1):
print(f"\n{i}. 相关度分数: {score:.4f}")
print(f" 文档内容: {doc}")
# 测试分数获取功能
print(f"\n🔄 测试分数获取功能...")
scores = reranker.get_scores(query, documents[:3]) # 只测试前3个文档
print("✅ 分数获取成功！")
print(f"📊 分数结果:")
for i, (doc, score) in enumerate(zip(documents[:3], scores), 1):
print(f"\n{i}. 分数: {score:.4f}")
print(f" 文档: {doc[:50]}...")
return True
except Exception as e:
print(f"❌ 重排测试失败: {e}")
print(f"错误类型: {type(e).__name__}")
# 尝试直接测试API连接
print(f"\n🔄 测试API连接...")
try:
import openai
client = openai.OpenAI(
api_key="sk-ygqlrgrxrypykiiskuspuahkwihhbhhjhazqokntwdzfwqdv",
base_url="https://api.siliconflow.cn/v1"
)
# 测试模型列表
models = client.models.list()
print("✅ API连接正常")
print(f"可用模型数量: {len(models.data)}")
# 查找重排模型
rerank_models = [m for m in models.data if "rerank" in m.id.lower() or "bge" in m.id.lower()]
if rerank_models:
print("🎯 找到重排模型:")
for model in rerank_models[:5]: # 只显示前5个
print(f" - {model.id}")
else:
print("⚠️ 未找到重排模型，可能需要使用不同的API端点")
except Exception as api_e:
print(f"❌ API连接失败: {api_e}")
return False
if __name__ == "__main__":
success = test_siliconflow_reranker()
if success:
print("\n🎉 SiliconFlow重排器测试成功！")
print("现在可以在MemOS中使用SiliconFlow重排功能了。")
else:
print("\n❌ SiliconFlow重排器测试失败")
print("可能需要调整API调用格式或检查模型可用性")
sys.exit(1)