在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能

發(fā)布于 2024-5-17 11:07

瀏覽

0收藏

TLDR

本文將展示如何使用 DuckDB 的數(shù)組類(lèi)型和余弦相似度函數(shù)實(shí)現(xiàn)高效的基于文本嵌入的向量搜索，并通過(guò)電影數(shù)據(jù)集實(shí)例演示其應(yīng)用。

簡(jiǎn)介

現(xiàn)在，許多分析任務(wù)涉及文本數(shù)據(jù)，例如電子商務(wù)商店的產(chǎn)品評(píng)論。這些任務(wù)包括但不限于分類(lèi)、聚類(lèi)和相似性比較。它們主要使用文本數(shù)據(jù)的向量嵌入表示來(lái)實(shí)現(xiàn)矢量搜索功能。

DuckDB 提供了 Array[2] 和 List[3] 數(shù)據(jù)類(lèi)型，可以用來(lái)存儲(chǔ)和處理向量嵌入，以實(shí)現(xiàn)向量搜索。在這個(gè)系列博客的第一篇中，我們將探討相似性比較，學(xué)習(xí)如何在 DuckDB 中使用向量嵌入。我們將涵蓋向量嵌入、余弦相似度以及基于嵌入的向量搜索。

什么向量搜索？

在自然語(yǔ)言處理（NLP）的世界中，向量嵌入或稱為向量搜索，指的是文本數(shù)據(jù)的數(shù)值表示。這些嵌入將單詞、短語(yǔ)甚至整個(gè)文檔轉(zhuǎn)換為實(shí)數(shù)向量，捕捉詞之間的關(guān)系和文本數(shù)據(jù)的語(yǔ)義含義。通過(guò)將文本表示為向量嵌入，我們能夠應(yīng)用數(shù)學(xué)運(yùn)算，如相似性比較、聚類(lèi)和分類(lèi)。讓我們通過(guò)一個(gè)例子來(lái)進(jìn)一步理解這一點(diǎn)。

這里有四個(gè)詞的向量嵌入，使用了一個(gè)簡(jiǎn)單的向量嵌入模型：

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

注意：上述向量嵌入是通過(guò) mixedbread-ai/mxbai-embed-large-v1[4]

將它們?cè)趫D表上可視化后，我們得到：

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

graph visualization

根據(jù)我們的語(yǔ)義知識(shí)（常識(shí)），我們知道“狗”和“吠叫”之間的關(guān)系類(lèi)似于“貓”和“喵喵叫”。乍一看，我們?cè)?x 軸的左側(cè)看到“狗”和“吠叫”，而在右側(cè)看到“貓”和“喵喵叫”。為了定量分析這些詞的關(guān)系，我們需要使用諸如余弦相似度之類(lèi)的指標(biāo)。

什么是余弦相似度？

余弦相似度是一種計(jì)算向量嵌入語(yǔ)義相似性的指標(biāo)。它也常用于信息的語(yǔ)義檢索中。我們通過(guò)取兩個(gè)歸一化的向量之間的點(diǎn)積（ dot product[5]

? 這個(gè)指標(biāo)的值為 1 表示兩個(gè)向量完全相同

? 值為 0 意味著它們是獨(dú)立的（正交）

? 值為-1 表示它們是對(duì)立的（相反）

下面列出的是詞對(duì)的余弦值:

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

通過(guò)比較余弦值(cosine(“dog”, “meow”)) 和 cosine(“cat”, “bark”))，我們可以推斷出"meow"幾乎與"dog"相反，而"cat"和"bark"也是如此。在我們擁有的詞匯中，可以看到"dog"與"bark"的關(guān)系最為密切，而"cat"與"meow"的關(guān)系最為緊密。有趣的是，盡管"meow"和"bark"是對(duì)立的，但"dog"和"cat"并不是。也許這個(gè)模型捕捉到了它們都是被馴化的動(dòng)物、寵物，并且非常可愛(ài)這一共同點(diǎn)。??

DuckDB 的數(shù)組類(lèi)型和余弦相似度函數(shù)

自 0.10 版本以來(lái)，DuckDB 提供了 ARRAY 類(lèi)型來(lái)存儲(chǔ)固定大小的數(shù)組，非常適合存儲(chǔ)向量嵌入。這意味著 ARRAY 類(lèi)型列中的所有字段都具有相同的長(zhǎng)度和相同的底層類(lèi)型。要使用此數(shù)據(jù)類(lèi)型初始化表，您需要指定數(shù)組中每個(gè)元素的數(shù)據(jù)類(lèi)型，后跟方括號(hào)和數(shù)組大小；例如， FLOAT[2] 將初始化一個(gè)大小為 2 的數(shù)組，其中每個(gè)元素都是 FLOAT 類(lèi)型。

讓我們來(lái)看看如何將上述數(shù)據(jù)轉(zhuǎn)化為表格形式：

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]);

這里我們提供了一個(gè)表格，其中包含單詞及其向量嵌入。DuckDB還提供了一個(gè)函數(shù) array_cosine_similarity(array1, array2) ，用于計(jì)算兩個(gè)向量之間的余弦相似度指標(biāo)。

對(duì)于上述表格，讓我們計(jì)算詞對(duì)之間的余弦相似度指標(biāo)。

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;

這為我們提供了與上述部分相同的結(jié)果：

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

注意：在 DuckDB v0.10.0 之前，用于存儲(chǔ)可變大小數(shù)組的數(shù)據(jù)類(lèi)型 LIST 也可用于存儲(chǔ)嵌入。在這種情況下，您需要使用 list_cosine_similarity 函數(shù)。

基于嵌入的檢索是如何實(shí)現(xiàn)向量搜索的？

嵌入式檢索背后的核心思想是，將查詢輸入和數(shù)據(jù)集中的項(xiàng)目都表示為高維空間中向量嵌入，這樣在對(duì)查詢與項(xiàng)目之間的余弦相似度進(jìn)行排名時(shí)，語(yǔ)義相似性得到了體現(xiàn)。因此，根據(jù)給定查詢與數(shù)據(jù)集中項(xiàng)目的余弦相似度進(jìn)行排名，排名靠前的項(xiàng)目是最相關(guān)的。讓我們使用 Kaggle 上的電影數(shù)據(jù)集[6]來(lái)說(shuō)明這一點(diǎn)。

數(shù)據(jù)集中包含電影的標(biāo)題和概述，我已經(jīng)通過(guò) sentence-transformers 包中的 mxbai-embed-large-v1 模型計(jì)算了標(biāo)題和概述的向量嵌入。

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

現(xiàn)在，假設(shè)我想要搜索一部與以下描述非常相似的電影：“一部關(guān)于戰(zhàn)士為他的部落而戰(zhàn)的電影”。為了檢索相關(guān)的電影，我們將計(jì)算描述的嵌入，并將其與數(shù)據(jù)集中條目的嵌入進(jìn)行比較。

標(biāo)題嵌入的相似性

下面的 SQL 查詢實(shí)現(xiàn)了上述描述嵌入的相似檢索，針對(duì)標(biāo)題嵌入。查詢計(jì)算余弦相似度，并按降序排列結(jié)果，選擇前 5 項(xiàng)。我們看到這些項(xiàng)目的標(biāo)題中都包含“戰(zhàn)士”這個(gè)詞。

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 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

概述嵌入的相似性

當(dāng)將上述描述的嵌入與電影概述的嵌入進(jìn)行相似性檢索時(shí)，結(jié)果完全不同，因?yàn)樗鼈兤ヅ淞烁攀觥＿@是因?yàn)槊坎侩娪暗母攀鰧傩园嗪筒煌脑~語(yǔ)，這些詞語(yǔ)與電影相關(guān)，而標(biāo)題中這些詞語(yǔ)有時(shí)并不是電影本身的描述。概述嵌入將更好地與電影描述嵌入語(yǔ)義相似，而標(biāo)題嵌入有時(shí)可能對(duì)電影的描述不夠充分。

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

嵌入組合的相似性

由于我們的數(shù)據(jù)集中有2個(gè)嵌入，我們可以通過(guò)將兩個(gè)分?jǐn)?shù)相加來(lái)計(jì)算它們的組合，并根據(jù)總和進(jìn)行排名。

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;

這一輪我們得到了一個(gè)不同的結(jié)果，它與第一個(gè)結(jié)果有些相似。

在 DUCKDB 中構(gòu)建人工智能驅(qū)動(dòng)的搜索功能-AI.x社區(qū)

結(jié)論

文本數(shù)據(jù)包含的信息深度和豐富性使其非常有價(jià)值。通過(guò)向量嵌入，我們將語(yǔ)言轉(zhuǎn)換為數(shù)學(xué)空間，從而能夠進(jìn)行多種操作，這些操作提供了提取和轉(zhuǎn)換其中存儲(chǔ)的信息的機(jī)會(huì)，從而解鎖向量搜索功能。DuckDB 憑借其高效的處理能力和用戶友好的 SQL 界面，使得與向量嵌入的交互過(guò)程變得更加輕松。無(wú)論是在進(jìn)行相似性搜索、聚類(lèi)還是執(zhí)行其他任何基于向量的操作時(shí)，DuckDB 都能提供無(wú)縫連接，讓您更接近文本數(shù)據(jù)進(jìn)行分析實(shí)驗(yàn)。