Milvus是一个开源向量数据库,适用于各种规模的AI应用程序,从在Jupyter笔记本中运行演示聊天机器人到构建服务数十亿用户的Web规模搜索。本文将介绍如何在几分钟内在本地设置Milvus,并使用Python客户端库来生成、存储和搜索向量。
一、安装Milvus
下面将使用MilvusLite嵌入到客户端应用程序中的python库。Milvus还支持在Docker和Kubernetes上进行部署,用于生产用例。
在开始之前,请确保本地环境中有可用的Python3.8+。安装其中包含python客户端库和MilvusLite:
$ pip install -U pymilvus
如果使用的是GoogleColab,要启用刚刚安装的依赖项,可能需要重新启动运行时。(单击屏幕顶部的“运行时”菜单,然后从下拉菜单中选择“重新启动会话”。)
二、设置向量数据库
要创建本地Milvus向量数据库,只需指定一个文件名来存储所有数据,例如“milvus_demo.db”。
from pymilvus import MilvusClient client = MilvusClient("milvus_demo.db")
三、创建集合
在Milvus中,我们需要一个集合来存储向量及其相关的元数据。可以将其视为传统SQL数据库中的表。创建集合时,可以定义架构和索引参数来配置向量规格,例如维度、索引类型和距离指标。此外,还有一些复杂的概念可用于优化索引以实现向量搜索性能。现在,让我们只关注基础知识,并对所有可能的事物使用默认值。至少,只需要设置集合名称和集合的向量字段的维度。
if client.has_collection(collection_name="demo_collection"): client.drop_collection(collection_name="demo_collection") client.create_collection( collection_name="demo_collection", dimension=768, # The vectors we will use in this demo has 768 dimensions )
在上述设置中,主键和向量字段使用其默认名称(“id”和“vector”)。度量类型(向量距离定义)设置为其默认值(COSINE)。主键字段接受整数,不会自动递增(即不使用auto-id功能)。或者,可以按照此说明正式定义集合的架构。
四、准备数据
接下来使用向量对文本执行语义搜索,需要通过下载嵌入模型来生成文本向量,可以通过使用库中的实用函数轻松完成。
五、用向量表示文本
首先,安装模型库。此包包括PyTorch等基本ML工具。如果本地环境从未安装过PyTorch,则包下载可能需要一些时间。
$ pip install "pymilvus[model]"
使用默认模型生成向量嵌入。Milvus期望数据以字典列表的形式输入,其中每个字典代表一个数据记录,称为一个实体。
from pymilvus import model embedding_fn = model.DefaultEmbeddingFunction() docs = [ "Artificial intelligence was founded as an academic discipline in 1956.", "Alan Turing was the first person to conduct substantial research in AI.", "Born in Maida Vale, London, Turing was raised in southern England.", ] vectors = embedding_fn.encode_documents(docs) print("Dim:", embedding_fn.dim, vectors[0].shape) # Dim: 768 (768,) data = [ {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"} for i in range(len(vectors)) ] print("Data has", len(data), "entities, each with fields: ", data[0].keys()) print("Vector dim:", len(data[0]["vector"]))
Dim: 768 (768,) Data has 3 entities, each with fields: dict_keys(['id', 'vector', 'text', 'subject']) Vector dim: 768
或者使用随机向量的虚假表示:
如果由于网络问题而无法下载模型,则可以使用随机向量来表示文本,同时仍能完成示例。请注意,搜索结果不会反映语义相似性,因为向量是假的。
import random # Text strings to search from. docs = [ "Artificial intelligence was founded as an academic discipline in 1956.", "Alan Turing was the first person to conduct substantial research in AI.", "Born in Maida Vale, London, Turing was raised in southern England.", ] # Use fake representation with random vectors (768 dimension). vectors = [[random.uniform(-1, 1) for _ in range(768)] for _ in docs] data = [ {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"} for i in range(len(vectors)) ] print("Data has", len(data), "entities, each with fields: ", data[0].keys()) print("Vector dim:", len(data[0]["vector"]))
Data has 3 entities, each with fields: dict_keys(['id', 'vector', 'text', 'subject']) Vector dim: 768
六、插入数据
将数据插入到集合中:
res = client.insert(collection_name="demo_collection", data=data) print(res)
{'insert_count': 3, 'ids': [0, 1, 2], 'cost': 0}
七、语义搜索与向量搜索
现在可以通过将搜索查询文本表示为向量来进行语义搜索,并在Milvus上执行向量相似性搜索。
Milvus支持单个或多个向量的搜索请求。query_vectors变量包含一个向量列表,其中每个向量是一个浮点数数组。
query_vectors = embedding_fn.encode_queries(["Who is Alan Turing?"]) res = client.search( collection_name="demo_collection", # target collection data=query_vectors, # query vectors limit=2, # number of returned entities output_fields=["text", "subject"], # specifies fields to be returned ) print(res)
data: ["[{'id': 2, 'distance': 0.5859944820404053, 'entity': {'text': 'Born in Maida Vale, London, Turing was raised in southern England.', 'subject': 'history'}}, {'id': 1, 'distance': 0.5118255615234375, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.', 'subject': 'history'}}]"] , extra_info: {'cost': 0}
输出是一个结果列表,每个结果对应于一个搜索查询。每个查询的结果列表包括实体的主键、到查询向量的距离,以及指定的.output_fields。
八、具有元数据过滤功能的向量搜索
你可以在进行向量搜索时考虑元数据(在Milvus中称为“标量”字段)。这通过指定筛选表达式来实现。以下示例展示了如何使用字段进行搜索和筛选。
docs = [ "Machine learning has been used for drug design.", "Computational synthesis with AI algorithms predicts molecular properties.", "DDR1 is involved in cancers and fibrosis.", ] vectors = embedding_fn.encode_documents(docs) data = [ {"id": 3 + i, "vector": vectors[i], "text": docs[i], "subject": "biology"} for i in range(len(vectors)) ] client.insert(collection_name="demo_collection", data=data) res = client.search( collection_name="demo_collection", data=embedding_fn.encode_queries(["tell me AI related information"]), filter="subject == 'biology'", limit=2, output_fields=["text", "subject"], ) print(res)
data: ["[{'id': 4, 'distance': 0.27030569314956665, 'entity': {'text': 'Computational synthesis with AI algorithms predicts molecular properties.', 'subject': 'biology'}}, {'id': 3, 'distance': 0.16425910592079163, 'entity': {'text': 'Machine learning has been used for drug design.', 'subject': 'biology'}}]"] , extra_info: {'cost': 0}
九、查询
query()操作用于检索与条件匹配的所有实体,例如根据过滤表达式或特定ID进行检索。
检索标量字段具有特定值的所有实体:
res = client.query( collection_name="demo_collection", filter="subject == 'history'", output_fields=["text", "subject"], )
通过主键直接检索实体:
res = client.query( collection_name="demo_collection", ids=[0, 2], output_fields=["vector", "text", "subject"], )
十、删除实体
要清除数据,你可以通过主键删除实体,也可以删除与特定筛选表达式匹配的所有实体。
通过主键删除实体:
res = client.delete(collection_name="demo_collection", ids=[0, 2]) print(res) res = client.delete( collection_name="demo_collection", filter="subject == 'biology'", ) print(res)
[0, 2] [3, 4, 5]
十一、加载现有数据
由于MilvusLite将所有数据存储在本地文件中,你可以在程序终止后使用现有文件创建一个文件,并将所有数据加载到内存中。例如,从“milvus_demo.db”文件中恢复集合,并继续写入数据。
from pymilvus import MilvusClient client = MilvusClient("milvus_demo.db")
十二、删除集合
要删除集合中的所有数据,可以使用以下命令:
client.drop_collection(collection_name="demo_collection")
十三、了解更多信息
MilvusLite非常适合开始使用本地Python程序。如果你处理大量数据或希望在生产环境中使用Milvus,可以学习如何在Docker和Kubernetes上部署Milvus。Milvus的所有部署模式共享相同的API,因此切换到其他部署模式时,客户端代码无需做太大修改。只需指定部署在任何地方的Milvus服务器的URI和Token:
client = MilvusClient(uri="http://localhost:19530", token="root:Milvus")
Milvus提供REST和gRPCAPI,以及Python、Java、Go、C#和Node.js等语言的客户端库。