qwen_embedding_config.py•3.63 kB
#!/usr/bin/env python3
"""
Qwen嵌入模型配置
使用Qwen/Qwen3-Embedding-0.6B和Qwen/Qwen3-Reranker-0.6B
"""
import os
import sys
from pathlib import Path
# 添加src路径
sys.path.insert(0, str(Path(__file__).parent / "src"))
from memos.configs.embedder import EmbedderConfigFactory
from memos.configs.reranker import RerankerConfigFactory
from memos.embedders.factory import EmbedderFactory
from memos.rerankers.factory import RerankerFactory
def create_qwen_embedder_config():
"""创建Qwen嵌入器配置"""
config = EmbedderConfigFactory.model_validate({
"backend": "siliconflow",
"config": {
"model_name_or_path": "Qwen/Qwen3-Embedding-0.6B",
"model_name": "Qwen/Qwen3-Embedding-0.6B",
"api_key": os.getenv("SILICONFLOW_API_KEY", "sk-ygqlrgrxrypykiiskuspuahkwihhbhhjhazqokntwdzfwqdv"),
"api_base": os.getenv("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn/v1"),
"embedding_dims": 1024
}
})
return config
def create_qwen_reranker_config():
"""创建Qwen重排器配置"""
config = RerankerConfigFactory(
backend="siliconflow",
config={
"model_name_or_path": "Qwen/Qwen3-Reranker-0.6B",
"model_name": "Qwen/Qwen3-Reranker-0.6B",
"api_key": os.getenv("SILICONFLOW_API_KEY", "sk-ygqlrgrxrypykiiskuspuahkwihhbhhjhazqokntwdzfwqdv"),
"api_base": os.getenv("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn/v1"),
"top_k": 10,
"max_chunks_per_query": 100
}
)
return config
def create_qwen_embedder():
"""创建Qwen嵌入器实例"""
config = create_qwen_embedder_config()
embedder = EmbedderFactory.from_config(config)
return embedder
def create_qwen_reranker():
"""创建Qwen重排器实例"""
config = create_qwen_reranker_config()
reranker = RerankerFactory.from_config(config)
return reranker
def test_qwen_models():
"""测试Qwen模型"""
print("🧪 测试Qwen嵌入和重排模型")
print("=" * 50)
try:
# 测试嵌入器
print("🔄 测试Qwen嵌入器...")
embedder = create_qwen_embedder()
test_texts = [
"这是一个测试文本",
"MemOS是一个智能记忆管理系统",
"Qwen模型提供高质量的嵌入向量"
]
embeddings = embedder.embed(test_texts)
print(f"✅ 嵌入器测试成功！向量维度: {len(embeddings[0])}")
# 测试重排器
print("🔄 测试Qwen重排器...")
reranker = create_qwen_reranker()
query = "智能记忆管理系统"
documents = [
"MemOS是一个智能记忆管理系统，支持向量搜索",
"数据分析平台需要考虑存储和计算",
"机器学习模型需要大量训练数据"
]
reranked_results = reranker.rerank(query, documents, top_k=3)
print(f"✅ 重排器测试成功！返回结果数: {len(reranked_results)}")
# 显示重排结果
print("\n📝 重排结果:")
for i, (doc, score) in enumerate(reranked_results, 1):
print(f"{i}. 分数: {score:.4f} - {doc[:50]}...")
return True
except Exception as e:
print(f"❌ 测试失败: {e}")
return False
if __name__ == "__main__":
success = test_qwen_models()
if success:
print("\n🎉 Qwen模型测试成功！")
else:
print("\n❌ Qwen模型测试失败")
sys.exit(1)