Milvus - 构建向量数据库并进行相似度查询

向量相似度检索在大规模数据分析和机器学习应用中是一个非常关键的任务，特别是在处理文本、图像或其他嵌入向量时。Milvus 是一个高性能的开源向量数据库，专为存储和检索大规模向量数据设计。本文将介绍如何在 Docker 中安装 Milvus，并展示如何使用 Python 插入向量数据及进行向量相似度查询。

一、在 Docker 中安装并运行 Milvus

Milvus 提供了简便的安装方式，可以通过 Docker 快速启动 Milvus 实例。

1. 安装 Docker

确保你已在系统上安装了 Docker。如果未安装，可以访问 Docker 官方网站 获取安装指南。

2. 在 Docker 中安装 Milvus

通过以下步骤下载并启动 Milvus 实例。

# 下载 Milvus 安装脚本
curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh

# 启动 Milvus 容器
bash standalone_embed.sh start

运行该脚本后，一个名为 milvus-standalone 的 Docker 容器将启动，并在 19530 端口提供 Milvus 服务。嵌入式的 etcd 也将在同一容器中运行，使用 2379 端口进行服务。

3. 停止和删除 Milvus

你可以随时停止和删除 Milvus 容器及其数据：

# 停止 Milvus
bash standalone_embed.sh stop

# 删除 Milvus 容器和数据
bash standalone_embed.sh delete

二、创建向量集合

在 Milvus 中，数据存储在集合（Collection）中，类似于传统数据库中的表。集合可以包含多个字段，包括嵌入向量字段。接下来，我们将通过 Python 创建一个存储文本及其对应向量的集合。

1. 连接 Milvus 并定义集合

首先，我们通过 Python 的 pymilvus 库连接 Milvus 并创建集合。

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection

# 连接 Milvus
connections.connect("default", host="localhost", port="19530")

# 定义集合 schema
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="user_id", dtype=DataType.INT64),
    FieldSchema(name="file_id", dtype=DataType.INT64),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=384),  # 384维向量
]

schema = CollectionSchema(fields, description="知识库向量数据")

# 创建集合
collection = Collection("knowledge_vector_data", schema=schema)

2. 创建索引

为提高查询效率，Milvus 支持为向量字段创建索引。下面是为向量字段创建索引的示例：

# 创建索引
index_params = {
    "metric_type": "IP",  # 内积，用于向量相似度
    "index_type": "IVF_FLAT",  # 索引类型
    "params": {"nlist": 128}
}
collection.create_index(field_name="embedding", index_params=index_params)
print("Index created.")

三、向量化、插入数据并进行相似度查询

在这个部分，我们将通过 Python 脚本，将测试句子向量化后插入 Milvus，并进行相似度查询。

1. 安装必要的 Python 库

首先，确保安装了 sentence-transformers 和 pymilvus 库：

pip install sentence-transformers pymilvus numpy

2. 完整的 Python 脚本

以下是一个完整的 Python 脚本，包含了向量化句子、插入数据和相似度查询的功能。

import numpy as np
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
from sentence_transformers import SentenceTransformer

# 连接 Milvus
connections.connect("default", host="localhost", port="19530")

# 定义集合 schema
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="user_id", dtype=DataType.INT64),
    FieldSchema(name="file_id", dtype=DataType.INT64),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=384),  # 384维向量
]

schema = CollectionSchema(fields, description="知识库向量数据")

# 创建集合
collection = Collection("knowledge_vector_data", schema=schema)


# 向量化句子
def vectorize_sentences(sentences):
    model = SentenceTransformer('all-MiniLM-L6-v2')
    return model.encode(sentences)


# 插入向量到 Milvus
def insert_vectors_to_milvus(sentences, vectors):
    data = [
        # [None] * len(sentences),  # ID 自动生成
        [1] * len(sentences),  # user_id
        [1] * len(sentences),  # file_id
        sentences,  # content
        vectors.tolist()  # 向量
    ]

    collection.insert(data)
    print(f"Inserted {len(sentences)} sentences into Milvus")


# 创建索引以加速检索
def create_index():
    index_params = {
        "metric_type": "IP",  # 内积，用于向量相似度
        "index_type": "IVF_FLAT",  # 索引类型
        "params": {"nlist": 128}
    }
    collection.create_index(field_name="embedding", index_params=index_params)
    print("Index created.")


# 查询相似句子
def query_similar_sentences(target_vector, top_k=5):

    collection.load()

    search_params = {"metric_type": "IP", "params": {"nprobe": 10}}

    # 执行查询
    results = collection.search(
        data=[target_vector],
        anns_field="embedding",
        param=search_params,
        limit=top_k,
        output_fields=["content"]
    )

    for result in results[0]:
        print(f"Content: {result.entity.get('content')}, Similarity score: {result.score}")


if __name__ == "__main__":
    # 测试句子
    sentences = [
        "This is a test sentence.",
        "Another sentence for testing.",
        "Milvus vector database integration."
    ]

    # 向量化测试句子
    vectors = vectorize_sentences(sentences)

    # 插入向量到 Milvus
    insert_vectors_to_milvus(sentences, vectors)

    # 创建索引
    create_index()

    # 目标句子，进行相似度查询
    target_sentence = "This is a test sentence."
    target_vector = vectorize_sentences([target_sentence])[0]

    # 查询与目标句子最相似的句子
    query_similar_sentences(target_vector)

四、总结

通过本文的步骤，我们学习了如何在 Docker 中安装 Milvus，并使用 Python 向量化文本数据、插入数据和执行相似度查询。Milvus 是处理大规模向量数据的理想工具，特别适合于多媒体数据的相似度搜索。

原文地址：https://blog.csdn.net/fenglingguitar/article/details/142357290

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！