使用 Redis 进行嵌入搜索

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

什么是向量数据库

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

为什么使用向量数据库

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

演示流程

演示流程如下:

  • 设置:导入包并设置任何必需的变量
  • 加载数据:加载数据集并使用 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'

# 该文件约 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()
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,因此这应该对大多数人来说都很熟悉。

设置

有许多方法可以使用 RediSearch 部署 Redis。最简单的入门方法是使用 Docker,但也有许多潜在的部署选项。有关其他部署选项,请参阅此仓库中的 redis 目录(./redis)。

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

通过运行以下 docker 命令来启动一个包含 RediSearch(Redis Stack)的 Redis 版本

$ 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"                      # 向量的距离度量(例如 COSINE、IP、L2)
# 为数据集中每个列定义 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("索引已存在")
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"已加载 {redis_client.info()['db0']['keys']} 个文档到名为 {INDEX_NAME} 的 Redis 搜索索引中")
已加载 25000 个文档到名为 embeddings-index 的 Redis 搜索索引中

运行搜索查询

现在我们有了搜索索引并将文档加载到其中,我们可以运行搜索查询了。下面我们将提供一个运行搜索查询并返回结果的函数。使用此函数,我们将运行几个查询,展示如何将 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} (分数: {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. 现代艺术博物馆 (分数: 0.875)
1. 西欧 (分数: 0.867)
2. 文艺复兴艺术 (分数: 0.864)
3. 波普艺术 (分数: 0.86)
4. 北欧 (分数: 0.855)
5. 希腊化艺术 (分数: 0.853)
6. 现代主义文学 (分数: 0.847)
7. 艺术电影 (分数: 0.843)
8. 中欧 (分数: 0.843)
9. 欧洲 (分数: 0.841)
results = search_redis(redis_client, 'Famous battles in Scottish history', vector_field='content_vector', k=10)
0. 班诺克本战役 (分数: 0.869)
1. 苏格兰独立战争 (分数: 0.861)
2. 1651 (分数: 0.853)
3. 苏格兰独立战争第一阶段 (分数: 0.85)
4. 罗伯特一世 (苏格兰国王) (分数: 0.846)
5. 841 (分数: 0.844)
6. 1716 (分数: 0.844)
7. 1314 (分数: 0.837)
8. 1263 (分数: 0.836)
9. 威廉·华莱士 (分数: 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. 苏格兰独立战争第一阶段 (分数: 0.892)
1. 苏格兰独立战争 (分数: 0.889)
2. 苏格兰独立战争第二阶段 (分数: 0.879)
3. 苏格兰君主列表 (分数: 0.873)
4. 苏格兰边境区 (分数: 0.863)
# 运行混合查询,搜索标题向量中关于艺术的文章,并且只包含文本中包含“Leonardo da Vinci”一词的结果
results = search_redis(redis_client,
                       "Art",
                       vector_field="title_vector",
                       k=5,
                       hybrid_fields=create_hybrid_field("text", "Leonardo da Vinci")
                       )

# 查找我们全文搜索查询返回的文本中关于 Leonardo da Vinci 的具体提及
mention = [sentence for sentence in results[0].text.split("\n") if "Leonardo da Vinci" in sentence][0]
mention
0. 艺术 (分数: 1.0)
1. 绘画 (分数: 0.896)
2. 文艺复兴艺术 (分数: 0.88)
3. 绘画 (分数: 0.874)
4. 文艺复兴 (分数: 0.846)





'在欧洲,中世纪之后,出现了“文艺复兴”,意思是“重生”。人们重新发现了科学,艺术家们被允许描绘除宗教题材以外的主题。米开朗基罗和达·芬奇等艺术家仍然画宗教画,但他们现在也可以画神话画了。这些艺术家还发明了透视法,即远处的东西在画面中看起来更小。这很新颖,因为在中世纪,人们会把所有人物都画得很近,只是相互重叠。这些艺术家在他们的艺术中经常使用裸体。'

有关将 Redis 用作向量数据库的更多示例,请参阅本仓库 vector_databases/redis 目录中的 README 和示例。