在 DUCKDB 中构建人工智能驱动的搜索功能

发布于 2024-5-17 11:07

浏览

0收藏

TLDR

本文将展示如何使用 DuckDB 的数组类型和余弦相似度函数实现高效的基于文本嵌入的向量搜索，并通过电影数据集实例演示其应用。

简介

现在，许多分析任务涉及文本数据，例如电子商务商店的产品评论。这些任务包括但不限于分类、聚类和相似性比较。它们主要使用文本数据的向量嵌入表示来实现矢量搜索功能。

DuckDB 提供了 Array[2] 和 List[3] 数据类型，可以用来存储和处理向量嵌入，以实现向量搜索。在这个系列博客的第一篇中，我们将探讨相似性比较，学习如何在 DuckDB 中使用向量嵌入。我们将涵盖向量嵌入、余弦相似度以及基于嵌入的向量搜索。

什么向量搜索？

在自然语言处理（NLP）的世界中，向量嵌入或称为向量搜索，指的是文本数据的数值表示。这些嵌入将单词、短语甚至整个文档转换为实数向量，捕捉词之间的关系和文本数据的语义含义。通过将文本表示为向量嵌入，我们能够应用数学运算，如相似性比较、聚类和分类。让我们通过一个例子来进一步理解这一点。

这里有四个词的向量嵌入，使用了一个简单的向量嵌入模型：

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

注意：上述向量嵌入是通过 mixedbread-ai/mxbai-embed-large-v1[4]

将它们在图表上可视化后，我们得到：

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

graph visualization

根据我们的语义知识（常识），我们知道“狗”和“吠叫”之间的关系类似于“猫”和“喵喵叫”。乍一看，我们在 x 轴的左侧看到“狗”和“吠叫”，而在右侧看到“猫”和“喵喵叫”。为了定量分析这些词的关系，我们需要使用诸如余弦相似度之类的指标。

什么是余弦相似度？

余弦相似度是一种计算向量嵌入语义相似性的指标。它也常用于信息的语义检索中。我们通过取两个归一化的向量之间的点积（ dot product[5]

• 这个指标的值为 1 表示两个向量完全相同

• 值为 0 意味着它们是独立的（正交）

• 值为-1 表示它们是对立的（相反）

下面列出的是词对的余弦值:

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

通过比较余弦值(cosine(“dog”, “meow”)) 和 cosine(“cat”, “bark”))，我们可以推断出"meow"几乎与"dog"相反，而"cat"和"bark"也是如此。在我们拥有的词汇中，可以看到"dog"与"bark"的关系最为密切，而"cat"与"meow"的关系最为紧密。有趣的是，尽管"meow"和"bark"是对立的，但"dog"和"cat"并不是。也许这个模型捕捉到了它们都是被驯化的动物、宠物，并且非常可爱这一共同点。😆

DuckDB 的数组类型和余弦相似度函数

自 0.10 版本以来，DuckDB 提供了 ARRAY 类型来存储固定大小的数组，非常适合存储向量嵌入。这意味着 ARRAY 类型列中的所有字段都具有相同的长度和相同的底层类型。要使用此数据类型初始化表，您需要指定数组中每个元素的数据类型，后跟方括号和数组大小；例如， FLOAT[2] 将初始化一个大小为 2 的数组，其中每个元素都是 FLOAT 类型。

让我们来看看如何将上述数据转化为表格形式：

CREATE TABLE word_embeddings 
(word VARCHAR, embedding FLOAT[2]);

INSERT INTO word_embeddings
VALUES ("dog", [ 0.23, 0.37]),
       ("cat", [-0.27, 0.29]),
       ("bark", [ 0.35, -0.02]),
       ("meow", [-0.32, -0.09]);

这里我们提供了一个表格，其中包含单词及其向量嵌入。DuckDB还提供了一个函数 array_cosine_similarity(array1, array2) ，用于计算两个向量之间的余弦相似度指标。

对于上述表格，让我们计算词对之间的余弦相似度指标。

SELECT x.word as word_1,
       y.word as word_2,
       array_cosine_similarity(x.embedding, y.embedding) AS similarity_metric
FROM word_embeddings AS x
CROSS JOIN word_embeddings AS y
WHERE word_1 > word_2
ORDER BY similarity_metric DESC;

这为我们提供了与上述部分相同的结果：

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

注意：在 DuckDB v0.10.0 之前，用于存储可变大小数组的数据类型 LIST 也可用于存储嵌入。在这种情况下，您需要使用 list_cosine_similarity 函数。

基于嵌入的检索是如何实现向量搜索的？

嵌入式检索背后的核心思想是，将查询输入和数据集中的项目都表示为高维空间中向量嵌入，这样在对查询与项目之间的余弦相似度进行排名时，语义相似性得到了体现。因此，根据给定查询与数据集中项目的余弦相似度进行排名，排名靠前的项目是最相关的。让我们使用 Kaggle 上的电影数据集[6]来说明这一点。

数据集中包含电影的标题和概述，我已经通过 sentence-transformers 包中的 mxbai-embed-large-v1 模型计算了标题和概述的向量嵌入。

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

现在，假设我想要搜索一部与以下描述非常相似的电影：“一部关于战士为他的部落而战的电影”。为了检索相关的电影，我们将计算描述的嵌入，并将其与数据集中条目的嵌入进行比较。

标题嵌入的相似性

下面的 SQL 查询实现了上述描述嵌入的相似检索，针对标题嵌入。查询计算余弦相似度，并按降序排列结果，选择前 5 项。我们看到这些项目的标题中都包含“战士”这个词。

SELECT title, overview
    FROM (
        SELECT *, 
        array_cosine_similarity(title_embeddings, [0.7058067321777344, -0.0012793205678462982, -0.08653011173009872...]) AS score
        FROM movies_embeddings
    ) sq
    WHERE score IS NOT NULL
    ORDER BY score DESC LIMIT 5;

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

概述嵌入的相似性

当将上述描述的嵌入与电影概述的嵌入进行相似性检索时，结果完全不同，因为它们匹配了概述。这是因为每部电影的概述属性包含更多和不同的词语，这些词语与电影相关，而标题中这些词语有时并不是电影本身的描述。概述嵌入将更好地与电影描述嵌入语义相似，而标题嵌入有时可能对电影的描述不够充分。

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

嵌入组合的相似性

由于我们的数据集中有2个嵌入，我们可以通过将两个分数相加来计算它们的组合，并根据总和进行排名。

SELECT title, overview
    FROM (
        SELECT *,
        array_cosine_similarity(title_embeddings, [0.7058067321777344, ...]) AS score_1,
        array_cosine_similarity(overview_embeddings, [0.7058067321777344, ...]) AS score_2
        FROM movies_embeddings
    ) sq
    WHERE score_1 IS NOT NULL AND score_2 IS NOT NULL
    ORDER BY score_1+score_2 DESC LIMIT 5;

这一轮我们得到了一个不同的结果，它与第一个结果有些相似。

在 DUCKDB 中构建人工智能驱动的搜索功能-AI.x社区

结论

文本数据包含的信息深度和丰富性使其非常有价值。通过向量嵌入，我们将语言转换为数学空间，从而能够进行多种操作，这些操作提供了提取和转换其中存储的信息的机会，从而解锁向量搜索功能。DuckDB 凭借其高效的处理能力和用户友好的 SQL 界面，使得与向量嵌入的交互过程变得更加轻松。无论是在进行相似性搜索、聚类还是执行其他任何基于向量的操作时，DuckDB 都能提供无缝连接，让您更接近文本数据进行分析实验。