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

May 1, 2023
Open in Github

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

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

什么是MyScale

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

部署选项

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

安装要求

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

! pip install openai clickhouse-connect wget pandas
import openai

# get API key from on OpenAI website
openai.api_key = "OPENAI_API_KEY"

# check we have authenticated
openai.Engine.list()

连接到MyScale

按照连接详情部分从MyScale控制台获取集群主机、用户名和密码信息，并使用它创建到集群的连接，如下所示：

import clickhouse_connect

# initialize client
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"

# The file is ~700 MB so this will take some time
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

# read data from csv
article_df = pd.read_csv('../data/vector_database_wikipedia_articles_embedded.csv')
article_df = article_df[['id', 'url', 'title', 'text', 'content_vector']]

# read vectors from strings back into a list
article_df["content_vector"] = article_df.content_vector.apply(literal_eval)
article_df.head()

索引数据

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

# create articles table with vector index
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
""")

# insert data into the table in batches
from tqdm.auto import tqdm

batch_size = 100
total_records = len(article_df)

# upload data in batches
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)

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

# check count of inserted data
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"

# creates embedding vector from user query
embed = openai.Embedding.create(
    input=query,
    model="text-embedding-3-small",
)["data"][0]["embedding"]

# query the database to find the top K similar content to the given query
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}
""")

# display results
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