使用 MyScale 作为 OpenAI 嵌入的向量数据库

本笔记本提供了使用 MyScale 作为 OpenAI 嵌入的向量数据库的分步指南。该过程包括：

利用 OpenAI API 生成的预计算嵌入。
将这些嵌入存储在 MyScale 的云实例中。
使用 OpenAI API 将原始文本查询转换为嵌入。
利用 MyScale 在创建的集合中执行最近邻搜索。

什么是 MyScale

MyScale 是一个构建在 Clickhouse 之上的数据库，它结合了向量搜索和 SQL 分析，提供了高性能、简化的全托管体验。它旨在促进结构化数据和向量数据上的联合查询和分析，并为所有数据处理提供全面的 SQL 支持。

部署选项

通过使用 MyScale Console，在两分钟内部署并在集群上执行向量搜索和 SQL 分析。

先决条件

要遵循本指南，您需要具备以下条件：

通过遵循快速入门指南部署的 MyScale 集群。
clickhouse-connect 库，用于与 MyScale 进行交互。
OpenAI API 密钥，用于查询的向量化。

安装要求

此笔记本需要 openai、clickhouse-connect 以及其他一些依赖项。使用以下命令安装它们：

! pip install openai clickhouse-connect wget pandas

准备您的 OpenAI API 密钥

要使用 OpenAI API，您需要设置一个 API 密钥。如果您还没有，可以从 OpenAI 获取。

import openai

# 从 OpenAI 网站获取 API 密钥
openai.api_key = "OPENAI_API_KEY"

# 检查我们是否已通过身份验证
openai.Engine.list()

连接到 MyScale

请遵循连接详细信息部分，从 MyScale 控制台检索集群主机、用户名和密码信息，并使用它来创建到集群的连接，如下所示：

import clickhouse_connect

# 初始化客户端
client = clickhouse_connect.get_client(host='YOUR_CLUSTER_HOST', port=8443, username='YOUR_USERNAME', password='YOUR_CLUSTER_PASSWORD')

加载数据

我们需要加载 OpenAI 提供的预计算的维基百科文章向量嵌入数据集。使用 wget 包下载数据集。

import wget

embeddings_url = "https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip"

# 该文件约为 700 MB，因此需要一些时间
wget.download(embeddings_url)

下载完成后，使用 zipfile 包提取文件：

import zipfile

with zipfile.ZipFile("vector_database_wikipedia_articles_embedded.zip", "r") as zip_ref:
    zip_ref.extractall("../data")

现在，我们可以将数据从 vector_database_wikipedia_articles_embedded.csv 加载到 Pandas DataFrame 中：

import pandas as pd

from ast import literal_eval

# 从 csv 读取数据
article_df = pd.read_csv('../data/vector_database_wikipedia_articles_embedded.csv')
article_df = article_df[['id', 'url', 'title', 'text', 'content_vector']]

# 将向量从字符串读回列表
article_df["content_vector"] = article_df.content_vector.apply(literal_eval)
article_df.head()

索引数据

我们将在 MyScale 中创建一个名为 articles 的 SQL 表来存储嵌入数据。该表将包含一个具有余弦距离度量的向量索引，以及对嵌入长度的约束。使用以下代码创建表并将数据插入表中：

# 创建带有向量索引的文章表
embedding_len=len(article_df['content_vector'][0]) # 1536

client.command(f"""
CREATE TABLE IF NOT EXISTS default.articles
(
    id UInt64,
    url String,
    title String,
    text String,
    content_vector Array(Float32),
    CONSTRAINT cons_vector_len CHECK length(content_vector) = {embedding_len},
    VECTOR INDEX article_content_index content_vector TYPE HNSWFLAT('metric_type=Cosine')
)
ENGINE = MergeTree ORDER BY id
""")

# 分批插入数据到表中
from tqdm.auto import tqdm

batch_size = 100
total_records = len(article_df)

# 批量上传数据
data = article_df.to_records(index=False).tolist()
column_names = article_df.columns.tolist() 

for i in tqdm(range(0, total_records, batch_size)):
    i_end = min(i + batch_size, total_records)
    client.insert("default.articles", data[i:i_end], column_names=column_names)

我们需要在继续搜索之前检查向量索引的构建状态，因为它是在后台自动构建的。

# 检查插入数据的数量
print(f"articles count: {client.command('SELECT count(*) FROM default.articles')}")

# 检查向量索引的状态，确保向量索引已准备就绪并处于“Built”状态
get_index_status="SELECT status FROM system.vector_indices WHERE name='article_content_index'"
print(f"index build status: {client.command(get_index_status)}")

articles count: 25000
index build status: Built

搜索数据

一旦在 MyScale 中建立索引，我们就可以执行向量搜索来查找相似内容。首先，我们将使用 OpenAI API 为我们的查询生成嵌入。然后，我们将使用 MyScale 执行向量搜索。

import openai

query = "Famous battles in Scottish history"

# 从用户查询创建嵌入向量
embed = openai.Embedding.create(
    input=query,
    model="text-embedding-3-small",
)["data"][0]["embedding"]

# 查询数据库以查找与给定查询最相似的 K 个内容
top_k = 10
results = client.query(f"""
SELECT id, url, title, distance(content_vector, {embed}) as dist
FROM default.articles
ORDER BY dist
LIMIT {top_k}
""")

# 显示结果
for i, r in enumerate(results.named_results()):
    print(i+1, r['title'])

1 Battle of Bannockburn
2 Wars of Scottish Independence
3 1651
4 First War of Scottish Independence
5 Robert I of Scotland
6 841
7 1716
8 1314
9 1263
10 William Wallace