语义搜索

语义搜索旨在通过理解搜索查询和要搜索的语料库的语义来提高搜索准确性。与只能根据词汇匹配查找文档的关键词搜索引擎不同，语义搜索在处理同义词、缩写和拼写错误方面也表现良好。

背景

语义搜索背后的思想是将语料库中的所有条目（无论是句子、段落还是文档）嵌入到向量空间中。在搜索时，查询被嵌入到相同的向量空间中，并找到语料库中最接近的嵌入。这些条目应与查询具有高度的语义相似性。

对称与非对称语义搜索

对于您的设置，一个关键的区别是对称与非对称语义搜索

• 对于对称语义搜索，您的查询和语料库中的条目长度大致相同，内容量也相同。一个例子是搜索相似的问题：例如，您的查询可能是“如何在网上学习 Python？”，您想要找到像“如何在网络上学习 Python？”这样的条目。对于对称任务，您可以潜在地翻转查询和语料库中的条目。

  • 相关训练示例：Quora 重复问题。

  • 合适的模型：预训练句子嵌入模型

• 对于非对称语义搜索，您通常有一个简短的查询（例如问题或一些关键词），并且您想要找到一个较长的段落来回答该查询。一个例子是像“什么是 Python”这样的查询，您想要找到段落“Python 是一种解释型的、高级的和通用的编程语言。Python 的设计理念……”。对于非对称任务，翻转查询和语料库中的条目通常没有意义。

  • 相关训练示例：MS MARCO

  • 合适的模型：预训练 MS MARCO 模型

您为任务类型选择正确的模型至关重要。

手动实现

对于小型语料库（最多约 100 万个条目），我们可以通过手动实现语义搜索，方法是计算语料库和查询的嵌入向量，然后使用 语义文本相似度 和 SentenceTransformer.similarity 计算 语义文本相似度。

有关简单示例，请参阅 semantic_search.py

输出

Query: A man is eating pasta.
Top 5 most similar sentences in corpus:
A man is eating food. (Score: 0.7035)
A man is eating a piece of bread. (Score: 0.5272)
A man is riding a horse. (Score: 0.1889)
A man is riding a white horse on an enclosed ground. (Score: 0.1047)
A cheetah is running behind its prey. (Score: 0.0980)

Query: Someone in a gorilla costume is playing a set of drums.
Top 5 most similar sentences in corpus:
A monkey is playing drums. (Score: 0.6433)
A woman is playing violin. (Score: 0.2564)
A man is riding a horse. (Score: 0.1389)
A man is riding a white horse on an enclosed ground. (Score: 0.1191)
A cheetah is running behind its prey. (Score: 0.1080)

Query: A cheetah chases prey on across a field.
Top 5 most similar sentences in corpus:
A cheetah is running behind its prey. (Score: 0.8253)
A man is eating food. (Score: 0.1399)
A monkey is playing drums. (Score: 0.1292)
A man is riding a white horse on an enclosed ground. (Score: 0.1097)
A man is riding a horse. (Score: 0.0650)

"""
This is a simple application for sentence embeddings: semantic search

We have a corpus with various sentences. Then, for a given query sentence,
we want to find the most similar sentence in this corpus.

This script outputs for various queries the top 5 most similar sentences in the corpus.
"""

import torch

from sentence_transformers import SentenceTransformer

embedder = SentenceTransformer("all-MiniLM-L6-v2")

# Corpus with example sentences
corpus = [
    "A man is eating food.",
    "A man is eating a piece of bread.",
    "The girl is carrying a baby.",
    "A man is riding a horse.",
    "A woman is playing violin.",
    "Two men pushed carts through the woods.",
    "A man is riding a white horse on an enclosed ground.",
    "A monkey is playing drums.",
    "A cheetah is running behind its prey.",
]
# Use "convert_to_tensor=True" to keep the tensors on GPU (if available)
corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)

# Query sentences:
queries = [
    "A man is eating pasta.",
    "Someone in a gorilla costume is playing a set of drums.",
    "A cheetah chases prey on across a field.",
]

# Find the closest 5 sentences of the corpus for each query sentence based on cosine similarity
top_k = min(5, len(corpus))
for query in queries:
    query_embedding = embedder.encode(query, convert_to_tensor=True)

    # We use cosine-similarity and torch.topk to find the highest 5 scores
    similarity_scores = embedder.similarity(query_embedding, corpus_embeddings)[0]
    scores, indices = torch.topk(similarity_scores, k=top_k)

    print("\nQuery:", query)
    print("Top 5 most similar sentences in corpus:")

    for score, idx in zip(scores, indices):
        print(corpus[idx], f"(Score: {score:.4f})")

    """
    # Alternatively, we can also use util.semantic_search to perform cosine similarty + topk
    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=5)
    hits = hits[0]      #Get the hits for the first query
    for hit in hits:
        print(corpus[hit['corpus_id']], "(Score: {:.4f})".format(hit['score']))
    """

优化实现

您可以直接使用 util.semantic_search 函数，而不是自己实现语义搜索。

该函数接受以下参数

sentence_transformers.util.semantic_search(query_embeddings: ~torch.Tensor, corpus_embeddings: ~torch.Tensor, query_chunk_size: int = 100, corpus_chunk_size: int = 500000, top_k: int = 10, score_function: ~typing.Callable[[~torch.Tensor, ~torch.Tensor], ~torch.Tensor] = <function cos_sim>) → list[list[dict[str, int | float]]]

此函数在查询嵌入列表和语料库嵌入列表之间执行余弦相似度搜索。它可用于信息检索/语义搜索，适用于最多约 100 万个条目的语料库。

参数:

• query_embeddings (Tensor) – 包含查询嵌入的二维张量。

• corpus_embeddings (Tensor) – 包含语料库嵌入的二维张量。

• query_chunk_size (int, 可选) – 同时处理 100 个查询。增加此值可以提高速度，但需要更多内存。默认为 100。

• corpus_chunk_size (int, 可选) – 每次扫描语料库 10 万个条目。增加此值可以提高速度，但需要更多内存。默认为 500000。

• top_k (int, 可选) – 检索前 k 个匹配条目。默认为 10。

• score_function (Callable[[Tensor, Tensor], Tensor], 可选) – 用于计算分数的函数。默认情况下，为余弦相似度。

返回:

一个列表，其中每个查询对应一个条目。每个条目都是一个字典列表，其中包含键 ‘corpus_id’ 和 ‘score’，并按余弦相似度分数降序排序。

返回类型:

List[List[Dict[str, Union[int, float]]]]

默认情况下，最多并行处理 100 个查询。此外，语料库被分块为最多 50 万个条目的集合。您可以增加 query_chunk_size 和 corpus_chunk_size，这可以提高大型语料库的速度，但也会增加内存需求。

速度优化

为了获得 util.semantic_search 方法的最佳速度，建议将 query_embeddings 和 corpus_embeddings 放在同一个 GPU 设备上。这可以显著提高性能。此外，我们可以对语料库嵌入进行归一化，使每个语料库嵌入的长度为 1。在这种情况下，我们可以使用点积来计算分数。

corpus_embeddings = corpus_embeddings.to("cuda")
corpus_embeddings = util.normalize_embeddings(corpus_embeddings)

query_embeddings = query_embeddings.to("cuda")
query_embeddings = util.normalize_embeddings(query_embeddings)
hits = util.semantic_search(query_embeddings, corpus_embeddings, score_function=util.dot_score)

Elasticsearch

Elasticsearch 可以 索引密集向量 并将其用于文档评分。我们可以轻松地索引嵌入向量，将其他数据与向量一起存储，最重要的是，使用 近似最近邻搜索 (HNSW，另见下文) 在嵌入上高效检索相关条目。

有关更多详细信息，请参阅 semantic_search_quora_elasticsearch.py。

近似最近邻

如果使用精确最近邻搜索（如 util.semantic_search 所使用的方法），搜索包含数百万个嵌入的大型语料库可能非常耗时。

在这种情况下，近似最近邻 (ANN) 会很有帮助。在这里，数据被划分为较小的相似嵌入部分。可以有效地搜索此索引，并且可以在几毫秒内检索到具有最高相似度（最近邻）的嵌入，即使您有数百万个向量。但是，结果不一定是精确的。可能会遗漏一些具有高相似度的向量。

对于所有 ANN 方法，通常有一个或多个参数需要调整，这些参数决定了召回率-速度的权衡。如果您想要最高速度，则很有可能错过匹配项。如果您想要高召回率，则搜索速度会降低。

三个流行的近似最近邻库是 Annoy、FAISS 和 hnswlib。

示例

• semantic_search_quora_hnswlib.py

• semantic_search_quora_annoy.py

• semantic_search_quora_faiss.py

检索 & 重排序

对于复杂的语义搜索场景，建议使用两阶段检索 & 重排序 Pipeline：

有关更多详细信息，请参阅 检索 & 重排序。

示例

我们列出了一些常见的用例

相似问题检索

semantic_search_quora_pytorch.py [ Colab 版本 ] 展示了一个基于 Quora 重复问题 数据集的示例。用户可以输入一个问题，代码使用 util.semantic_search 从数据集中检索最相似的问题。作为模型，我们使用 distilbert-multilingual-nli-stsb-quora-ranking，该模型经过训练以识别相似问题并支持 50 多种语言。因此，用户可以使用 50 多种语言中的任何一种输入问题。这是一个对称搜索任务，因为搜索查询与语料库中的问题具有相同的长度和内容。

相似出版物检索

semantic_search_publications.py [ Colab 版本 ] 展示了如何查找相似的科学出版物的示例。作为语料库，我们使用在 EMNLP 2016 - 2018 会议上发表的所有出版物。作为搜索查询，我们输入较新出版物的标题和摘要，并从我们的语料库中查找相关的出版物。我们使用 SPECTER 模型。这是一个对称搜索任务，因为语料库中的论文由标题和摘要组成，我们搜索标题和摘要。

问题 & 答案检索

semantic_search_wikipedia_qa.py [ Colab 版本 ]：此示例使用在 Natural Questions 数据集 上训练的模型。它包含约 10 万个真实的 Google 搜索查询，以及来自维基百科的带注释的段落，其中提供了答案。这是一个非对称搜索任务的示例。作为语料库，我们使用较小的 Simple English Wikipedia，以便它可以轻松地放入内存中。

retrieve_rerank_simple_wikipedia.ipynb [ Colab 版本 ]：此脚本使用 检索 & 重排序 策略，并且是非对称搜索任务的一个示例。我们将所有维基百科文章分成段落，并使用 bi-encoder 对其进行编码。如果输入新的查询/问题，它将由相同的 bi-encoder 编码，并检索具有最高余弦相似度的段落。接下来，检索到的候选段落由 Cross-Encoder 重排序器评分，并将 Cross-Encoder 中得分最高的 5 个段落呈现给用户。我们使用的模型在 MS Marco Passage Reranking 数据集上进行了训练，该数据集包含约 50 万个来自 Bing 搜索的真实查询。