Milvus 向量检索

1

文档 → Embedding 模型 → 文档向量 → 与原文、类别组成记录 → Milvus 集合

1

查询文本 → Embedding 模型 → 查询向量 → Milvus 搜索 → 文档、相似度分数

1
2
3

python3 -m venv .venv
source .venv/bin/activate
python -m pip install "pymilvus==3.0.2" "milvus-lite==3.2.2" "sentence-transformers==5.1.2" "faiss-cpu==1.12.0"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

import numpy as np
from sentence_transformers import SentenceTransformer

docs = [
    {"id": 101, "text": "Faiss 提供向量索引和相似度搜索，原始文档由应用程序管理。", "category": "数据库"},
    {"id": 102, "text": "Milvus 支持在向量搜索中按类别等字段过滤文档。", "category": "数据库"},
    {"id": 103, "text": "MySQL 使用关系表保存数据，通过 SQL 查询和更新记录。", "category": "数据库"},
    {"id": 104, "text": "梯度下降根据损失函数的梯度调整模型参数。", "category": "机器学习"},
    {"id": 105, "text": "文本分类模型为文档预测类别，用于组织搜索结果。", "category": "机器学习"},
    {"id": 106, "text": "Docker Compose 使用配置文件管理多个容器及其数据卷。", "category": "运维"},
]

model = SentenceTransformer("BAAI/bge-m3", device="cpu")
texts = [doc["text"] for doc in docs]
vectors = model.encode(texts, normalize_embeddings=True).astype(np.float32)

ids = np.array([doc["id"] for doc in docs], dtype=np.int64)
categories = [doc["category"] for doc in docs]
np.savez("vectors.npz", vectors=vectors, ids=ids, texts=texts, categories=categories)
print(vectors.shape)
# 输出：(6, 1024)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

from pymilvus import DataType, MilvusClient

client = MilvusClient(uri="./milvus_bge_m3.db")
collection = "article_demo"

schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=vectors.shape[1])
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1024)
schema.add_field(field_name="category", datatype=DataType.VARCHAR, max_length=64)

1
2
3
4
5
6
7
8
9

index_params = client.prepare_index_params()
index_params.add_index(field_name="vector", index_type="FLAT", metric_type="COSINE")
client.create_collection(collection_name=collection, schema=schema, index_params=index_params,
                         consistency_level="Strong")

records = [{**doc, "vector": vector.tolist()} for doc, vector in zip(docs, vectors)]
result = client.insert(collection_name=collection, data=records)
print(result["insert_count"])
# 输出：6

1
2
3
4

query = "检索结果只保留指定类别的文档"
query_vector = model.encode([query], normalize_embeddings=True).astype(np.float32)
print(query_vector.shape)
# 输出：(1, 1024)

1
2
3
4
5

hits = client.search(collection_name=collection, data=query_vector.tolist(), anns_field="vector",
                     limit=3, output_fields=["text", "category"])

for hit in hits[0]:
    print(hit["id"], round(hit["distance"], 4), hit["entity"]["text"])

1
2
3
4
5

hits = client.search(collection_name=collection, data=query_vector.tolist(), anns_field="vector",
                     filter='category == "数据库"', limit=2, output_fields=["text", "category"])

for hit in hits[0]:
    print(hit["id"], hit["entity"]["category"], hit["entity"]["text"])

1
2
3
4
5
6

rows = client.query(collection_name=collection, filter="id in [101, 102]", output_fields=["text", "category"])
for row in sorted(rows, key=lambda row: row["id"]):
    print(row["id"], row["category"], row["text"])
# 输出：
# 101 数据库 Faiss 提供向量索引和相似度搜索，原始文档由应用程序管理。
# 102 数据库 Milvus 支持在向量搜索中按类别等字段过滤文档。

1
2
3
4

updated_record = {**records[0], "category": "检索工具"}
result = client.upsert(collection_name=collection, data=[updated_record])
print(result["upsert_count"])
# 输出：1

1
2
3

client.delete(collection_name=collection, ids=[106])
print(client.query(collection_name=collection, filter="id == 106", output_fields=["text"]))
# 输出：data: [], extra_info: {}

1
2
3
4
5
6
7
8

client.close()
client = MilvusClient(uri="./milvus_bge_m3.db")
client.load_collection(collection_name="article_demo")
rows = client.query(collection_name="article_demo", filter="id == 101", output_fields=["text", "category"])
for row in rows:
    print(row["id"], row["category"], row["text"])
client.close()
# 输出：101 检索工具 Faiss 提供向量索引和相似度搜索，原始文档由应用程序管理。

1
2
3
4
5
6
7
8
9

mkdir milvus-demo
cd milvus-demo

curl -L --fail \
  https://github.com/milvus-io/milvus/releases/download/v3.0.2/milvus-standalone-docker-compose.yaml \
  -o docker-compose.yaml

docker compose up -d
docker compose ps -a

1
2

docker compose logs --tail=100 standalone
curl --fail http://localhost:9091/healthz

1
2
3
4
5

from pymilvus import MilvusClient

client = MilvusClient(uri="http://localhost:19530")
print(client.list_collections())
# 新实例输出：[]

1
2
3
4
5
6

import numpy as np

data = np.load("vectors.npz")
vectors = data["vectors"]
docs = [{"id": int(doc_id), "text": str(text), "category": str(category)}
        for doc_id, text, category in zip(data["ids"], data["texts"], data["categories"])]

1

docker compose down

1

docker compose up -d

1
2
3
4

client = MilvusClient(uri="http://localhost:19530")
client.load_collection(collection_name="article_demo")
print(client.query(collection_name="article_demo", filter="id == 101", output_fields=["text", "category"]))
client.close()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

from milvus_client import MilvusCollectionClient

store = MilvusCollectionClient(
    uri="http://localhost:19530", collection_name="article_demo", vector_field="vector",
    metric_type="COSINE", output_fields=["text", "category"], timeout=30,
)
try:
    store.load()
    hits = store.search(query_vector[0].tolist(), limit=3, filter='category == "数据库"')
    for hit in hits:
        print(hit["id"], hit["distance"], hit["entity"]["text"])
    print(store.query("id == 101"))
finally:
    store.close()