如何在 Python 中执行向量搜索并查找文档的语义相似性？

客户端库
python
向量数据库

最后更新日期：2024 年 4 月 20 日

问题

如何在 Python 中执行向量搜索并查找文档的语义相似性？

答案

要在 Python 中执行向量相似性搜索，首先创建索引以执行相似文档的推荐。对于模型 `all-distilroberta-v1`，请确保 `DIM` 为 `768`（参见示例）。

FT.CREATE vss_index ON HASH PREFIX 1 "doc:" SCHEMA name TEXT content TEXT creation NUMERIC SORTABLE update NUMERIC SORTABLE content_embedding VECTOR FLAT 6 TYPE FLOAT32 DIM 768 DISTANCE_METRIC COSINE

文档建模

然后导入建模库，为了使用 `all-distilroberta-v1`，您必须包含库 `SentenceTransformer`。

from sentence_transformers import SentenceTransformer

现在我们需要生成文档的向量表示。使用合适的模型计算以下内容的向量嵌入：

content = "This is an arbitrary content"
model = SentenceTransformer('sentence-transformers/all-distilroberta-v1')
embedding = model.encode(content).astype(np.float32).tobytes()

现在您可以将嵌入存储在 Hash 中，该 Hash

doc = { "content_embedding" : embedding, 
        "name" : "Document's title",
        "state" : document.state}

conn.hset("doc:{}".format(pk), mapping=doc)

搜索相似文档

为了搜索与给定文档相似的文档，您将按照之前创建向量嵌入数据库时的方式对文档进行建模。

model = SentenceTransformer('sentence-transformers/all-distilroberta-v1')
new_embedding = model.encode(content).astype(np.float32).tobytes()

然后执行相似性搜索。

q = Query("*=>[KNN 3 @v $vec]").return_field("__v_score").dialect(2)
res = conn.ft("vss_index").search(q, query_params={"vec": new_embedding})

参考

https://hugging-face.cn/sentence-transformers