在RAG應用中，有一個我們可以去提升的環節就是——Embedding模型，我在之前的文章《大模型主流應用RAG的介紹——從架構到技術細節》也說過可以去微調embedding模型以便增強我們整體的檢索能力。
最早我們用的是AI target=_blank class=infotextkey>OpenAI的Embedding模型text-embedding-ada-002，但這個模型后面不一定可以在正式環境中使用，而且我們也沒辦法去微調，因此讓我們在本文中探索對開源Embedding模型進行微調。
BAAI/bge-small-en
目前HuggingFace的MTEB(海量文本Embedding基準)排行榜上排名第一的Embedding模型是big-large-en，它由北京人工智能研究院(BAAI，智源)開發。它是一種預訓練的transformer模型，可用于各種自然語言處理任務，如文本分類、問答、文本生成等。該模型在海量文本和代碼數據集上進行訓練，并在海量文本Embedding基準(MTEB)上進行了微調。
在本文中，我們將使用 big-large-en的縮小版big-small-en，這是一個384維的小規模模型（OpenAI是1500+維），具有競爭力的性能，非常適合在google Colab中運行。大家也可以選擇中文版的bge-base-zh-v1.5，只有0.1G。當然你的硬件環境允許，也可以使用1.3G的bge-large-zh-v1.5等embedding模型。
微調Embedding模型與微調LLM
與LLM（大語言模型）微調相比，big-small-en微調的實現有一些不一樣，下面簡單說一下異同點：
相似點
-
兩種類型的微調都遵循相同的方法，即生成用于訓練和評估的數據集，微調模型，最后評估基本模型和微調模型之間的性能。 -
使用LLM自動生成訓練和評估數據集。
不同點
-
數據集內容在LLM微調和Embedding模型微調之間有所不同。用于LLM微調的數據集包含LLM生成的問題。在微調過程中，包括問題、答案、系統prompt等在內的一系列數據將以JSON行( jsonl)文件的形式傳遞給要進行微調的模型。
不同的是，用于Embedding模型微調的數據集包含以下三組:
-
queries：node_id映射和LLM生成的問題的集合。 -
corpus：node_id映射和相應節點中的文本的集合。 -
relevant_docs：查詢的node_id和語料庫node_id之間的交叉引用映射的集合。給定一個查詢，它告訴Embedding模型要查找哪個文本節點/語料庫。
-
由于我們使用開源Embedding模型 bge-small-en，微調的前提就是要先把它下載到您的本地環境。以Google Colab為例，經過微調的模型將被下載到筆記本的根目錄中。 -
評估方法在微調Embedding模型和微調LLM之間有所不同，我們可以使用Ragas框架來衡量精準度和答案相關性。然而，當使用Embedding模型微調時，我們無法測量答案的正確性，因為我們只能為我們的問題檢索相關節點。相反，我們使用一個稱為“命中率”的簡單度量，這意味著對于每個 (query, relevant_doc)對，我們用查詢檢索top-k文檔，如果結果包含relevant_doc，則它被認為是“命中”的。該指標可用于專有Embeddings，如OpenAI的Embedding模型和開源Embedding模型。對于開源Embedding模型，我們還可以使用來自sentence_transformers的InformationRetrievalEvaluator進行評估，因為它提供了一套更全面的指標。
微調Embedding模型似乎涉及到很多問題。幸運的是，LlamaIndex（我個人感覺LlamaIndex目前的發展可能會在RAG方面打敗LangChain）在最近的0.8.21版本中引入以下關鍵類/函數，使得微調Embedding模型變得超級簡單:
-
SentenceTransformersF.NETuneEngine -
generate_qa_embedding_pairs -
EmbeddingQAFinetuneDataset
這些類和函數為我們抽象了底層的詳細集成邏輯，使開發人員能夠非常直觀地調用它。
微調方法
為了可視化微調BAAI/big-small-en所涉及的主要任務，讓我們看看下圖:
如圖中的數值所示，主要任務包括:
-
通過調用 EmbeddingQAFinetuneDataset函數generate_qa_embedding_pairs，自動生成評估和訓練數據集的數據。 -
通過傳入基本模型和訓練數據集來構造 SentenceTransformersFinetuneEngine，然后調用其finetune函數來訓練基本模型。 -
創建經過微調的模型。 -
調用向量存儲索引檢索器檢索相關節點并評估基本模型的命中率。 -
調用 InformationRetrievalEvaluator來評估基本模型。 -
調用向量存儲索引檢索器檢索相關節點并評估微調模型的命中率。 -
調用 InformationRetrievalEvaluator來評估經過微調的模型。
基于LlamaIndex的微調Embeddings指南（文末有鏈接），我們將在我們的用例中微調bge-small-en模型。
實現細節
Step 1: 生成數據集
讓我們使用LLM來自動生成訓練和評估的數據集。
-
Load corpus
在我們的用例中NVIDIA的SEC 10-K文件（代碼中和文末都有鏈接）是一個169頁的PDF文檔（你可以用你自己的中文PDF），所以我們需要在生成數據集時將文檔分成兩部分——一部分用于訓練數據集，另一部分用于evalals數據集。
使用單獨的數據集進行訓練和評估被認為是一種很好的ML實踐。可以調用load_corpus函數來收集訓練數據集(前90頁)或eval數據集(其余頁面)的節點。下面是load_corpus的代碼片段:
!curl https://d18rn0p25nwr6d.cloudfront.net/CIK-0001045810/4e9abe7b-fdc7-4cd2-8487-dc3a99f30e98.pdf --output nvidia-sec-10k-2022.pdf
def load_corpus(docs, for_training=False, verbose=False):
parser = SimpleNodeParser.from_defaults()
if for_training:
nodes = parser.get_nodes_from_documents(docs[:90], show_progress=verbose)
else:
nodes = parser.get_nodes_from_documents(docs[91:], show_progress=verbose)
if verbose:
print(f'Parsed {len(nodes)} nodes')
return nodes
SEC_FILE = ['nvidia-sec-10k-2022.pdf']
print(f"Loading files {SEC_FILE}")
reader = SimpleDirectoryReader(input_files=SEC_FILE)
docs = reader.load_data()
print(f'Loaded {len(docs)} docs')
train_nodes = load_corpus(docs, for_training=True, verbose=True)
val_nodes = load_corpus(docs, for_training=False, verbose=True)
請記住，在LlamaIndex中，節點和頁面并不完全匹配。對于一個169頁的文檔，結果顯示它為訓練數據集解析了97個節點，為evals數據集解析了91個節點。這兩個數據集的節點數量足夠接近。讓我們繼續。
-
生成合成查詢和數據集
現在，讓我們生成訓練和評估的數據集。請注意，我們這里沒有傳遞LLM (gpt-3.5-turbo-0613)，只有OpenAI API密鑰。這是因為LlamaIndex的默認LLM是gpt-3.5-turbo-0613;如果沒有定義LLM，只要提供OpenAI API密鑰，則默認為它。
generate_qa_embedding_pairs是一個生成數據集的方便函數。基于上面load_corpus函數返回的節點，它為每個節點生成問題(默認為每個節點兩個問題，可以自定義)，然后用所有三組數據構建數據集:queries，corpus和relevant_docs(queries與corpus之間的映射對應的node_id)。
from llama_index.finetuning import (
generate_qa_embedding_pairs,
EmbeddingQAFinetuneDataset,
)
from llama_index.llms import OpenAI
os.environ["OPENAI_API_KEY"] = "sk-############"
openai.api_key = os.environ["OPENAI_API_KEY"]
train_dataset = generate_qa_embedding_pairs(train_nodes)
val_dataset = generate_qa_embedding_pairs(val_nodes)
train_dataset.save_json("train_dataset.json")
val_dataset.save_json("val_dataset.json")
train_dataset = EmbeddingQAFinetuneDataset.from_json("train_dataset.json")
val_dataset = EmbeddingQAFinetuneDataset.from_json("val_dataset.json")
下面是樣本訓練數據集的樣子。注意queries和corpus在截圖中是折疊的，因為每個都有超過100個數據對:
Step 2: 微調Embedding模型
SentenceTransformersFinetuneEngine就是為這個任務設計的。在底層，它執行多個子任務:
-
通過構建 SentenceTransformer加載預訓練模型，傳入BAAI/big-small-en模型id。 -
定義數據加載器。它加載我們的訓練數據集，將其解析為 查詢，語料庫和relevant_docs。然后循環查詢，將relevant_docs中的node_id與corpus中的文本節點進行映射，構造InputExample，其列表依次傳遞到創建DataLoader中. -
定義loss（損失函數）。它使用 sentence_transformersmultiplenegativerankingloss來訓練檢索設置的Embeddings。 -
定義評估器。它設置了一個帶有eval數據集的評估器來監控Embedding模型在訓練期間的表現。 -
運行訓練。它插入上面定義的數據加載器、損失函數和評估器來運行訓練。
LlamaIndex將微調Embedding模型的所有詳細子任務封裝在一個SentenceTransformersFinetuneEngine中，我們所需要做的就是調用它的finetune函數。下面，您可以看到展示LlamaIndex的代碼片段：
from llama_index.finetuning import SentenceTransformersFinetuneEngine
finetune_engine = SentenceTransformersFinetuneEngine(
train_dataset,
model_id="BAAI/bge-small-en",
model_output_path="test_model",
val_dataset=val_dataset,
)
finetune_engine.finetune()
embed_model = finetune_engine.get_finetuned_model()
Step 3: 評估微調后的模型
如上所述，我們使用兩種不同的評估方法:
-
命中率:對每個
query/relevant_doc對進行簡單的top-k檢索。如果搜索結果包含relevant_doc，那么它就是一個“命中”。這可以用于專有的Embeddings，例如OpenAI的Embedding模型和開源Embedding模型。請參閱下面代碼片段中的evaluate函數。 -
InformationRetrievalEvaluator:一個更全面的用于評估開源Embeddings的度量套件。請參閱下面代碼片段中的evaluate_st函數。
from llama_index.embeddings import OpenAIEmbedding
from llama_index import ServiceContext, VectorStoreIndex
from llama_index.schema import TextNode
from tqdm.notebook import tqdm
import pandas as pd
# function for hit rate evals
def evaluate(
dataset,
embed_model,
top_k=5,
verbose=False,
):
corpus = dataset.corpus
queries = dataset.queries
relevant_docs = dataset.relevant_docs
service_context = ServiceContext.from_defaults(embed_model=embed_model)
nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()]
index = VectorStoreIndex(nodes, service_context=service_context, show_progress=True)
retriever = index.as_retriever(similarity_top_k=top_k)
eval_results = []
for query_id, query in tqdm(queries.items()):
retrieved_nodes = retriever.retrieve(query)
retrieved_ids = [node.node.node_id for node in retrieved_nodes]
expected_id = relevant_docs[query_id][0]
is_hit = expected_id in retrieved_ids # assume 1 relevant doc
eval_result = {
"is_hit": is_hit,
"retrieved": retrieved_ids,
"expected": expected_id,
"query": query_id,
}
eval_results.Append(eval_result)
return eval_results
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from sentence_transformers import SentenceTransformer
def evaluate_st(
dataset,
model_id,
name,
):
corpus = dataset.corpus
queries = dataset.queries
relevant_docs = dataset.relevant_docs
evaluator = InformationRetrievalEvaluator(queries, corpus, relevant_docs, name=name)
model = SentenceTransformer(model_id)
return evaluator(model, output_path="results/")
-
評測OpenAI
現在，讓我們評估一下OpenAI的Embedding模型text-embedding-ada-002。代碼如下:
ada = OpenAIEmbedding()
ada_val_results = evaluate(val_dataset, ada)
df_ada = pd.DataFrame(ada_val_results)
hit_rate_ada = df_ada['is_hit'].mean()
結果：
-
評測 BAAI/bge-small-en
bge = "local:BAAI/bge-small-en"
bge_val_results = evaluate(val_dataset, bge)
df_bge = pd.DataFrame(bge_val_results)
hit_rate_bge = df_bge['is_hit'].mean()
evaluate_st(val_dataset, "BAAI/bge-small-en", name='bge')
結果：
-
評估微調后的model
finetuned = "local:test_model"
val_results_finetuned = evaluate(val_dataset, finetuned)
df_finetuned = pd.DataFrame(val_results_finetuned)
hit_rate_finetuned = df_finetuned['is_hit'].mean()
evaluate_st(val_dataset, "test_model", name='finetuned')
查看結果:
-
Summary of results
把評測結果放在一起，讓我們仔細看看。
命中率：我們的微調模型比其基本模型bge-small-en的性能提高了1.29%。與OpenAI的Embedding模型相比，我們的微調模型的性能僅低了4.85%。
InformationRetrievalEvaluator結果:經過微調的模型比其基本模型的性能提高了5.81%。與基本模型相比，微調模型對這30多個指標列中的每一個都有更好的數字。
總結
在本文中，我們探討了微調RAG管道的Embedding模型所涉及的步驟。我們使用開源的sentence_transformers模型BAAI/big-small-en作為我們的基本Embedding模型，介紹了如何生成用于訓練和評估的數據集，如何對其進行微調，以及如何評估基本模型和微調模型之間的性能差異。
評估結果表明，微調Embedding模型的性能比基本模型提高了1-6%，與OpenAI的Embedding模型相比，微調模型的性能損失僅為4.85%。這種性能提升可能因數據集的質量和數量而異。
我們還簡要探討了LlamaIndex的最新版本，該版本對任何Embedding模型的線性適配器進行了微調，從而提高了性能并避免了在RAG管道中重新嵌入文檔。
