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

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

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

什么是MyScale

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

部署选项

• 通过使用MyScale控制台，在两分钟内在您的集群上部署和执行带有SQL的向量搜索。

先决条件

要按照本指南操作，您需要具备以下条件：

1. 通过按照快速入门指南部署的 MyScale 集群。
2. ‘clickhouse-connect’ 库，用于与 MyScale 进行交互。
3. 用于查询向量化的 OpenAI API 密钥。

安装要求

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

! pip install openai clickhouse-connect wget pandas

准备你的OpenAI API密钥

要使用OpenAI API，你需要设置一个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"

# 文件大小约为700MB，因此需要一些时间来完成。
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表，用于存储嵌入数据。该表将包括一个带有余弦距离度量的向量索引，并且会有一个用于嵌入长度的约束。使用以下代码来创建并插入数据到articles表中：

# 创建带有向量索引的文章表
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')}")

# check the status of the vector index, make sure vector index is ready with 'Built' status
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