语义搜索

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

背景

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

对称与非对称语义搜索

您的设置的**关键区别**是*对称*与*非对称语义搜索*

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

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

  • 适用模型：预训练 Sentence Embedding 模型

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

  • 相关训练示例：MS MARCO

  • 适用模型：预训练 MS MARCO 模型

**选择正确的模型**以适应您的任务类型至关重要。

提示

对于非对称语义搜索，建议您使用 SentenceTransformer.encode_query 对您的查询进行编码，并使用 SentenceTransformer.encode_document 对您的语料库进行编码。

更通用的 SentenceTransformer.encode 方法与 SentenceTransformer.encode_query 和 SentenceTransformer.encode_document 的区别在于两个方面：

1. 如果没有提供 prompt_name 或 prompt，它会使用模型 prompts 字典中定义的“query”或“document”提示。

2. 它将 task 设置为“document”。如果模型包含 Router 模块，它将使用“query”或“document”任务类型将输入路由到适当的子模块。

请注意，SentenceTransformer.encode 是最通用的方法，可用于包括信息检索在内的任何任务；如果模型未经预定义提示和/或任务类型训练，则所有三种方法将返回相同的嵌入。

手动实现

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

一个简单的示例，请参见 semantic_search.py

输出

Query: How do artificial neural networks work?
Top 5 most similar sentences in corpus:
(Score: 0.5926) Neural networks are computing systems vaguely inspired by the biological neural networks that constitute animal brains.
(Score: 0.5288) Deep learning is part of a broader family of machine learning methods based on artificial neural networks with representation learning.
(Score: 0.4647) Machine learning is a field of study that gives computers the ability to learn without being explicitly programmed.
(Score: 0.1381) Mars rovers are robotic vehicles designed to travel on the surface of Mars to collect data and perform experiments.
(Score: 0.0912) Carbon capture technologies aim to collect CO2 emissions before they enter the atmosphere and store them underground.

Query: What technology is used for modern space exploration?
Top 5 most similar sentences in corpus:
(Score: 0.3754) Mars rovers are robotic vehicles designed to travel on the surface of Mars to collect data and perform experiments.
(Score: 0.3669) SpaceX's Starship is designed to be a fully reusable transportation system capable of carrying humans to Mars and beyond.
(Score: 0.3452) The James Webb Space Telescope is the largest optical telescope in space, designed to conduct infrared astronomy.
(Score: 0.2625) Renewable energy sources include solar, wind, hydro, and geothermal power that naturally replenish over time.
(Score: 0.2275) Carbon capture technologies aim to collect CO2 emissions before they enter the atmosphere and store them underground.

Query: How can we address climate change challenges?
Top 5 most similar sentences in corpus:
(Score: 0.3760) Global warming is the long-term heating of Earth's climate system observed since the pre-industrial period due to human activities.
(Score: 0.3144) Carbon capture technologies aim to collect CO2 emissions before they enter the atmosphere and store them underground.
(Score: 0.2948) Renewable energy sources include solar, wind, hydro, and geothermal power that naturally replenish over time.
(Score: 0.0420) Machine learning is a field of study that gives computers the ability to learn without being explicitly programmed.
(Score: 0.0411) Deep learning is part of a broader family of machine learning methods based on artificial neural networks with representation learning.

"""
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 documents
corpus = [
    "Machine learning is a field of study that gives computers the ability to learn without being explicitly programmed.",
    "Deep learning is part of a broader family of machine learning methods based on artificial neural networks with representation learning.",
    "Neural networks are computing systems vaguely inspired by the biological neural networks that constitute animal brains.",
    "Mars rovers are robotic vehicles designed to travel on the surface of Mars to collect data and perform experiments.",
    "The James Webb Space Telescope is the largest optical telescope in space, designed to conduct infrared astronomy.",
    "SpaceX's Starship is designed to be a fully reusable transportation system capable of carrying humans to Mars and beyond.",
    "Global warming is the long-term heating of Earth's climate system observed since the pre-industrial period due to human activities.",
    "Renewable energy sources include solar, wind, hydro, and geothermal power that naturally replenish over time.",
    "Carbon capture technologies aim to collect CO2 emissions before they enter the atmosphere and store them underground.",
]
# Use "convert_to_tensor=True" to keep the tensors on GPU (if available)
corpus_embeddings = embedder.encode_document(corpus, convert_to_tensor=True)

# Query sentences:
queries = [
    "How do artificial neural networks work?",
    "What technology is used for modern space exploration?",
    "How can we address climate change challenges?",
]

# 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(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(f"(Score: {score:.4f})", corpus[idx])

    """
    # 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, optional) – 同时处理 100 个查询。增加此值可提高速度，但需要更多内存。默认为 100。

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

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

• score_function (Callable[[Tensor, Tensor], Tensor], optional) – 计算分数的函数。默认为余弦相似度。

返回:

一个列表，其中包含每个查询的一个条目。每个条目都是一个字典列表，包含“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。

OpenSearch

OpenSearch 是一个社区驱动的开源搜索引擎，支持向量搜索功能。它允许您索引密集向量并使用近似最近邻算法执行高效的相似度搜索。OpenSearch 可用于实现传统的关键词搜索 (BM25) 和语义搜索，从而可以比较和组合这两种方法。

有关示例实现，请参阅 semantic_search_nq_opensearch.py，其中展示了如何将 OpenSearch 与 Natural Questions 数据集结合使用，演示了语义搜索和 BM25 搜索功能。

近似最近邻

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

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

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

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

示例

• semantic_search_quora_hnswlib.py

• semantic_search_quora_annoy.py

• semantic_search_quora_faiss.py

检索与重排

对于复杂的语义搜索场景，建议采用两阶段的检索与重排管道：

更多详情，请参阅 检索与重排。

示例

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

相似问题检索

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 搜索查询，以及来自维基百科的带有答案的标注段落。这是一个**非对称搜索任务**的示例。作为语料库，我们使用较小的 简易英语维基百科，以便轻松地加载到内存中。

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