ColBERT——以詞元級別的向量嵌入提升信息檢索效果详情 - 向量,人工智能,檢索系統,llm,編碼 Momodel 日志

介紹

檢索增強一代（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

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")

我們首先從 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])

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
   )

在這裏，我們調用 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"])

在這裏，首先，我們調用 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"])

在上面的代碼中，我們提出了一個非常具體的問題，即 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,
)

我們首先通過 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})

我們首先從 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)

我們調用檢索器對象的 .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})

這裏的代碼與我們剛剛編寫的代碼非常相似

唯一的區別是，我們傳入 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)