PostgreSQL上的语义搜索

Question

我知道PostgreSQL有三元组相似性搜索，甚至有针对它优化的索引(CREATE INDEX trgm_idx ON table USING gist (column gist_trgm_ops);(，可以直接从Django(Web框架(使用：

Model.objects.filter(attribute__trigram_similar=query_string)

但是，如果我想对数据库对象执行语义相似性查询，而不是表面相似性，该怎么办？(这显然与经典的三元组相似性有很大不同(。

很好的例子是谷歌的通用句子编码器，我将把所有字符串转换为 512 维嵌入向量(使用库(，并通过计算规范化点积(余弦相似性(来执行查询，并生成具有最高相似度的对象(或者可能是 n 个具有相似性的对象>=0.50(。

最简单的做法是在数据库对象上进行迭代(在框架级别(，但这效率非常低(特别是如果数据库很大(，因此我宁愿找到一种可以在数据库级别执行查询的方法(如果可能的话，也许为语义搜索设置最佳索引？

---

对预矢量化对象的数据库执行此自定义相似性搜索的最佳方法是什么？

如果我手动获得预矢量化数据库中所有对象的点积怎么办？

谢谢！

Answer 1

这是对 PostgreSQL 的查询。

您可以将 ARRAY[1,2,3] 替换为要搜索的文本的嵌入。长度必须与 documents.embeddings 字段相同。 例如，如果使用 OpenAI ada 模型来计算嵌入，则必须对搜索查询使用相同的模型。

CREATE OR REPLACE FUNCTION cosine_similarity(a float[], b float[]) RETURNS float AS $$
DECLARE
dot_product float = 0;
norm_a float = 0;
norm_b float = 0;
BEGIN
FOR i IN 1 .. array_length(a, 1)
LOOP
dot_product = dot_product + a[i] * b[i];
norm_a = norm_a + a[i] * a[i];
norm_b = norm_b + b[i] * b[i];
END LOOP;
RETURN dot_product / (sqrt(norm_a) * sqrt(norm_b));
END;
$$ LANGUAGE plpgsql;

SELECT *, cosine_similarity(documents.embeddings, ARRAY[1,2,3]) as cosine_similarity
FROM documents
ORDER BY cosine_similarity DESC