AI 日报

使用 Redis 进行高效的向量相似性搜索:分步教程

  • By aihubon
  • Nov 08, 2023 - 2 min read



使用 Redis 进行高效的向量相似性搜索:分步教程

使用矢量嵌入和 Redis 增强搜索结果

搜索信息的能力在当今的数字环境中至关重要,用户希望几乎每个应用程序和网站都具有搜索功能。为了改进搜索结果,架构师和开发人员必须不断探索新的方法和架构。其中一种方法是利用深度学习模型生成的向量嵌入,这可以提高搜索结果的准确性和相关性。

为此,许多组织正在利用索引技术将其数据转换为向量空间。通过将数据表示为向量,可以执行返回最相关结果的相似性搜索。

在本教程中,我们将探索如何使用深度学习模型来创建向量嵌入,然后可以在 Redis 的帮助下对其进行索引以进行高效准确的搜索。通过对这种方法的透彻理解,架构师和开发人员可以更好地了解 AI 驱动的搜索功能的潜力,并找到改善用户搜索体验的最佳方法。

我们将经历以下过程:

为亚马逊产品数据集创建向量嵌入,使用 Redissearching for similar vectors 为它们建立索引。

我们还将探讨不同索引方法的优缺点以及如何使用它们来提高搜索性能。

让我们开始吧!

创建一个新目录并创建一个新的 Jupyter notebook。从此处获取数据集 CSV 文件。将其存储在./data/目录中。我们将使用 Python 3.8。在第一个单元格中安装以下依赖项:

%pip install redis numpy pandas%pip install -U sentence-transformers

安装依赖项后,您可以开始导入必要的库并定义任何必要的类或函数。在这种情况下,我们可以导入以下库并定义一个颜色类以备后用:

import randomimport numpy as npimport pandas as pdimport timefrom redis import Redisfrom redis.commands.search.field import VectorFieldfrom redis.commands.search.field import TextFieldfrom redis.commands.search.field import TagFieldfrom redis.commands.search.query import Queryfrom redis.commands.search.result import Resultclass color:   PURPLE = ''   CYAN = ''   DARKCYAN = ''   BLUE = ''   GREEN = ''   YELLOW = ''   RED = ''   BOLD = ''   UNDERLINE = ''   END = ''

导入 Redis 库是为了与 Redis 交互,Redis 是一种内存中的数据结构存储,通常用作数据库、缓存和消息代理。redis.commands.search.field我们还从和模块中导入以下类redis.commands.search.query

VectorField:用于表示Redis中的向量场,比如embeddings。TextField: 用于表示 Redis 中的文本字段。TagField: 用于表示 Redis 中的标签字段。Query:用于为 Redis 创建搜索查询。Result:用于表示Redis返回的搜索结果。

此外,我们定义了一个颜色类,可用于将彩色文本打印到控制台。颜色类有几个属性,如 PURPLE、CYAN、BOLD 等,可用于对控制台中的文本输出进行着色。

下一步涉及将亚马逊产品数据加载到 Pandas DataFrame 中,并将长文本字段截断为最大长度 512 个字符,这是我们稍后将使用的预训练句子嵌入生成器支持的最大长度。

下面是加载产品数据和截断长文本字段的单元格代码:

MAX_TEXT_LENGTH = 512NUMBER_PRODUCTS = 1000# define a function to auto-truncate long text fieldsdef auto_truncate(val):    return val[:MAX_TEXT_LENGTH]# load the product data and truncate long text fieldsall_prods_df = pd.read_csv("data/product_data.csv", converters={'bullet_point': auto_truncate, 'item_keywords': auto_truncate, 'item_name': auto_truncate})all_prods_df['primary_key'] = all_prods_df['item_id'] + '-' + all_prods_df['domain_name']all_prods_df['item_keywords'].replace('', np.nan, inplace=True)all_prods_df.dropna(subset=['item_keywords'], inplace=True)all_prods_df.reset_index(drop=True, inplace=True)# get the first 1000 products with non-empty item keywordsproduct_metadata = all_prods_df.head(NUMBER_PRODUCTS).to_dict(orient='index')

该代码从 CSV 文件中加载产品数据,截断长文本字段,在 DataFrame 中添加一个新的 primary_key 列,过滤掉任何没有关键字的产品,并提取前 1000 个具有非空项目关键字的产品的元数据并存储它在名为 product_metadata 的字典中。

让我们看一下产品数据的前几行:

all_prods_df.head()
使用 Redis 进行高效的向量相似性搜索:分步教程

连接到 Redis

将产品数据加载到 Pandas DataFrame 并提取 1000 个产品的元数据后,下一步是连接到 Redis。我们将在其云中使用 RedisLabs 提供的 Redis 实例,该实例提供免费套餐。您可以在 redis.com/try-free/ 注册一个免费帐户。

启动一个新的 Redis 实例并复制连接详细信息。连接到 Redis 实例时需要提供密码。您可以在连接详细信息页面中找到密码。密码与默认用户的密码相同。

redis_conn = Redis(  host='XXX',  port=XXXXX,  password="")print ('Connected to redis')

创建嵌入

使用 SentenceTransformer

下一步涉及使用名为 distilroberta-v1 的预训练 Sentence Transformer 模型为项目关键字生成嵌入(向量)。该模型可从 Sentence Transformer 库中获得。我们将使用 SentenceTransformer 类加载模型并生成嵌入。

from sentence_transformers import SentenceTransformermodel = SentenceTransformer('sentence-transformers/all-distilroberta-v1')%%timeitem_keywords =  [product_metadata[i]['item_keywords']  for i in product_metadata.keys()]item_keywords_vectors = [ model.encode(sentence) for sentence in item_keywords]

然后我们将检查嵌入的维度:

len(item_keywords_vectors)len(product_metadata)# Check one of the productsproduct_metadata[0]

准备效用函数

现在我们有 1000 个产品的 1000 个嵌入。接下来,我们将定义 3 个效用函数。一个用于加载产品数据,两个用于在 Vector 字段上创建索引。

def load_vectors(client:Redis, product_metadata, vector_dict, vector_field_name):    p = client.pipeline(transaction=False)    for index in product_metadata.keys():            #hash key        key='product:'+ str(index)+ ':' + product_metadata[index]['primary_key']                #hash values        item_metadata = product_metadata[index]        item_keywords_vector = vector_dict[index].astype(np.float32).tobytes()        item_metadata[vector_field_name]=item_keywords_vector                # HSET        p.hset(key,mapping=item_metadata)                p.execute()def create_flat_index (redis_conn,vector_field_name,number_of_vectors, vector_dimensions=512, distance_metric='L2'):    redis_conn.ft().create_index([        VectorField(vector_field_name, "FLAT", {"TYPE": "FLOAT32", "DIM": vector_dimensions, "DISTANCE_METRIC": distance_metric, "INITIAL_CAP": number_of_vectors, "BLOCK_SIZE":number_of_vectors }),        TagField("product_type"),        TextField("item_name"),        TextField("item_keywords"),        TagField("country")            ])def create_hnsw_index (redis_conn,vector_field_name,number_of_vectors, vector_dimensions=512, distance_metric='L2',M=40,EF=200):    redis_conn.ft().create_index([        VectorField(vector_field_name, "HNSW", {"TYPE": "FLOAT32", "DIM": vector_dimensions, "DISTANCE_METRIC": distance_metric, "INITIAL_CAP": number_of_vectors, "M": M, "EF_CONSTRUCTION": EF}),        TagField("product_type"),        TextField("item_keywords"),                TextField("item_name"),        TagField("country")         ])    

平面索引和 HNSW 都是在高维空间进行近似最近邻搜索的方法,但是它们在构造和搜索索引的方式上有所不同。

平面索引是一种直接的方法,其中所有数据点都被索引并存储在单个列表或树结构中。为了找到查询点的最近邻居,通过计算查询点与索引中所有其他点之间的距离来进行强力搜索。虽然简单且易于实现,但对于大型数据集和高维空间,平面索引的计算成本很高且不切实际。

另一方面,HNSW 代表 Hierarchical Navigable Small World,是一种更复杂的索引算法,可将数据点组织成层次图结构。该图是通过将每个点连接到其最近的邻居,然后以分层方式递归连接附近的点来构建的。这导致了一个高度聚集的图结构,并通过仅探索图的一小部分来实现快速近似最近邻搜索。

HNSW 相对于平面索引的主要优势在于它能够扩展到大型数据集和高维空间。然而,HNSW 需要更仔细地调整其参数,并且与平面索引相比,索引构建时间可能更长。

索引并查询新数据

接下来,我们将首先使用平面索引加载和索引产品数据。

%%timeITEM_KEYWORD_EMBEDDING_FIELD='item_keyword_vector'TEXT_EMBEDDING_DIMENSION=768NUMBER_PRODUCTS=1000print ('Loading and Indexing + ' +  str(NUMBER_PRODUCTS) + ' products')#flush all dataredis_conn.flushall()#create flat index & load vectorscreate_flat_index(redis_conn, ITEM_KEYWORD_EMBEDDING_FIELD,NUMBER_PRODUCTS,TEXT_EMBEDDING_DIMENSION,'COSINE')load_vectors(redis_conn,product_metadata,item_keywords_vectors,ITEM_KEYWORD_EMBEDDING_FIELD)

查询我们的 FLAT 索引

在此之后,我们可以查询我们的索引并搜索给定查询向量的前 5 (topK) 个最近邻居:

%%timetopK=5product_query='beautifully crafted present for her. a special occasion'#product_query='cool way to pimp up my cell'#vectorize the queryquery_vector = model.encode(product_query).astype(np.float32).tobytes()#prepare the queryq = Query(f'*=>[KNN {topK} @{ITEM_KEYWORD_EMBEDDING_FIELD} $vec_param AS vector_score]').sort_by('vector_score').paging(0,topK).return_fields('vector_score','item_name','item_id','item_keywords').dialect(2)params_dict = {"vec_param": query_vector}#Execute the queryresults = redis_conn.ft().search(q, query_params = params_dict)#Print similar products foundfor product in results.docs:    print ('***************Product  found ************')    print (color.BOLD + 'hash key = ' +  color.END + product.id)    print (color.YELLOW + 'Item Name = ' +  color.END  + product.item_name)    print (color.YELLOW + 'Item Id = ' +  color.END  + product.item_id)    print (color.YELLOW + 'Item keywords = ' +  color.END  + product.item_keywords)    print (color.YELLOW + 'Score = ' +  color.END  + product.vector_score)
***************Product  found ************hash key = product:597:B076ZYG35R-amazon.comItem Name = Amazon Brand - The Fix Women's Lizzie Block Heel Ruffled Sandal Heeled, Dove Suede, 9 B USItem Id = B076ZYG35RItem keywords = gifts her zapatos shoe ladies mujer womans designer spring summer date night dressy fancy high heelsScore = 0.498145997524***************Product  found ************hash key = product:112:B07716JGFN-amazon.comItem Name = Amazon Brand - The Fix Women's Jackelyn Kitten Heel Bow Sandal HeeledItem Id = B07716JGFNItem keywords = zapatos shoe ladies mujer womans spring summer casual date night gifts for herScore = 0.613550662994***************Product  found ************hash key = product:838:B0746M8ZY9-amazon.comItem Name = Amazon Brand - 206 Collective Women's Roosevelt Shearling Slide Slipper Shoe, Chestnut, 7.5 B USItem Id = B0746M8ZY9Item keywords = zapatos shoe para de ladies mujer womans mocasines designer  clothing  work wear office top gifts gifts for gifts for herScore = 0.623450696468***************Product  found ************hash key = product:800:B0065HI4KA-amazon.comItem Name = Sterling Silver Ribbon Black and White Diamond Earrings (1/5 cttw, I-J Color, I2-I3 Clarity)Item Id = B0065HI4KAItem keywords = classics with a twist braid twist rope color diamonds black diamonds valentines day gifts for her valentines day jewelry valentines day love gifts for women valentine jewelry valentines jewelry gifts for her braided earrings rope earrings black diamond earrings color diamond earrings braid twist rope color diamonds black diamonds valentines day gifts for her valentines day jewelry valentines day love gifts for women valentine jewelry valentines jewelry gifts for her braided earrings rope earrings black diamScore = 0.628378272057***************Product  found ************hash key = product:475:B07DJ2M9YR-amazon.comItem Name = Essentials Women's Pom Knit Hat and Scarf SetItem Id = 0mB07DJ2M9YRItem keywords = women woman ladies winter gift cozyScore = 0.644373476505

查询我们的 HNSW 指数

它的工作!那太棒了!现在让我们对 HNSW 索引进行同样的尝试。首先加载和索引产品数据。

%%timeprint ('Loading and Indexing + ' +  str(NUMBER_PRODUCTS) + ' products')ITEM_KEYWORD_EMBEDDING_FIELD='item_keyword_vector'NUMBER_PRODUCTS=1000TEXT_EMBEDDING_DIMENSION=768#flush all dataredis_conn.flushall()#create flat index & load vectorscreate_hnsw_index(redis_conn, ITEM_KEYWORD_EMBEDDING_FIELD,NUMBER_PRODUCTS,TEXT_EMBEDDING_DIMENSION,'COSINE',M=40,EF=200)load_vectors(redis_conn,product_metadata,item_keywords_vectors,ITEM_KEYWORD_EMBEDDING_FIELD)

现在我们可以再次查询:

%%timetopK=5product_query='beautifully crafted present for her. a special occasion'#product_query='cool way to pimp up my cell'#vectorize the queryquery_vector = model.encode(product_query).astype(np.float32).tobytes()#prepare the queryq = Query(f'*=>[KNN {topK} @{ITEM_KEYWORD_EMBEDDING_FIELD} $vec_param AS vector_score]').sort_by('vector_score').paging(0,topK).return_fields('vector_score','item_name','item_id','item_keywords','country').dialect(2)params_dict = {"vec_param": query_vector}#Execute the queryresults = redis_conn.ft().search(q, query_params = params_dict)#Print similar products foundfor product in results.docs:    print ('***************Product  found ************')    print (color.BOLD + 'hash key = ' +  color.END + product.id)    print (color.YELLOW + 'Item Name = ' +  color.END  + product.item_name)    print (color.YELLOW + 'Item Id = ' +  color.END  + product.item_id)    print (color.YELLOW + 'Item keywords = ' +  color.END  + product.item_keywords)    print (color.YELLOW + 'Country = ' +  color.END  + product.country)    print (color.YELLOW + 'Score = ' +  color.END  + product.vector_score)

了解更多!

您可以在这个 Github 存储库中找到完整的代码。我们强烈建议查看此存储库以获取更多信息以及有关如何进行相似性搜索的指南,就像我们在此处对图像所做的那样。

谢谢你!如果您喜欢本教程,您可以在我们的教程页面上找到更多信息并继续阅读。如果您想立即检查自己学到了什么,我强烈建议您参加我们的 AI 编程马拉松之一。AI未来百科 ; 探索AI的边界与未来! 懂您的AI未来站