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

ops/2024/9/20 8:13:39/ 标签: 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_idsentences,  # contentvectors.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)# 插入向量到 Milvusinsert_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)