You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.8 KiB
78 lines
2.8 KiB
import time
|
|
import jieba # 导入 jieba 分词库
|
|
from WxMini.Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
|
|
from WxMini.Milvus.Utils.MilvusConnectionPool import *
|
|
from WxMini.Milvus.Config.MulvusConfig import *
|
|
from gensim.models import KeyedVectors
|
|
|
|
# 1. 加载预训练的 Word2Vec 模型
|
|
model_path = MS_MODEL_PATH # 替换为你的 Word2Vec 模型路径
|
|
model = KeyedVectors.load_word2vec_format(model_path, binary=False, limit=MS_MODEL_LIMIT)
|
|
print(f"模型加载成功，词向量维度: {model.vector_size}")
|
|
|
|
|
|
# 将文本转换为嵌入向量
|
|
def text_to_embedding(text):
|
|
words = jieba.lcut(text) # 使用 jieba 分词
|
|
print(f"文本: {text}, 分词结果: {words}")
|
|
embeddings = [model[word] for word in words if word in model] # 获取有效词向量
|
|
print(f"有效词向量数量: {len(embeddings)}")
|
|
if embeddings:
|
|
avg_embedding = sum(embeddings) / len(embeddings) # 计算平均向量
|
|
print(f"生成的平均向量: {avg_embedding[:5]}...") # 打印前 5 维
|
|
return avg_embedding
|
|
else:
|
|
print("未找到有效词，返回零向量")
|
|
return [0.0] * model.vector_size
|
|
|
|
|
|
# 2. 使用连接池管理 Milvus 连接
|
|
milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)
|
|
|
|
# 3. 从连接池中获取一个连接
|
|
connection = milvus_pool.get_connection()
|
|
|
|
# 4. 初始化集合管理器
|
|
collection_name = MS_COLLECTION_NAME
|
|
collection_manager = MilvusCollectionManager(collection_name)
|
|
|
|
# 5. 加载集合到内存
|
|
collection_manager.load_collection()
|
|
|
|
# 6. 输入一句话
|
|
input_text = input("请输入一句话：") # 例如：“我今天心情不太好”
|
|
|
|
# 7. 将文本转换为嵌入向量
|
|
current_embedding = text_to_embedding(input_text)
|
|
print(f"当前文本的向量: {current_embedding[:5]}...") # 打印前 5 维
|
|
|
|
# 8. 查询与当前对话最相关的历史对话
|
|
search_params = {
|
|
"metric_type": "L2", # 使用 L2 距离度量方式
|
|
"params": {"nprobe": MS_NPROBE} # 设置 IVF_FLAT 的 nprobe 参数
|
|
}
|
|
start_time = time.time()
|
|
results = collection_manager.search(current_embedding, search_params, limit=2) # 返回 2 条结果
|
|
end_time = time.time()
|
|
|
|
# 9. 输出查询结果
|
|
print("最相关的历史对话:")
|
|
if results:
|
|
for hits in results:
|
|
for hit in hits:
|
|
try:
|
|
text = collection_manager.query_text_by_id(hit.id)
|
|
print(f"- {text} (距离: {hit.distance})")
|
|
except Exception as e:
|
|
print(f"查询失败: {e}")
|
|
else:
|
|
print("未找到相关历史对话，请检查查询参数或数据。")
|
|
|
|
# 10. 输出查询耗时
|
|
print(f"查询耗时: {end_time - start_time:.4f} 秒")
|
|
|
|
# 11. 释放连接
|
|
milvus_pool.release_connection(connection)
|
|
|
|
# 12. 关闭连接池
|
|
milvus_pool.close() |