【拥抱AI】如何使用Milvus向量数据库进行数据库检索？

news2024/11/25 9:55:36

使用向量数据库进行数据库检索是一种高效的方法，特别是在处理大规模、高维度的数据时。以下是详细的步骤，帮助你理解和实施这一过程：

1. 准备环境

选择向量数据库

常见的向量数据库有：

Faiss：由Facebook AI Research开发，适用于大规模向量搜索。
Milvus：开源的向量数据库，支持多种索引类型和相似度度量。
Elasticsearch：虽然主要是全文搜索引擎，但也支持向量搜索功能。

安装和配置

根据你选择的向量数据库，安装并配置相应的环境。例如，对于Milvus，你可以使用Docker来快速部署：

docker pull milvusdb/milvus:latest
docker run -d --name milvus_cpu -p 19530:19530 -p 19121:19121 -v /path/to/data:/var/lib/milvus milvusdb/milvus:latest

2. 数据准备

生成Embedding

使用预训练模型或自训练模型生成数据的embedding。例如，使用BERT生成文本的embedding：

from transformers import BertTokenizer, BertModel
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

def get_embedding(text):
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()

格式化数据

将生成的embedding和相关元数据（如ID、标签等）格式化为适合向量数据库的格式。例如：

data = [
    {"id": 1, "vector": get_embedding("这是一个示例文本")},
    {"id": 2, "vector": get_embedding("另一个示例文本")},
    # 更多数据...
]

3. 插入数据

连接到向量数据库

使用向量数据库的客户端库连接到数据库。例如，使用Milvus的Python SDK：

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

connections.connect("default", host="localhost", port="19530")

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=768)
]

schema = CollectionSchema(fields, "示例集合")
collection = Collection("example_collection", schema)

# 插入数据
entities = [
    [item["id"] for item in data],
    [item["vector"] for item in data]
]

collection.insert(entities)

4. 构建索引

为了提高搜索效率，需要构建索引。例如，在Milvus中：

index_params = {
    "index_type": "IVF_FLAT",
    "params": {"nlist": 1024},
    "metric_type": "L2"
}

collection.create_index(field_name="vector", index_params=index_params)

5. 执行相似性搜索

生成查询向量

使用同样的方法生成查询向量：

query_vector = get_embedding("查询文本")

发起查询

使用向量数据库的API发起相似性搜索：

search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
results = collection.search([query_vector], "vector", search_params, limit=5)

for result in results[0]:
    print(f"ID: {result.id}, 距离: {result.distance}")

6. 处理返回结果

根据返回的结果，可以进一步处理和展示。例如，从传统数据库中获取更多详细信息：

# 假设你有一个传统数据库，存储了更多的元数据
import sqlite3

conn = sqlite3.connect('example.db')
cursor = conn.cursor()

for result in results[0]:
    cursor.execute("SELECT * FROM items WHERE id=?", (result.id,))
    item = cursor.fetchone()
    print(item)