使用 MyScale 进行嵌入搜索

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

什么是向量数据库

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

为什么使用向量数据库

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

演示流程

演示流程如下：

• 设置：导入包并设置任何必需的变量
• 加载数据：加载数据集并使用 OpenAI 嵌入对其进行嵌入
• MyScale
  • 设置：设置 MyScale Python 客户端。有关更多详细信息，请访问此处
  • 索引数据：我们将创建一个表并为 content 建立索引。
  • 搜索数据：运行几个示例查询，并考虑各种目标。

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

设置

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

# 我们需要安装 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 的 Python 客户端库
import clickhouse_connect

# 我已将其设置为我们新的嵌入模型，可以将其更改为您选择的嵌入模型
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'

# 该文件约为 700 MB，因此需要一些时间
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()

# 将字符串中的向量读回列表
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 Console，您可以在两分钟内部署并执行集群上的向量搜索。

连接到 MyScale

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

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

索引数据

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

# 使用向量索引创建 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')}")

# 检查向量索引的状态，确保向量索引已准备就绪并处于“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: 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