|
|
|
@ -10,6 +10,7 @@ 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 分词
|
|
|
|
@ -24,6 +25,7 @@ def text_to_embedding(text):
|
|
|
|
|
print("未找到有效词,返回零向量")
|
|
|
|
|
return [0.0] * model.vector_size
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 2. 使用连接池管理 Milvus 连接
|
|
|
|
|
milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)
|
|
|
|
|
|
|
|
|
@ -45,12 +47,21 @@ input_text = "小学数学中有哪些模型?"
|
|
|
|
|
current_embedding = text_to_embedding(input_text)
|
|
|
|
|
|
|
|
|
|
# 8. 查询与当前对话最相关的历史对话
|
|
|
|
|
start_time = time.time()
|
|
|
|
|
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=10) # 返回 2 条结果
|
|
|
|
|
expr = "document_id == 'MATH_DATA_1'" # 这回我只想查找 document_id='MATH_DATA_2' 的数据
|
|
|
|
|
#expr = "document_id == 'MATH_DATA_2'" # 这回我只想查找 document_id='MATH_DATA_2' 的数据
|
|
|
|
|
results = collection_manager.search(
|
|
|
|
|
current_embedding,
|
|
|
|
|
search_params,
|
|
|
|
|
expr=expr, # 新增条件表达式
|
|
|
|
|
limit=10
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
#results = collection_manager.search(current_embedding, search_params, limit=10) # 返回 2 条结果
|
|
|
|
|
end_time = time.time()
|
|
|
|
|
|
|
|
|
|
# 9. 输出查询结果
|
|
|
|
@ -79,4 +90,4 @@ print(f"查询耗时: {end_time - start_time:.4f} 秒")
|
|
|
|
|
milvus_pool.release_connection(connection)
|
|
|
|
|
|
|
|
|
|
# 12. 关闭连接池
|
|
|
|
|
milvus_pool.close()
|
|
|
|
|
milvus_pool.close()
|
|
|
|
|