milvus lite快速实践-了解RAG落地背后的机制

Milvus Lite 是Milvus 的轻量级版本，Milvus 是一个开源向量数据库，通过向量嵌入和相似性搜索为人工智能应用提供支持，最典型的应用场景就是 RAG（Retrieval-Augmented Generation，检索增强生成），为 RAG 系统提供了强大的向量存储和检索能力。通过下面的实践，可以了解文本向量化与相似度匹配（语义匹配）的大概过程，了解RAG落地背后的机制。

安装 milvus lite

Milvus Lite 已包含在Milvus 的 Python SDK 中。它可以通过pip install pymilvus 简单地部署。

部署前提：

• python 3.8+;
• Ubuntu >= 20.04（x86_64 或者 arm64）
• MacOS >= 11.0（苹果 M1/M2 或者 x86_64）

安装命令如下：

pip install -U pymilvus

文本向量化

创建向量数据库

通过实例化MilvusClient ，指定一个存储所有数据的文件名来创建本地的Milvus向量数据库。例如下：

from pymilvus import MilvusClient

client = MilvusClient("milvus_demo.db")

创建 Collections

Collections类似于传统 SQL 数据库中的表格。在 Milvus 中Collections 用来存储向量及其相关元数据。创建 Collections 时，可以定义 Schema 和索引参数来配置向量规格，如维度、索引类型和远距离度量。

创建 Collections 时，至少需要设置名称和向量场的维度，未指定的参数使用默认值。

if client.has_collection(collection_name="demo_collection"):
    client.drop_collection(collection_name="demo_collection")
client.create_collection(
    collection_name="demo_collection",
    dimension=768,  # 该demo中的向量规格为768维
)

文本向量化

下载 embedding模型 为测试的文本生成向量。

首先，安装模型库，该软件包包含 PyTorch 等基本 ML 工具。如果本地环境从未安装过 PyTorch，则软件包下载可能需要一些时间。

pip install "pymilvus[model]"

用默认模型生成向量 Embeddings。Milvus 希望数据以字典列表的形式插入，每个字典代表一条数据记录，称为实体。

from pymilvus import model

# 如果访问 https://huggingface.co/ 失败，请取消注释以下路径
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 这将下载一个较小的嵌入模型 "paraphrase-albert-small-v2"（~50MB）。
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)

# 输出向量有768维，与创建的集合匹配。
print("Dim:", embedding_fn.dim, vectors[0].shape)  # Dim: 768 (768,)


# 每个实体（entity）都有 id、向量表示、原始文本和主题标签，用于演示元数据过滤。
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

插入数据

下面把把数据插入 Collections 中：

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",  # 目标 collection
    data=query_vectors,  # query vectors 查询的向量
    limit=2,  # 返回的entity数量
    output_fields=["text", "subject"],  # 指定查询返回的实体字段
)

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 的实体详细信息。

从distance值可以评估搜索结果的相关度，越接近0，则搜索结果更相关。关于相似度量更多说明，请参考相似度量。

带元数据过滤的向量搜索

你还可以在考虑元数据值（在 Milvus 中称为 "标量 "字段，因为标量指的是非向量数据）的同时进行向量搜索。这可以通过指定特定条件的过滤表达式来实现。让我们在下面的示例中看看如何使用subject 字段进行搜索和筛选。

# 在另外一个主题中插入更多文档。
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)

# 通过filter过滤subject。这将排除任何主题为“历史”的文本，尽管与查询向量非常接近。
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}

默认情况下，标量字段不编制索引。如果需要在大型数据集中执行元数据过滤搜索，可以考虑使用固定 Schema，同时打开索引以提高搜索性能。

除了向量搜索，还可以执行其他类型的搜索。

查询

查询()是一种操作符，用于检索与某个条件（如过滤表达式或与某些 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)

# 使用filter 表达式删除所有 subject 为 "biology" 的实体。
res = client.delete(
    collection_name="demo_collection",
    filter="subject == 'biology'",
)

print(res)

输出：

[0, 2]
[3, 4, 5]

加载现有数据

由于 Milvus Lite 的所有数据都存储在本地文件中，因此即使在程序终止后，你也可以通过创建一个带有现有文件的MilvusClient ，将所有数据加载到内存中。例如，这将恢复 "milvus_demo.db "文件中的 Collections，并继续向其中写入数据。

from pymilvus import MilvusClient

client = MilvusClient("milvus_demo.db")

删除 Collections

如果想删除某个 Collections 中的所有数据，可以通过以下方法丢弃该 Collections

# 删除 Collections
client.drop_collection(collection_name="demo_collection")

参考资料

1. https://milvus.io/docs/zh/quickstart.md

2. https://milvus.io/docs/zh/milvus_lite.md