210 lines
7.8 KiB
Python
210 lines
7.8 KiB
Python
import json
|
|
import os
|
|
import numpy as np
|
|
from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
|
|
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
|
|
|
|
# 1. 初始化设置
|
|
COLLECTION_NAME = "dragon_hybrid_demo"
|
|
MILVUS_URI = "http://localhost:19530" # 服务器模式
|
|
DATA_PATH = "../../data/C4/metadata/dragon.json" # 相对路径
|
|
BATCH_SIZE = 50
|
|
|
|
# 2. 连接 Milvus 并初始化嵌入模型
|
|
print(f"--> 正在连接到 Milvus: {MILVUS_URI}")
|
|
connections.connect(uri=MILVUS_URI)
|
|
|
|
print("--> 正在初始化 BGE-M3 嵌入模型...")
|
|
ef = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
|
|
print(f"--> 嵌入模型初始化完成。密集向量维度: {ef.dim['dense']}")
|
|
|
|
# 3. 创建 Collection
|
|
milvus_client = MilvusClient(uri=MILVUS_URI)
|
|
if milvus_client.has_collection(COLLECTION_NAME):
|
|
print(f"--> 正在删除已存在的 Collection '{COLLECTION_NAME}'...")
|
|
milvus_client.drop_collection(COLLECTION_NAME)
|
|
|
|
fields = [
|
|
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
|
FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
|
|
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=256),
|
|
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
|
|
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
|
|
FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
|
|
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
|
|
FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
|
|
FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
|
FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=ef.dim["dense"])
|
|
]
|
|
|
|
# 如果集合不存在，则创建它及索引
|
|
if not milvus_client.has_collection(COLLECTION_NAME):
|
|
print(f"--> 正在创建 Collection '{COLLECTION_NAME}'...")
|
|
schema = CollectionSchema(fields, description="关于龙的混合检索示例")
|
|
# 创建集合
|
|
collection = Collection(name=COLLECTION_NAME, schema=schema, consistency_level="Strong")
|
|
print("--> Collection 创建成功。")
|
|
|
|
# 4. 创建索引
|
|
print("--> 正在为新集合创建索引...")
|
|
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
|
|
collection.create_index("sparse_vector", sparse_index)
|
|
print("稀疏向量索引创建成功。")
|
|
|
|
dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
|
|
collection.create_index("dense_vector", dense_index)
|
|
print("密集向量索引创建成功。")
|
|
|
|
collection = Collection(COLLECTION_NAME)
|
|
|
|
# 5. 加载数据并插入
|
|
collection.load()
|
|
print(f"--> Collection '{COLLECTION_NAME}' 已加载到内存。")
|
|
|
|
if collection.is_empty:
|
|
print(f"--> Collection 为空，开始插入数据...")
|
|
if not os.path.exists(DATA_PATH):
|
|
raise FileNotFoundError(f"数据文件未找到: {DATA_PATH}")
|
|
with open(DATA_PATH, 'r', encoding='utf-8') as f:
|
|
dataset = json.load(f)
|
|
|
|
docs, metadata = [], []
|
|
for item in dataset:
|
|
parts = [
|
|
item.get('title', ''),
|
|
item.get('description', ''),
|
|
item.get('location', ''),
|
|
item.get('environment', ''),
|
|
# *item.get('combat_details', {}).get('combat_style', []),
|
|
# *item.get('combat_details', {}).get('abilities_used', []),
|
|
# item.get('scene_info', {}).get('time_of_day', '')
|
|
]
|
|
docs.append(' '.join(filter(None, parts)))
|
|
metadata.append(item)
|
|
print(f"--> 数据加载完成，共 {len(docs)} 条。")
|
|
|
|
print("--> 正在生成向量嵌入...")
|
|
embeddings = ef(docs)
|
|
print("--> 向量生成完成。")
|
|
|
|
print("--> 正在分批插入数据...")
|
|
# 为每个字段准备批量数据
|
|
img_ids = [doc["img_id"] for doc in metadata]
|
|
paths = [doc["path"] for doc in metadata]
|
|
titles = [doc["title"] for doc in metadata]
|
|
descriptions = [doc["description"] for doc in metadata]
|
|
categories = [doc["category"] for doc in metadata]
|
|
locations = [doc["location"] for doc in metadata]
|
|
environments = [doc["environment"] for doc in metadata]
|
|
|
|
# 获取向量
|
|
sparse_vectors = embeddings["sparse"]
|
|
dense_vectors = embeddings["dense"]
|
|
|
|
# 插入数据
|
|
collection.insert([
|
|
img_ids,
|
|
paths,
|
|
titles,
|
|
descriptions,
|
|
categories,
|
|
locations,
|
|
environments,
|
|
sparse_vectors,
|
|
dense_vectors
|
|
])
|
|
|
|
collection.flush()
|
|
print(f"--> 数据插入完成，总数: {collection.num_entities}")
|
|
else:
|
|
print(f"--> Collection 中已有 {collection.num_entities} 条数据，跳过插入。")
|
|
|
|
# 6. 执行搜索
|
|
search_query = "悬崖上的巨龙"
|
|
search_filter = 'category in ["western_dragon", "chinese_dragon", "movie_character"]'
|
|
top_k = 5
|
|
|
|
print(f"\n{'='*20} 开始混合搜索 {'='*20}")
|
|
print(f"查询: '{search_query}'")
|
|
print(f"过滤器: '{search_filter}'")
|
|
|
|
query_embeddings = ef([search_query])
|
|
dense_vec = query_embeddings["dense"][0]
|
|
sparse_vec = query_embeddings["sparse"]._getrow(0)
|
|
|
|
# 打印向量信息
|
|
print("\n=== 向量信息 ===")
|
|
print(f"密集向量维度: {len(dense_vec)}")
|
|
print(f"密集向量前5个元素: {dense_vec[:5]}")
|
|
print(f"密集向量范数: {np.linalg.norm(dense_vec):.4f}")
|
|
|
|
print(f"\n稀疏向量维度: {sparse_vec.shape[1]}")
|
|
print(f"稀疏向量非零元素数量: {sparse_vec.nnz}")
|
|
print("稀疏向量前5个非零元素:")
|
|
for i in range(min(5, sparse_vec.nnz)):
|
|
print(f" - 索引: {sparse_vec.indices[i]}, 值: {sparse_vec.data[i]:.4f}")
|
|
density = (sparse_vec.nnz / sparse_vec.shape[1] * 100)
|
|
print(f"\n稀疏向量密度: {density:.8f}%")
|
|
|
|
# 定义搜索参数
|
|
search_params = {"metric_type": "IP", "params": {}}
|
|
|
|
# 先执行单独的搜索
|
|
print("\n--- [单独] 密集向量搜索结果 ---")
|
|
dense_results = collection.search(
|
|
[dense_vec],
|
|
anns_field="dense_vector",
|
|
param=search_params,
|
|
limit=top_k,
|
|
expr=search_filter,
|
|
output_fields=["title", "path", "description", "category", "location", "environment"]
|
|
)[0]
|
|
|
|
for i, hit in enumerate(dense_results):
|
|
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
|
print(f" 路径: {hit.entity.get('path')}")
|
|
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
|
|
|
print("\n--- [单独] 稀疏向量搜索结果 ---")
|
|
sparse_results = collection.search(
|
|
[sparse_vec],
|
|
anns_field="sparse_vector",
|
|
param=search_params,
|
|
limit=top_k,
|
|
expr=search_filter,
|
|
output_fields=["title", "path", "description", "category", "location", "environment"]
|
|
)[0]
|
|
|
|
for i, hit in enumerate(sparse_results):
|
|
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
|
print(f" 路径: {hit.entity.get('path')}")
|
|
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
|
|
|
print("\n--- [混合] 稀疏+密集向量搜索结果 ---")
|
|
# 创建 RRF 融合器
|
|
rerank = RRFRanker(k=60)
|
|
|
|
# 创建搜索请求
|
|
dense_req = AnnSearchRequest([dense_vec], "dense_vector", search_params, limit=top_k)
|
|
sparse_req = AnnSearchRequest([sparse_vec], "sparse_vector", search_params, limit=top_k)
|
|
|
|
# 执行混合搜索
|
|
results = collection.hybrid_search(
|
|
[sparse_req, dense_req],
|
|
rerank=rerank,
|
|
limit=top_k,
|
|
output_fields=["title", "path", "description", "category", "location", "environment"]
|
|
)[0]
|
|
|
|
# 打印最终结果
|
|
for i, hit in enumerate(results):
|
|
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
|
print(f" 路径: {hit.entity.get('path')}")
|
|
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
|
|
|
# 7. 清理资源
|
|
milvus_client.release_collection(collection_name=COLLECTION_NAME)
|
|
print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
|
|
milvus_client.drop_collection(COLLECTION_NAME)
|
|
print(f"已删除 Collection: '{COLLECTION_NAME}'")
|