为了账号安全,请及时绑定邮箱和手机立即绑定

从卷轴到相似性搜索:用DuckDB VSS构建电影推荐系统

学习如何使用DuckDB的向量相似性搜索插件来构建一个利用语义搜索和Gemini嵌入的推荐系统

从古卷轴到DuckDB的相似性搜索,来源:由Adobe Firefly生成

从古时的卷轴到如今的数字电影的演变

你站在古埃及,大约在我们现代日历开始的250年前。当你的眼睛适应了从高窗照进的温暖光线时,你发现自己身处亚历山大图书馆中。眼前的景象令人震撼:书架向四面八方延伸,存放了将近五十万卷的卷轴。空气中弥漫着纸草的气味,你在寻找关于天文学的文献。你听到卡利马科斯轻轻翻动卷轴的声音,他正在研究解决你问题的方法,他是这座图书馆最著名的学者。

他不再仅仅按作者或标题来组织卷轴,而是开创了一套按主题和内容分类的系统,使学者发现相关作品。这些作品可能永远无法通过简单的字母排序找到。他通过创建Pinakes——世界上第一个图书馆目录系统——革新了信息检索。

卡利马丘斯的插图,来源:使用 DALL-E 3 生成

两千多年后,我们面对着一个类似的挑战。在Netflix早期的数字档案中,工程师们面临了一个类似亚历山大图书馆规模的问题。他们基于简单的评分匹配的电影推荐系统,难以捕捉电影之间真正相似的本质。一部关于婚礼的喜剧可能与浪漫剧情片有更多的共同点,而却可能与另一部关于体育的喜剧没有太多共同之处。传统分类方法——类似于根据卷轴的物理特征来组织卷轴——忽略了这些微妙的联系。语义理解的挑战——捕捉项目之间真正的含义和相似性——依然是现代搜索引擎和推荐系统的核心挑战。

奈飞个性化简史:从1998年创业至今,奈飞在个性化策略、指标和实验方面的详细历史…gibsonbiddle.medium.com

今天,我们站在一个有意思的岔路口。因为向量搜索在增强检索生成(RAG)系统中的广泛应用,它们承诺能解决这些语义匹配的问题,但有个问题:但大多数解决方案需要复杂的基础设施、大量的资源和精心的维护。但在某些情况下,有一个实用的解决方案可以应对这种复杂性:DuckDB的向量相似性搜索(VSS)扩展

在本指南中,我们将构建一个电影推荐引擎,这个引擎可以解决Netflix早期遇到的挑战,使用现代工具,这些工具可以轻松运行在您的笔记本电脑上。通过结合DuckDB的VSS扩展功能和Gemini的嵌入能力,我们将创建一个不仅理解电影元数据,更能理解电影本质的系统。无论您是构建下一代推荐引擎,还是仅仅想更好地了解向量搜索,这次实践之旅将帮助您掌握在自己的项目中应对语义搜索挑战所需的知识。

文本转数字:理解词嵌入

在我们开始探讨相似性搜索之前,先来了解一下我们是如何将电影描述转换为计算机可以处理的数值的。这就需要用到embeddings这一概念。

嵌入向量通过将文本、图像和视频转换为浮点数数组(即向量)来工作。这些向量旨在捕捉文本、图像和视频的意义和特征。嵌入向量的长度也被称为向量的维度。例如,一段文本可能被一个包含数百个维度的向量来表示。

向量相似搜索

一旦我们得到这些数值向量,我们就需要方法来衡量它们的接近程度和相似性。

在 v0.10.0 版本中,DuckDB 引入了 ARRAY 数据类型,它存储固定大小的数组,以补充可变大小的 LIST 类型。

他们还为这种新的_ARRAY_类型增加了一些距离函数:array_distancearray_negative_inner_productarray_cosine_distance。借助这些距离函数,可以测量相似度。

欧氏距离,来源:作者(原文)

余弦相似度 :来源:作者(原版)

DuckDB的VSS扩展之后添加了对分级可导航小世界(HNSW )索引的支持,以加快向量相似度搜索。

HNSW:了解小世界

让我们来了解一下HNSW索引的工作原理。想象你在纽约市想找一个既玩《魔兽世界》又教量子物理的人。以下是不同的搜索方法会是如何的:

策略1:蛮力法:

在纽约市里拦住每一个路人,问他们符不符合你的条件。

  • 时间:几个月
  • 准确性:100%
  • 效率:多余的数百万次检查
策略2:HNSW(智能分级搜索)

可以把它想象成一个巧妙组织的社会网络,比如有着不同的连接层次:

三级(顶级)——全球链接

  • 像认识主要游戏社区和大学物理系的负责人
  • 快速广泛的覆盖范围:“这里是一些主要机构和游戏中心,需要留意”

第二级(中级)——区域联系

  • 像认识当地《魔兽世界》公会的领袖和物理系系主任一样
  • 更有针对性:这三个物理系不仅有活跃的游戏爱好者群体,就像我们认识当地《魔兽世界》公会的领袖和物理系系主任一样。

第一层(地面层)——本地联系

  • 直接了解个别玩家和教授
  • 精确匹配合适的人选

HNSW 可视化图,来源:DALL-E 3 生成

搜索从最高层级开始,迅速锁定有潜力的区域,然后高效深入。对于100万个候选对象,只需检查大约20个,而不是全部,同时保持95%到99%的准确度。

构建电影推荐器
先决条件
开始设置
以下是一个Python代码示例,用于初始化谷歌AI平台的认证信息和配置。

from typing import List, Dict  
import duckdb  
import httpx  
from google.cloud import aiplatform  
from vertexai.language_models import TextEmbeddingModel  

tmdb_api_key: str = 'your-tmdb-api-key'  
# 服务账号凭据从JSON文件中读取
credentials = service_account.Credentials.from_service_account_file('your-sa.json')  
# 初始化AI平台,指定项目ID、位置和凭据
aiplatform.init(project='your-project', location='us-central1', credentials=credentials)
抓取电影数据

我们使用httpx从TMDB API获取电影数据。同时,我们可以通过设置最低平均评分和评分人数来缩小数据集规模,只保留更知名的电影。

    def _get_movies(page: int, vote_avg_min: float, vote_count_min: float) -> List[Dict]:  
        """ 从TMDB API获取电影数据 """  
        response = httpx.get('https://api.themoviedb.org/3/discover/movie', headers={  
            'Authorization': f'Bearer {tmdb_api_key}'  
        }, params={  
            'sort_by': 'popularity.desc',  
            'include_adult': 'false',  
            'include_video': 'false',  
            'language': 'en-US',  
            'with_original_language': 'en',  
            'vote_average.gte': vote_avg_min,  
            'vote_count.gte': vote_count_min,  
            'page': page  
        })  
        response.raise_for_status()  # 如果响应状态不好则抛出异常  
        return response.json()['results']  

    def get_movies(pages: int, 最低评分: float, 最低票数: float) -> List[Dict]:  
        """ 生成器,用于从多页中获取电影数据 """  
        for page in range(1, pages + 1):  
            yield from _get_movies(page, 最低评分, 最低票数)
生成嵌入表示

我们使用Gemini的text-embedding-004模型来生成嵌入向量,并将维度设置为256。

确保在生成嵌入和创建DuckDB表时,维度大小(256)要保持一致。

def 嵌入文本内容(文本列表: List[str]) -> List[List[float]]:
    """使用Gemini生成文本列表的嵌入。返回一个包含嵌入向量的列表"""
    模型 = TextEmbeddingModel.from_pretrained('text-embedding-004')
    输入数据 = [TextEmbeddingInput(文本, 'RETRIEVAL_DOCUMENT') for 文本 in 文本列表]
    嵌入 = 模型.get_embeddings(输入数据, output_dimensionality=256)
    return [嵌入向量.values for 嵌入向量 in 嵌入]

# 从TMDB API获取电影数据,并为这些电影生成嵌入
电影数据 = list(获取电影(3, 6.0, 1000))
待嵌入的电影数据 = [(电影['id'], 电影['title'], 电影['overview']) for 电影 in 电影数据]
电影嵌入 = 嵌入文本内容([概述 for _, _, 概述 in 待嵌入的电影数据])
DuckDB VSS 配置

接下来,在DuckDB中安装并启用VSS扩展,并启用持久化。这让我们可以将嵌入存储到数据库文件里。

    安装 vss;  
    加载 vss;  
    设置 hnsw_enable_experimental_persistence 为 true;

然后我们就按照之前的维度创建了表格。

    创建一个表 movies_vectors (  
        id 整型,  
        title 文本,  
        vector 浮点型数组[256]  
    )

插入嵌入之后,我们会在向量列上创建一个HNSW索引,从而加快向量相似度搜索的速度。

创建名为 movies_vector_index 的索引,在 movies_vectors 上使用 HNSW 方法 (向量)

我们接着准备一个函数,它接受一部电影的描述作为输入。这是用户的搜索查询。我们根据这个输入创建一个嵌入向量。最后,我们使用DuckDB的相似度函数来找到相似的电影。

    SELECT 电影名  
    FROM movies_vectors  
    ORDER BY array_distance(vector, array[{vector_array}]::FLOAT[256])  
    LIMIT 3

就这样,这是DuckDB VSS配置的样子。

    # 设置DuckDB以启用向量相似性搜索(VSS)扩展和持久性功能  
    # 参见:https://duckdb.org/docs/extensions/vss.html  
    with duckdb.connect(database='movies.duckdb') as conn:  
        conn.execute("""  
            INSTALL vss;  
            LOAD vss;  
            SET hnsw_enable_experimental_persistence = true;  
        """)  

        conn.execute("""  
            CREATE TABLE movies_vectors (  
                id INTEGER,  
                title VARCHAR,  
                vector FLOAT[256]  
            )  
        """)  

        # 将嵌入插入DuckDB  
        conn.executemany("INSERT INTO movies_vectors VALUES (?, ?, ?)", [  
            (movies_for_embedding[idx][0], movies_for_embedding[idx][1], embedding)  
            for idx, embedding in enumerate(embeddings) if len(embedding) == 256  
        ])  

        # 创建分层可导航小世界(HNSW)索引  
        conn.execute("CREATE INDEX movies_vector_index ON movies_vectors USING HNSW (vector)")  

        def search_similar_movies(query: str):  
            """ 查找相似的电影 """  
            query_vector = embed_text([query])  

            vector_array = ', '.join(str(num) for num in query_vector[0])  

            query = conn.sql(f"""  
                SELECT title  
                FROM movies_vectors  
                ORDER BY array_distance(vector, array[{vector_array}]::FLOAT[256])  -- 计算向量距离
                LIMIT 3  
            """)  

            print(query.explain())  # 打印查询计划以展示HNSW_INDEX_SCAN节点  
            return query.fetchall()

我们不仅返回相似度搜索的结果,还用 query.explain() 显示查询计划,表明 HNSW 索引确实被使用了。

使用推荐功能

把所有东西放在一起,这就是一个完整的例子。

from typing import List, Dict  
import duckdb  
import httpx  
from google.cloud import aiplatform  
from google.oauth2 import service_account  
from google.oauth2.service_account import Credentials  
from vertexai.language_models import TextEmbeddingModel, TextEmbeddingInput  

tmdb_api_key: str = 'your-tmdb-api-key'  
credentials = service_account.Credentials.from_service_account_file('your-sa.json')  
aiplatform.init(project='your-project', location='us-central1', credentials=credentials)  

def _get_movies(page: int, vote_avg_min: float, vote_count_min: float) -> List[Dict]:  
    """从TMDB API获取电影数据"""  
    response = httpx.get('https://api.themoviedb.org/3/discover/movie', headers={  
        'Authorization': f'Bearer {tmdb_api_key}'  
    }, params={  
        'sort_by': 'popularity.desc',  
        'include_adult': 'false',  
        'include_video': 'false',  
        'language': 'en-US',  
        'with_original_language': 'en',  
        'vote_average.gte': vote_avg_min,  
        'vote_count.gte': vote_count_min,  
        'page': page  
    })  
    response.raise_for_status()  # 对于错误的响应,抛出异常(raise an error for bad responses)  
    return response.json()['results']  

def get_movies(pages: int, vote_avg_min: float, vote_count_min: float) -> List[Dict]:  
    """生成器,从多页获取电影数据"""  
    for page in range(1, pages + 1):  
        yield from _get_movies(page, vote_avg_min, vote_count_min)  

def embed_text(texts: List[str]) -> List[List[float]]:  
    """使用Gemini生成文本列表的嵌入"""  
    model = TextEmbeddingModel.from_pretrained('text-embedding-004')  
    inputs = [TextEmbeddingInput(text, 'RETRIEVAL_DOCUMENT') for text in texts]  
    embeddings = model.get_embeddings(inputs, output_dimensionality=256)  
    return [embedding.values for embedding in embeddings]  

# 从TMDB API获取电影数据并生成嵌入  
movie_data = list(get_movies(3, 6.0, 1000))  
movies_for_embedding = [(movie['id'], movie['title'], movie['overview']) for movie in movie_data]  
embeddings = embed_text([overview for _, _, overview in movies_for_embedding])  

# 设置DuckDB,启用Vector Similarity Search (VSS) 扩展和持久化  
# 参阅:https://duckdb.org/docs/extensions/vss.html  
with duckdb.connect(database='movies.duckdb') as conn:  
    conn.execute("""  
        INSTALL vss;  
        LOAD vss;  
        SET hnsw_enable_experimental_persistence = true;  # 这是一个实验性功能  
    """)  

    conn.execute("""  
        CREATE TABLE movies_vectors (  
            id INTEGER,  
            title VARCHAR,  
            vector FLOAT[256]  
        )  
    """)  

    # 将嵌入插入到DuckDB中  
    conn.executemany("INSERT INTO movies_vectors VALUES (?, ?, ?)", [  
        (movies_for_embedding[idx][0], movies_for_embedding[idx][1], embedding)  
        for idx, embedding in enumerate(embeddings) if len(embedding) == 256  
    ])  

    # 创建HNSW索引  
    conn.execute("CREATE INDEX movies_vector_index ON movies_vectors USING HNSW (vector)")  

    def search_similar_movies(query: str):  
        """根据给定的查询描述搜索相似的电影"""  
        query_vector = embed_text([query])  

        vector_array = ', '.join(str(num) for num in query_vector[0])  

        query = conn.sql(f"""  
            SELECT title  
            FROM movies_vectors  
            ORDER BY array_distance(vector, array[{vector_array}]::FLOAT[256])  
            LIMIT 3  # 限制返回结果为3部电影  
        """)  

        print(query.explain())  # 显示HNSW_INDEX_SCAN节点的查询计划  
        return query.fetchall()  

    # 示例搜索  
    query_description = 'Movie with an action hero who drives fast cars'  
    similar_movies = search_similar_movies(query_description)  

    # 显示搜索结果  
    print(f"Movies similar to your query: '{query_description}':")  
    for movie in similar_movies:  
        print(f"标题: {movie[0]}")

你可以用自然语言描述来搜索与之类似的电影,例如:

    query_description = '关于动作英雄开快车的电影'  
    similar_movies = search_similar_movies(query_description)

系统将:

  1. 将你的描述转换为一个256维的向量
  2. 利用HNSW索引高效地找到相似的电影
  3. 返回最相关的前3个匹配

电影推荐演示程序,来源:作者

最后的结论

记得卡利马库斯和他的作品目录(Pinakes)——有时最优雅的解决方案也是最简单的。通常,科技行业对每个问题的回答是“添加更多的基础设施”。DuckDB VSS 让我们想起一个不变的事实:你并不总是需要分布式系统来有效解决问题。就像古代图书管理员用一些基本原理创造了开创性的系统一样,我们也可以用简单高效的工具解决现代推荐挑战。

无论你是构建电影推荐引擎,还是处理其他语义搜索挑战,这里介绍的原则和技术能为你的项目提供坚实的基础。结合向量相似性搜索的强大功能和DuckDB的简洁性,可以创建复杂的搜索和推荐系统,或是增强检索生成(RAG),而无需处理分布式架构带来的复杂性。

可收听音频版本 🎧

更喜欢通过听来学习?我根据这篇文章使用Google的AI技术(比如NotebookLM)创建了一个基于这篇文章的音频播客。

📌 虽然叙述是由AI生成的,内容、见解和技术解释均源自原文。你可以把它想象成有一个AI助手在帮你让内容更易懂。

喜欢这篇文章吗?🫶
  • 👏 如果你觉得这篇文章有价值,不妨多给它点几个赞(你可以点到50次哦!)
  • 💭 在下面的评论区分享你的想法吧——我很想听听你的见解
  • ✨ 标出你最喜欢的观点,方便以后回顾

    🙏 你的参与对我来说非常重要,它能让更多读者看到这些内容。你的参与意味着整个世界那么重要。

点击查看更多内容
TA 点赞

若觉得本文不错,就分享一下吧!

评论

作者其他优质文章

正在加载中
  • 推荐
  • 评论
  • 收藏
  • 共同学习,写下你的评论
感谢您的支持,我会继续努力的~
扫码打赏,你说多少就多少
赞赏金额会直接到老师账户
支付方式
打开微信扫一扫,即可进行扫码打赏哦
今天注册有机会得

100积分直接送

付费专栏免费学

大额优惠券免费领

立即参与 放弃机会
意见反馈 帮助中心 APP下载
官方微信

举报

0/150
提交
取消