使用MyScale进行嵌入式搜索

本笔记本将带您完成一个简单的流程，下载一些数据，对其进行嵌入，然后使用一些向量数据库对其进行索引和搜索。这是客户常见的需求，他们希望在安全环境中存储和搜索我们的嵌入，以支持生产用例，如聊天机器人、主题建模等。

什么是向量数据库

向量数据库是一种用于存储、管理和搜索嵌入向量的数据库。近年来，使用嵌入将非结构化数据（文本、音频、视频等）编码为向量，以供机器学习模型消费的做法已经蓬勃发展，这是由于人工智能在解决涉及自然语言、图像识别和其他非结构化数据形式的用例时日益有效。向量数据库已经成为企业提供和扩展这些用例的有效解决方案。

为什么使用向量数据库

向量数据库使企业能够利用我们在这个存储库中分享的许多嵌入用例（例如问答、聊天机器人和推荐服务），并在安全、可扩展的环境中使用它们。许多客户使用嵌入来解决其问题，但性能和安全性阻碍了它们投入生产 - 我们认为向量数据库是解决这一问题的关键组成部分，在本指南中，我们将介绍嵌入文本数据的基础知识，将其存储在向量数据库中，并将其用于语义搜索。

演示流程

演示流程如下： - 设置：导入包并设置任何必需的变量 - 加载数据：加载数据集并使用OpenAI嵌入对其进行嵌入 - MyScale - 设置：设置MyScale Python客户端。有关更多详细信息，请访问这里 - 索引数据：我们将创建一个表，并对其进行__内容__索引。 - 搜索数据：运行一些具有不同目标的示例查询。

完成本笔记后，您应该对如何设置和使用向量数据库有基本的了解，并可以继续进行更复杂的用例，利用我们的嵌入。

设置

导入所需的库并设置我们想要使用的嵌入模型。

# 我们需要安装MyScale客户端。
!pip install clickhouse-connect

#安装wget以下载zip文件
!pip install wget

import openai

from typing import List, Iterator
import pandas as pd
import numpy as np
import os
import wget
from ast import literal_eval

# MyScale's client library for Python
import clickhouse_connect

# I've set this to our new embeddings model, this can be changed to the embedding model of your choice
EMBEDDING_MODEL = "text-embedding-3-small"

# 忽略未关闭的SSL套接字警告 - 如果你遇到这些错误，可以选择忽略。
import warnings

warnings.filterwarnings(action="ignore", message="unclosed", category=ResourceWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning) 

加载数据

在这一部分，我们将加载之前准备好的嵌入数据。

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

# 文件大小约为700MB，因此需要一些时间来完成。
wget.download(embeddings_url)

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

article_df = pd.read_csv('../data/vector_database_wikipedia_articles_embedded.csv')

article_df.head()

	id	url	title	text	title_vector	content_vector	vector_id
0	1	https://simple.wikipedia.org/wiki/April	April	April is the fourth month of the year in the J...	[0.001009464613161981, -0.020700545981526375, ...	[-0.011253940872848034, -0.013491976074874401,...	0
1	2	https://simple.wikipedia.org/wiki/August	August	August (Aug.) is the eighth month of the year ...	[0.0009286514250561595, 0.000820168002974242, ...	[0.0003609954728744924, 0.007262262050062418, ...	1
2	6	https://simple.wikipedia.org/wiki/Art	Art	Art is a creative activity that expresses imag...	[0.003393713850528002, 0.0061537534929811954, ...	[-0.004959689453244209, 0.015772193670272827, ...	2
3	8	https://simple.wikipedia.org/wiki/A	A	A or a is the first letter of the English alph...	[0.0153952119871974, -0.013759135268628597, 0....	[0.024894846603274345, -0.022186409682035446, ...	3
4	9	https://simple.wikipedia.org/wiki/Air	Air	Air refers to the Earth's atmosphere. Air is a...	[0.02224554680287838, -0.02044147066771984, -0...	[0.021524671465158463, 0.018522677943110466, -...	4

# 从字符串中读取向量并将其转换回列表
article_df['title_vector'] = article_df.title_vector.apply(literal_eval)
article_df['content_vector'] = article_df.content_vector.apply(literal_eval)

# 将 vector_id 设置为一个字符串
article_df['vector_id'] = article_df['vector_id'].apply(str)

article_df.info(show_counts=True)

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 25000 entries, 0 to 24999
Data columns (total 7 columns):
 #   Column          Non-Null Count  Dtype 
---  ------          --------------  ----- 
 0   id              25000 non-null  int64 
 1   url             25000 non-null  object
 2   title           25000 non-null  object
 3   text            25000 non-null  object
 4   title_vector    25000 non-null  object
 5   content_vector  25000 non-null  object
 6   vector_id       25000 non-null  object
dtypes: int64(1), object(6)
memory usage: 1.3+ MB

MyScale

接下来我们将考虑的下一个向量数据库是MyScale。

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

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

连接到MyScale

按照连接详情部分的说明，从MyScale控制台中检索集群主机、用户名和密码信息，并使用这些信息创建到您的集群的连接，如下所示：

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

索引数据

我们将在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)

# 我们只需要部分列。
article_df = article_df[['id', 'url', 'title', 'text', 'content_vector']]

# 批量上传数据
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)

  0%|          | 0/250 [00:00<?, ?it/s]

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

# 检查插入数据的数量
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: InProgress

搜索数据

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

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