跳到主要内容

使用Redis进行嵌入式搜索

nbviewer

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

什么是向量数据库

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

为什么使用向量数据库

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

演示流程

演示流程如下: - 设置:导入包并设置任何必需的变量 - 加载数据:加载数据集并使用OpenAI嵌入对其进行嵌入 - Redis - 设置:设置Redis-Py客户端。更多详情请查看这里 - 索引数据:为向量搜索和混合搜索(向量 + 全文搜索)在所有可用字段上创建搜索索引。 - 搜索数据:运行一些示例查询,以达到不同的目标。

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

设置

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

# 我们需要安装Redis客户端。
!pip install redis

#安装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

# Python 的 Redis 客户端库
import redis

# 我已将其设置为我们新的嵌入模型,您可以根据需要更改为其他嵌入模型。
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

Redis

本教程中涵盖的下一个向量数据库是 Redis。您很可能已经了解Redis。您可能不知道的是RediSearch模块。多年来,企业一直在各大云提供商、Redis Cloud和本地使用Redis与RediSearch模块。最近,Redis团队在这个模块中添加了向量存储和搜索功能,除了RediSearch已经具有的功能。

考虑到Redis周围庞大的生态系统,很可能有您需要的语言的客户端库。您可以使用任何标准的Redis客户端库来运行RediSearch命令,但最简单的方法是使用封装了RediSearch API的库。下面是一些示例,但您可以在这里找到更多客户端库。

项目 语言 许可证 作者 星数
jedis Java MIT Redis Stars
redis-py Python MIT Redis Stars
node-redis Node.js MIT Redis Stars
nredisstack .NET MIT Redis Stars
redisearch-go Go BSD Redis redisearch-go-stars
redisearch-api-rs Rust BSD Redis redisearch-api-rs-stars

在下面的单元格中,我们将指导您如何将Redis用作向量数据库。由于您中的许多人很可能已经习惯了Redis API,这对大多数人来说应该是熟悉的。

设置

有许多部署 Redis 与 RediSearch 的方法。开始的最简单方法是使用 Docker,但是还有许多其他部署选项。有关其他部署选项,请参阅本存储库中的 redis 目录

在本教程中,我们将使用 Docker 上的 Redis Stack。

通过运行以下 docker 命令启动带有 RediSearch 的 Redis 版本(Redis Stack)

$ cd redis
$ docker compose up -d

这还包括用于管理 Redis 数据库的 RedisInsight GUI,您可以在启动 docker 容器后在 http://localhost:8001 查看。

您已经设置好并准备好开始了!接下来,我们将导入并创建用于与我们刚刚创建的 Redis 数据库通信的客户端。

import redis
from redis.commands.search.indexDefinition import (
IndexDefinition,
IndexType
)
from redis.commands.search.query import Query
from redis.commands.search.field import (
TextField,
VectorField
)

REDIS_HOST = "localhost"
REDIS_PORT = 6379
REDIS_PASSWORD = "" # Redis 无密码默认设置

# 连接到 Redis
redis_client = redis.Redis(
host=REDIS_HOST,
port=REDIS_PORT,
password=REDIS_PASSWORD
)
redis_client.ping()

True

创建搜索索引

下面的单元格将展示如何在Redis中指定和创建一个搜索索引。我们将:

  1. 设置一些常量来定义我们的索引,比如距离度量和索引名称
  2. 使用RediSearch字段定义索引模式
  3. 创建索引
# 常量
VECTOR_DIM = len(article_df['title_vector'][0]) # 向量的长度
VECTOR_NUMBER = len(article_df) # 初始向量数量
INDEX_NAME = "embeddings-index" # 搜索索引的名称
PREFIX = "doc" # 文档键的前缀
DISTANCE_METRIC = "COSINE" # 向量间的距离度量方法(例如:余弦距离、内积、欧氏距离)

# 为数据集中的每一列定义 RediSearch 字段
title = TextField(name="title")
url = TextField(name="url")
text = TextField(name="text")
title_embedding = VectorField("title_vector",
"FLAT", {
"TYPE": "FLOAT32",
"DIM": VECTOR_DIM,
"DISTANCE_METRIC": DISTANCE_METRIC,
"INITIAL_CAP": VECTOR_NUMBER,
}
)
text_embedding = VectorField("content_vector",
"FLAT", {
"TYPE": "FLOAT32",
"DIM": VECTOR_DIM,
"DISTANCE_METRIC": DISTANCE_METRIC,
"INITIAL_CAP": VECTOR_NUMBER,
}
)
fields = [title, url, text, title_embedding, text_embedding]

# 检查索引是否存在
try:
redis_client.ft(INDEX_NAME).info()
print("Index already exists")
except:
# 创建RediSearch索引
redis_client.ft(INDEX_NAME).create_index(
fields = fields,
definition = IndexDefinition(prefix=[PREFIX], index_type=IndexType.HASH)
)

将文档加载到索引中

现在我们有了一个搜索索引,我们可以将文档加载到其中。我们将使用在之前的示例中使用过的相同文档。在Redis中,可以使用Hash或JSON(如果除了RediSearch还使用RedisJSON)数据类型来存储文档。在本示例中,我们将使用HASH数据类型。下面的单元格将展示如何将文档加载到索引中。

def index_documents(client: redis.Redis, prefix: str, documents: pd.DataFrame):
records = documents.to_dict("records")
for doc in records:
key = f"{prefix}:{str(doc['id'])}"

# 为标题和内容创建字节向量
title_embedding = np.array(doc["title_vector"], dtype=np.float32).tobytes()
content_embedding = np.array(doc["content_vector"], dtype=np.float32).tobytes()

# 将浮点数列表替换为字节向量
doc["title_vector"] = title_embedding
doc["content_vector"] = content_embedding

client.hset(key, mapping = doc)

index_documents(redis_client, PREFIX, article_df)
print(f"Loaded {redis_client.info()['db0']['keys']} documents in Redis search index with name: {INDEX_NAME}")

Loaded 25000 documents in Redis search index with name: embeddings-index

运行搜索查询

现在我们已经有了一个搜索索引并加载了文档,我们可以运行搜索查询。下面我们将提供一个函数,该函数将运行一个搜索查询并返回结果。使用这个函数,我们运行一些查询,展示如何利用Redis作为向量数据库。每个示例将演示在开发使用Redis的搜索应用程序时需要牢记的特定功能。

  1. 返回字段:您可以指定要在搜索结果中返回哪些字段。如果您只想返回文档中的字段子集,并且不需要单独调用来检索文档,则这将非常有用。在下面的示例中,我们将只返回搜索结果中的title字段。
  2. 混合搜索:您可以将向量搜索与任何其他RediSearch字段结合使用,实现混合搜索,如全文搜索、标签、地理和数字。在下面的示例中,我们将向量搜索与全文搜索结合使用。
def search_redis(
redis_client: redis.Redis,
user_query: str,
index_name: str = "embeddings-index",
vector_field: str = "title_vector",
return_fields: list = ["title", "url", "text", "vector_score"],
hybrid_fields = "*",
k: int = 20,
) -> List[dict]:

# 从用户查询生成嵌入向量
embedded_query = openai.Embedding.create(input=user_query,
model=EMBEDDING_MODEL,
)["data"][0]['embedding']

# 准备查询
base_query = f'{hybrid_fields}=>[KNN {k} @{vector_field} $vector AS vector_score]'
query = (
Query(base_query)
.return_fields(*return_fields)
.sort_by("vector_score")
.paging(0, k)
.dialect(2)
)
params_dict = {"vector": np.array(embedded_query).astype(dtype=np.float32).tobytes()}

# 执行向量搜索
results = redis_client.ft(index_name).search(query, params_dict)
for i, article in enumerate(results.docs):
score = 1 - float(article.vector_score)
print(f"{i}. {article.title} (Score: {round(score ,3) })")
return results.docs

# 使用OpenAI生成查询嵌入
openai.api_key = os.getenv("OPENAI_API_KEY", "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
results = search_redis(redis_client, 'modern art in Europe', k=10)

0. Museum of Modern Art (Score: 0.875)
1. Western Europe (Score: 0.867)
2. Renaissance art (Score: 0.864)
3. Pop art (Score: 0.86)
4. Northern Europe (Score: 0.855)
5. Hellenistic art (Score: 0.853)
6. Modernist literature (Score: 0.847)
7. Art film (Score: 0.843)
8. Central Europe (Score: 0.843)
9. European (Score: 0.841)
results = search_redis(redis_client, 'Famous battles in Scottish history', vector_field='content_vector', k=10)

0. Battle of Bannockburn (Score: 0.869)
1. Wars of Scottish Independence (Score: 0.861)
2. 1651 (Score: 0.853)
3. First War of Scottish Independence (Score: 0.85)
4. Robert I of Scotland (Score: 0.846)
5. 841 (Score: 0.844)
6. 1716 (Score: 0.844)
7. 1314 (Score: 0.837)
8. 1263 (Score: 0.836)
9. William Wallace (Score: 0.835)

使用Redis进行混合查询

前面的示例展示了如何使用RediSearch运行向量搜索查询。在本节中,我们将展示如何将向量搜索与其他RediSearch字段结合起来进行混合搜索。在下面的示例中,我们将结合向量搜索和全文搜索。

def create_hybrid_field(field_name: str, value: str) -> str:
return f'@{field_name}:"{value}"'

# 在内容向量中搜索关于苏格兰历史上著名战役的文章,并仅包括标题中含有“Scottish”的结果。
results = search_redis(redis_client,
"Famous battles in Scottish history",
vector_field="title_vector",
k=5,
hybrid_fields=create_hybrid_field("title", "Scottish")
)

0. First War of Scottish Independence (Score: 0.892)
1. Wars of Scottish Independence (Score: 0.889)
2. Second War of Scottish Independence (Score: 0.879)
3. List of Scottish monarchs (Score: 0.873)
4. Scottish Borders (Score: 0.863)
# run a hybrid query for articles about Art in the title vector and only include results with the phrase "Leonardo da Vinci" in the text
results = search_redis(redis_client,
"Art",
vector_field="title_vector",
k=5,
hybrid_fields=create_hybrid_field("text", "Leonardo da Vinci")
)

# 在我们的全文搜索查询返回的文本中,找到了关于莱昂纳多·达·芬奇的具体提及。
mention = [sentence for sentence in results[0].text.split("\n") if "Leonardo da Vinci" in sentence][0]
mention

0. Art (Score: 1.0)
1. Paint (Score: 0.896)
2. Renaissance art (Score: 0.88)
3. Painting (Score: 0.874)
4. Renaissance (Score: 0.846)
'In Europe, after the Middle Ages, there was a "Renaissance" which means "rebirth". People rediscovered science and artists were allowed to paint subjects other than religious subjects. People like Michelangelo and Leonardo da Vinci still painted religious pictures, but they also now could paint mythological pictures too. These artists also invented perspective where things in the distance look smaller in the picture. This was new because in the Middle Ages people would paint all the figures close up and just overlapping each other. These artists used nudity regularly in their art.'

要了解更多关于将Redis用作向量数据库的示例,请参阅此存储库的vector_databases/redis目录中的README和示例。