发布于 2026-01-06 0 阅读
0

如何使用 Milvus 存储和查询向量嵌入

如何使用 Milvus 存储和查询向量嵌入

博客文章中的抽象图片

在当今数据驱动的世界中,管理和搜索大型数据集变得日益重要。Milvus 是一个强大的工具,可以应对这一挑战。Milvus 是一个专为人工智能应用设计的开源向量数据库。在本篇博文中,我们将探讨 Milvus 的一个实际 Python 实现,并展示如何将其与文本嵌入技术相结合,从而创建一个高效的搜索系统。

本文所有代码均可在配套的 GitHub 存储库中找到。

Milvus:向量数据库

Milvus旨在为矢量数据提供可扩展、可靠且快速的搜索功能。它尤其适用于图像和视频识别、自然语言处理和推荐系统等应用,在这些应用中,数据可以表示为高维矢量。

设置 Milvus

在深入代码之前,请确保已安装并运行 Milvus。我们的 Python 脚本的第一步是建立与 Milvus 服务器的连接:

from pymilvus import connections

def connect_to_milvus():
    try:
        connections.connect("default", host="localhost", port="19530")
        print("Connected to Milvus.")
    except Exception as e:
        print(f"Failed to connect to Milvus: {e}")
        raise
Enter fullscreen mode Exit fullscreen mode

此函数尝试连接到运行在本地计算机上的 Milvus 服务器。错误处理至关重要,它可以捕获并了解连接过程中可能出现的任何问题。

在 Milvus 中创建集合

在 Milvus 中,集合就像传统数据库中的表一样,用于存储数据。每个集合可以包含多个字段,类似于表中的列。在我们的示例中,我们创建了一个包含三个字段的集合:主键(pk)、源文本(source)和嵌入(embeddings):

from pymilvus import FieldSchema, CollectionSchema, DataType, Collection

def create_collection(name, fields, description):
    schema = CollectionSchema(fields, description)
    collection = Collection(name, schema, consistency_level="Strong")
    return collection

# Define fields for our collection
fields = [
    FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=100),
    FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=500),
    FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=768)
]

collection = create_collection("hello_milvus", fields, "Collection for demo purposes")
Enter fullscreen mode Exit fullscreen mode

在此代码片段中,嵌入的维度为 768(对于我们特定的自定义嵌入模型,将在下一节中提到),这应该与您使用的嵌入模型的输出一致。

在 Python 中生成文本嵌入

在将数据插入 Milvus 数据集之前,我们需要生成文本嵌入。这个过程需要使用 transformers 库中的预训练模型将文本转换为数值向量。在我们的代码中,我们使用thenlper/gte-base 模型embedding_util.py来实现这一目标。对于我们的应用程序而言,这个过程由一个专门处理向量嵌入生成的模块进行抽象。

有关我们的自定义embedding_util.py模块如何创建矢量嵌入的更多详细信息,请查看我的博客文章,了解如何使用 weaviate 存储和查询矢量嵌入

生成和插入数据

为了从文本生成词嵌入,我们使用了前面提到的库中的预训练模型transformers。该模型将文本转换为可以存储在我们的 Milvus 集合中的数值向量:

from embedding_util import generate_embeddings

documents = [...]
embeddings = [generate_embeddings(doc) for doc in documents]
entities = [
    [str(i) for i in range(len(documents))],
    [str(doc) for doc in documents],
    embeddings
]

insert_result = insert_data(collection, entities)
Enter fullscreen mode Exit fullscreen mode

insert_data函数将我们的数据插入到 Milvus 集合中,然后刷新操作以确保数据持久性。

创建索引以实现高效搜索

Milvus 使用索引来加快搜索过程。这里,我们在嵌入字段上创建IVF_FLAT索引:

def create_index(collection, field_name, index_type, metric_type, params):
    index = {"index_type": index_type, "metric_type": metric_type, "params": params}
    collection.create_index(field_name, index)

create_index(collection, "embeddings", "IVF_FLAT", "L2", {"nlist": 128})
Enter fullscreen mode Exit fullscreen mode

执行向量搜索

数据索引完成后,我们现在可以基于向量相似性进行搜索:

def search_and_query(collection, search_vectors, search_field, search_params):
    collection.load()
    result = collection.search(search_vectors, search_field, search_params, limit=3, output_fields=["source"])
    print_search_results(result, "Vector search results:")

query = "Give me some content about the ocean"
query_vector = generate_embeddings(query)
search_and_query(collection, [query_vector], "embeddings", {"metric_type": "L2", "params": {"nprobe": 10}})
Enter fullscreen mode Exit fullscreen mode

在此搜索中,我们正在寻找与查询“给我一些关于海洋的内容”最相似的前 3 个文档。

如果应用程序运行成功,您应该会看到以下向量搜索结果,按余弦距离排序(距离越小,语义越相似):

Vector search results:
Hit: id: 6, distance: 0.39819106459617615, entity: {'source': 'The sunset paints the sky with shades of orange, pink, and purple, reflecting on the calm sea.'}, source field: The sunset paints the sky with shades of orange, pink, and purple, reflecting on the calm sea.
Hit: id: 4, distance: 0.4780573844909668, entity: {'source': 'The ancient tree, with its gnarled branches and deep roots, whispers secrets of the past.'}, source field: The ancient tree, with its gnarled branches and deep roots, whispers secrets of the past.
Hit: id: 0, distance: 0.4835127890110016, entity: {'source': 'A group of vibrant parrots chatter loudly, sharing stories of their tropical adventures.'}, source field: A group of vibrant parrots chatter loudly, sharing stories of their tropical adventures.
Enter fullscreen mode Exit fullscreen mode

清理

操作完成后,最好通过删除实体和丢弃集合来进行清理:

delete_entities(collection, f'pk in ["{insert_result.primary_keys[0]}", "{insert_result.primary_keys[1]}"]')
drop_collection("hello_milvus")
Enter fullscreen mode Exit fullscreen mode

结论

Milvus 提供了一种强大而灵活的方式来处理矢量数据。通过将其与自然语言处理技术相结合,我们可以构建复杂的搜索和推荐系统。这里演示的 Python 脚本只是一个简单的示例,但它的潜在应用范围非常广泛。

无论您是处理大规模图像数据库、复杂的推荐系统还是高级 NLP 任务,Milvus 都可以成为您 AI 工具库中不可或缺的工具。

文章来源:https://dev.to/stephenc222/how-to-use-milvus-to-store-and-query-vector-embeddings-5hhl