跳到主要内容

SPLADE

SPLADE 嵌入是一个为文档和查询提供高度稀疏表示的模型,继承了词袋(BOW)模型的理想特性,如精确词匹配和效率。

Milvus 通过 SpladeEmbeddingFunction 类与 SPLADE 模型集成。此类提供了对文档和查询进行编码的方法,并返回与 Milvus 索引兼容的稀疏向量嵌入。

要使用此功能,请安装必要的依赖项:

pip install --upgrade pymilvus
pip install "pymilvus[model]"

要实例化 SpladeEmbeddingFunction,请使用以下命令:

from pymilvus import model

splade_ef = model.sparse.SpladeEmbeddingFunction(
model_name="naver/splade-cocondenser-selfdistil",
device="cpu"
)

参数

  • model_name (string)

    用于编码的 SPLADE 模型名称。有效选项为 naver/splade-cocondenser-ensembledistil(默认)、naver/splade_v2_maxnaver/splade_v2_distilnaver/splade-cocondenser-selfdistil。有关更多信息,请参考 Play with models

  • device (string)

    要使用的设备,CPU 使用 cpu,第 n 个 GPU 设备使用 cuda:n

要为文档创建嵌入向量,请使用 encode_documents() 方法:

docs = [
"Artificial intelligence was founded as an academic discipline in 1956.",
"Alan Turing was the first person to conduct substantial research in AI.",
"Born in Maida Vale, London, Turing was raised in southern England.",
]

docs_embeddings = splade_ef.encode_documents(docs)

# 打印嵌入向量
print("Embeddings:", docs_embeddings)
# 由于输出嵌入向量是 2D csr_array 格式,我们将其转换为列表以便于操作。
print("Sparse dim:", splade_ef.dim, list(docs_embeddings)[0].shape)

预期输出类似于以下内容:

Embeddings:   (0, 2001) 0.6392706036567688
(0, 2034) 0.024093208834528923
(0, 2082) 0.3230178654193878
...
(2, 23602) 0.5671860575675964
(2, 26757) 0.5770265460014343
(2, 28639) 3.1990697383880615
Sparse dim: 30522 (1, 30522)

要为查询创建嵌入向量,请使用 encode_queries() 方法:

queries = ["When was artificial intelligence founded", 
"Where was Alan Turing born?"]

query_embeddings = splade_ef.encode_queries(queries)

# 打印嵌入向量
print("Embeddings:", query_embeddings)
# 由于输出嵌入向量是 2D csr_array 格式,我们将其转换为列表以便于操作。
print("Sparse dim:", splade_ef.dim, list(query_embeddings)[0].shape)

预期输出类似于以下内容:

Embeddings:   (0, 2001)        0.6353746056556702
(0, 2194) 0.015553371049463749
(0, 2301) 0.2756537199020386
...
(1, 18522) 0.1282549500465393
(1, 23602) 0.13133203983306885
(1, 28639) 2.8150033950805664
Sparse dim: 30522 (1, 30522)