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

客户端库
python
vectordatabase

最后更新时间 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()

现在你可以在哈希表中存储嵌入，该哈希表

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