ColBERT——以词元级别的向量嵌入提升信息检索效果_colbert模型

介绍

检索增强一代 （RAG） 自成立以来就风靡全球。RAG 是大型语言模型 （LLM） 提供或生成准确和事实答案所必需的。我们通过RAG解决LLM的事实性，我们尝试为LLM提供一个与用户查询上下文相似的上下文，以便LLM将处理此上下文并生成事实正确的响应。我们通过以向量嵌入的形式表示我们的数据和用户查询并执行余弦相似性来做到这一点。但问题是，所有传统方法都以单个嵌入表示数据，这对于良好的检索系统来说可能并不理想。在本指南中，我们将研究 ColBERT，它比传统的双编码器模型更准确地执行检索。

学习目标

• 了解 RAG 中的检索工作原理。
• 了解检索中的单个嵌入限制。
• 使用 ColBERT 的令牌嵌入改进检索上下文。
• 了解 ColBERT的后期交互如何改善检索。
• 了解如何使用 ColBERT 进行准确检索。

本文是作为数据科学博客马拉松的一部分发表的。

什么是RAG？

LLM 虽然能够生成既有意义又语法正确的文本，但这些 LLM 存在一个称为幻觉的问题。LLM 中的幻觉是 LLM 自信地生成错误答案的概念，也就是说，它们以一种让我们相信这是真的的方式编造了错误的答案。自引入 LLM 以来，这一直是一个主要问题。这些幻觉会导致不正确和事实上错误的答案。因此，引入了检索增强生成。
在RAG中，我们获取文档/文档块的列表，并将这些文本文档编码为称为向量嵌入的数值表示，其中单个向量嵌入表示单个文档块，并将它们存储在称为向量存储的数据库中。将这些块编码到嵌入中所需的模型称为编码模型或双编码器。这些编码器在大量数据语料库上进行训练，因此使它们足够强大，可以在单个矢量嵌入表示中对文档块进行编码。

现在，当用户向 LLM 请求查询时，我们将此查询提供给同一个编码器以生成单个向量嵌入。然后，此嵌入用于计算与文档块的各种其他向量嵌入的相似性分数，以获得文档中最相关的块。最相关的块或最相关的块列表以及用户查询将提供给 LLM。然后，LLM 接收此额外的上下文信息，然后生成与从用户查询接收的上下文一致的答案。这确保了 LLM 生成的内容是真实的，并且在必要时可以追溯。

传统双编码器的问题

传统编码器模型（如 all-miniLM、OpenAI 嵌入模型和其他编码器模型）的问题在于，它们将整个文本压缩为单个矢量嵌入表示。这些单向量嵌入表示非常有用，因为它们有助于高效、快速地检索相似文档。但是，问题在于查询和文档之间的上下文。单个向量嵌入可能不足以存储文档块的上下文信息，从而造成信息瓶颈。
想象一下，500 个单词被压缩到一个大小为 782 的向量。用单个向量嵌入来表示这样的块可能还不够，因此在大多数情况下，检索结果不尽如人意。在复杂查询或文档的情况下，单向量表示也可能失败。一种这样的解决方案是将文档块或查询表示为嵌入向量列表，而不是单个嵌入向量，这就是 ColBERT 的用武之地。

什么是ColBERT？

ColBERT（Contextual Late Interactions BERT）是一种双编码器，它以多向量嵌入表示形式表示文本。它接受一个查询或一个文档/一个小文档的块，并在令牌级别创建向量嵌入。也就是说，每个令牌都有自己的向量嵌入，查询/文档被编码为令牌级向量嵌入列表。令牌级嵌入是从预训练的 BERT 模型生成的，因此得名 BERT。
然后将这些存储在向量数据库中。现在，当查询进入时，会为其创建一个令牌级嵌入列表，然后在用户查询和每个文档之间执行矩阵乘法，从而生成包含相似性分数的矩阵。总体相似性是通过取每个查询令牌的文档令牌的最大相似度之和来实现的。其公式如下图所示：

在上面的等式中，我们看到我们在查询令牌矩阵（包含 N 个令牌级向量嵌入）和文档令牌矩阵的转置（包含 M 个令牌级向量嵌入）之间做一个点积，然后我们取每个查询令牌的文档令牌的最大相似度。然后，我们取所有这些最大相似性的总和，这为我们提供了文档和查询之间的最终相似性分数。这产生有效和准确检索的原因是，在这里我们有一个令牌级别的交互，这为查询和文档之间的更多上下文理解提供了空间。

为什么叫ColBERT？

由于我们在自身之前计算嵌入向量列表，并且仅在模型推理期间执行此 MaxSim（最大相似度）操作，因此将其称为后期交互步骤，并且由于我们通过令牌级交互获得更多上下文信息，因此称为上下文后期交互。因此，名称为Contextual Late Interactions BERT或ColBERT。这些计算可以并行执行，因此可以有效地计算。最后，一个问题是空间，也就是说，它需要大量的空间来存储这个令牌级向量嵌入列表。这个问题在 ColBERTv2 中得到了解决，其中嵌入通过称为残余压缩的技术进行压缩，从而优化了使用的空间。

ColBERT 动手实践示例

在本节中，我们将亲身体验 ColBERT，甚至检查它在常规嵌入模型中的性能。

第 1 步：下载库

我们将从下载以下库开始：

!pip install ragatouille langchain langchain_openai chromadb einops sentence-transformers tiktoken
1

• RAGatouille：该库使我们能够以易于使用的方式使用最先进的 （SOTA） 检索方法，例如ColBERT。它提供了在数据集上创建索引、查询索引的选项，甚至允许我们在数据上训练 ColBERT 模型。
• LangChain：这个库将允许我们使用开源嵌入模型，以便我们可以测试其他嵌入模型与 ColBERT 相比的工作情况。
• langchain_openai：安装 OpenAI 的 LangChain 依赖项。我们甚至将与 OpenAI Embedding 模型合作，以 ColBERT检查其性能。
• ChromaDB的：该库将允许我们在环境中创建一个向量存储，以便我们可以保存在数据上创建的嵌入，并在查询和存储的嵌入之间执行语义搜索。
• EINOPS：高效的张量矩阵乘法需要此库。 句子转换器和 tiktoken 库是开源嵌入模型正常工作所必需的。

第 2 步：下载预训练模型

在下一步中，我们将下载预训练的 ColBERT 模型。为此，代码将是

from ragatouille import RAGPretrainedModel

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
1
2
3

• 我们首先从 RAGatouille 库导入 RAGPretrainedModel 类。
• 然后我们调用 .from_pretrained（）并给出模型名称，即“colbert-ir/colbertv2.0”。

运行上述代码将实例化 ColBERT RAG 模型。现在让我们下载一个维基百科页面并从中执行检索。为此，代码将是：

from ragatouille.utils import get_wikipedia_page

document = get_wikipedia_page("Elon_Musk")
print("Word Count:",len(document))
print(document[:1000])
1
2
3
4
5

RAGatouille 带有一个名为 get_wikipedia_page 的方便函数，它接收字符串并获取相应的维基百科页面。在这里，我们下载了埃隆·马斯克（Elon Musk）的维基百科内容，并将其存储在变量文档中。让我们打印文档中存在的单词数和文档的前几行。

在这里，我们可以看到图片中的输出。我们可以看到，埃隆·马斯克（Elon Musk）的维基百科页面上共有64,668个单词。

第 3 步：编入索引

现在，我们将在此文档上创建一个索引。

RAG.index(
   # List of Documents
   collection=[document],
   # List of IDs for the above Documents
   document_ids=['elon_musk'],
   # List of Dictionaries for the metadata for the above Documents
   document_metadatas=[{"entity": "person", "source": "wikipedia"}],
   # Name of the index
   index_name="Elon2",
   # Chunk Size of the Document Chunks
   max_document_length=256,
   # Wether to Split Document or Not
   split_documents=True
   )
1
2
3
4
5
6
7
8
9
10
11
12
13
14

在这里，我们调用 RAG 的 .index（） 来索引我们的文档。为此，我们传递以下内容：

• 收集：这是我们要索引的文档列表。在这里，我们只有一个文档，因此是一个文档的列表。
• document_ids：每个文档都需要一个唯一的文档 ID。在这里，我们将其命名为elon_musk，因为该文件是关于埃隆·马斯克（Elon Musk）的。
• document_metadatas：每个文档都有其元数据。这又是一个字典列表，其中每个字典都包含特定文档的键值对元数据。
• index_name：我们正在创建的索引的名称。我们将其命名为 Elon2。
• max_document_size： 这与块大小类似。我们指定每个文档块应该有多少。在这里，我们给它的值是 256。如果我们不指定任何值，则 256 将作为默认块大小。
• split_documents：它是一个布尔值，其中 True 表示我们要根据给定的块大小拆分文档，而 False 表示我们希望将整个文档存储为单个块。

运行上面的代码将以每个块 256 个大小的文档进行分块，然后通过 ColBERT 模型嵌入它们，该模型将为每个块生成一个令牌级向量嵌入列表，最后将它们存储在索引中。此步骤需要一些时间才能运行，如果有 GPU，则可以加速。最后，它创建一个存储索引的目录。这里的目录将是“.ragatouille/colbert/indexes/Elon2”

第 4 步：常规查询

现在，我们将开始搜索。为此，代码将是

results = RAG.search(query="What companies did Elon Musk find?", k=3, index_name='Elon2')
for i, doc, in enumerate(results):
   print(f"---------------------------------- doc-{i} ------------------------------------")
   print(doc["content"])
1
2
3
4

• 在这里，首先，我们调用 RAG 对象的 .search（） 方法
• 为此，我们给出了包含查询名称、k（要检索的文档数）和要搜索的索引名称的变量
• 在这里，我们提供查询“埃隆·马斯克（Elon Musk）找到了哪些公司？获得的结果将采用字典格式的列表，其中包含内容、分数、排名、document_id、passage_id和document_metadata等键
• 因此，我们使用下面的代码以整洁的方式打印检索到的文档
• 在这里，我们浏览词典列表并打印文档的内容

运行代码将产生以下结果：

在图片中，我们可以看到第一份和最后一份文件完全涵盖了埃隆·马斯克（Elon Musk）创立的不同公司。ColBERT 能够正确检索回答查询所需的相关块。

第 5 步：特定查询

现在让我们更进一步，问它一个具体的问题。

results = RAG.search(query="How much Tesla stocks did Elon sold in \
Decemeber 2022?", k=3, index_name='Elon2')


for i, doc, in enumerate(results):
   print(f"---------------
   ------------------- doc-{i} ------------------------------------")
   print(doc["content"])
1
2
3
4
5
6
7
8

在上面的代码中，我们提出了一个非常具体的问题，即 2022 年 12 月售出了多少价值特斯拉 Elon 的股票。我们可以在这里看到输出。doc-1 包含问题的答案。埃隆已经出售了价值36亿美元的特斯拉股票。同样，ColBERT 能够成功检索给定查询的相关块。

第 6 步：测试其他模型

现在让我们尝试使用其他嵌入模型（包括开源和封闭的）来尝试同样的问题：

from langchain_community.embeddings import HuggingFaceEmbeddings
from transformers import AutoModel

model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)

model_name = "jinaai/jina-embeddings-v2-base-en"
model_kwargs = {'device': 'cpu'}

embeddings = HuggingFaceEmbeddings(
   model_name=model_name,
   model_kwargs=model_kwargs,
)
1
2
3
4
5
6
7
8
9
10
11
12

• 我们首先通过 Transformer 库中的 AutoModel 类下载模型。
• 然后，我们将model_name和model_kwargs存储在各自的变量中。
• 现在，为了在LangChain中使用这个模型，我们从LangChain导入HuggingFaceEmbeddings，并为其指定模型名称和model_kwargs。

运行此代码将下载并加载 Jina 嵌入模型，以便我们可以使用它

第 7 步：创建嵌入

现在，我们需要开始拆分我们的文档，然后从中创建嵌入并将它们存储在色度矢量存储中。为此，我们使用以下代码：

from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=256, 
    chunk_overlap=0)
splits = text_splitter.split_text(document)
vectorstore = Chroma.from_texts(texts=splits,
                                embedding=embeddings,
                                collection_name="elon")
retriever = vectorstore.as_retriever(search_kwargs = {'k':3})
1
2
3
4
5
6
7
8
9
10
11

• 我们首先从 LangChain 库导入 Chroma 和 RecursiveCharacterTextSplitter
• 然后，我们通过调用 RecursiveCharacterTextSplitter 的.from_tiktoken_encoder并向其传递 chunk_size 和 chunk_overlap 来实例化text_splitter
• 在这里，我们将使用提供给 ColBERT 的相同chunk_size
• 然后我们调用这个text_splitter的 .split_text（） 方法，并为其提供包含有关 Elon Musk 的维基百科信息的文档。然后，它根据给定的块大小拆分文档，最后，文档块列表存储在变量拆分中
• 最后，我们调用 Chroma 类的 .from_texts（） 函数来创建一个向量存储。对于这个函数，我们给出了拆分、嵌入模型和collection_name
• 现在，我们通过调用向量存储对象的 .as_retriever（） 函数从中创建一个检索器。我们给 k 值 3

运行此代码将获取我们的文档，将其拆分为每个块大小为 256 的较小文档，然后使用 Jina 嵌入模型嵌入这些较小的块，并将这些嵌入向量存储存储在色度向量存储中。

第 8 步：创建猎犬

最后，我们从中创建一个检索器。现在我们将执行向量搜索并检查结果。

docs = retriever.get_relevant_documents("What companies did Elon Musk find?",)

for i, doc in enumerate(docs):
 print(f"---------------------------------- doc-{i} ------------------------------------")
 print(doc.page_content)
1
2
3
4
5

• 我们调用检索器对象的 .get_relevent_documents（） 函数，并给它相同的查询。
• 然后，我们整齐地打印出检索到的前 3 个文档。
• 在图片中，我们可以看到 Jina Embedder 尽管是一个流行的嵌入模型，但我们查询的检索效果很差。它没有成功获得正确的文档块。

我们可以清楚地发现 Jina 和 ColBERT 模型之间的区别，前者是将每个区块表示为单个向量嵌入的嵌入模型，后者是将每个区块表示为令牌级嵌入向量列表的 ColBERT 模型。在这种情况下，ColBERT 的表现明显优于此。

第 9 步：测试 OpenAI 的嵌入模型

现在让我们尝试使用像 OpenAI 嵌入模型这样的闭源嵌入模型。

import os

os.environ["OPENAI_API_KEY"] = "Your API Key"

from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()

text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
              model_name = "gpt-4",
              chunk_size = 256,
              chunk_overlap  = 0,
              )

splits = text_splitter.split_text(document)
vectorstore = Chroma.from_texts(texts=splits,
                                embedding=embeddings,
                                collection_name="elon_collection")

retriever = vectorstore.as_retriever(search_kwargs = {'k':3})
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

这里的代码与我们刚刚编写的代码非常相似

• 唯一的区别是，我们传入 OpenAI API 密钥来设置环境变量。
• 然后，我们通过从LangChain导入OpenAI嵌入模型来创建一个实例。
• 在创建集合名称时，我们给出一个不同的集合名称，以便 OpenAI 嵌入模型中的嵌入存储在不同的集合中。
  运行此代码将再次获取我们的文档，将它们分块为大小为 256 的较小文档，然后使用 OpenAI 嵌入模型将它们嵌入到单向量嵌入表示中，最后将这些嵌入存储在色度矢量存储中。现在让我们尝试检索另一个问题的相关文档。

docs = retriever.get_relevant_documents("How much Tesla stocks did Elon sold in Decemeber 2022?",)

for i, doc in enumerate(docs):
  print(f"---------------------------------- doc-{i} ------------------------------------")
  print(doc.page_content)
1
2
3
4
5

• 我们看到，在检索到的块中找不到我们期望的答案。
• 第一部分包含有关 2022 年特斯拉股票的信息，但没有谈论埃隆出售它们。
• 剩下的两个文档块也可以看到同样的情况，它们包含的信息是关于特斯拉及其股票的，但这不是我们期望的信息。
• 上面检索到的块不会为 LLM 提供上下文来回答我们提供的查询。

即使在这里，我们也可以看到单向量嵌入表示与多向量嵌入表示之间的明显区别。多重嵌入表示可以清楚地捕获复杂的查询，从而实现更准确的检索。

结论

总之，ColBERT 通过在标记级别将文本表示为多向量嵌入，展示了比传统双编码器模型在检索性能方面的显着进步。这种方法允许在查询和文档之间更细致地理解上下文，从而获得更准确的检索结果，并减轻 LLM 中常见的幻觉问题。

关键要点

• RAG 通过提供用于生成事实答案的上下文信息来解决 LLM 中的幻觉问题。
• 传统的双编码器由于将整个文本压缩为单个矢量嵌入而存在信息瓶颈，导致检索精度低于标准。
• ColBERT具有令牌级嵌入表示形式，有助于更好地理解查询和文档之间的上下文，从而提高检索性能。
• ColBERT中的后期交互步骤与令牌级交互相结合，通过考虑上下文的细微差别来提高检索准确性。
• ColBERTv2通过残余压缩优化存储空间，同时保持检索效率。
• 动手实验表明，与传统和开源嵌入模型（如 Jina 和 OpenAIEmbedding）相比，ColBERT 在检索性能方面具有优势。

来源：https://www.analyticsvidhya.com/blog/2024/04/colbert-improve-retrieval-performance-with-token-level-vector-embeddings/