博客 / 詳情

返回

解決語義搜索痛點,基於對比學習的領域特定文本嵌入模型微調實踐

文本嵌入模型能夠將文本轉換為具有語義意義的向量表示,廣泛應用於檢索、分類等多種任務場景。然而,通用嵌入模型在處理特定領域任務時往往存在性能瓶頸。微調技術為解決這一問題提供了有效途徑。本文將深入探討嵌入模型微調的核心原理,並以AI職位匹配為例,詳細闡述基於對比學習的微調實現過程。

檢索增強生成中的嵌入應用

檢索增強生成(Retrieval-Augmented Generation, RAG)是文本嵌入模型的重要應用場景之一。在RAG系統中,當接收到用户輸入(如客户詢問)時,系統會自動從知識庫中檢索相關上下文信息(如FAQ條目),並將其傳遞給大型語言模型進行後續處理。

基於嵌入的檢索過程遵循標準的三步驟流程:首先為知識庫中的所有文檔計算向量表示,然後使用相同的嵌入模型將輸入文本轉換為向量,最後通過計算輸入向量與知識庫向量間的相似度來識別最相關的文檔。

這種語義搜索方法為處理任意文本內容提供了簡潔而靈活的解決方案,但在實際應用中仍面臨關鍵挑戰。

語義相似性與相關性的偏差問題

語義搜索的核心侷限在於相似性與相關性之間的偏差。即使查詢與知識庫項目在語義上高度相似(表現為嵌入向量間的小夾角),這種相似性也不能保證檢索結果能夠有效回答用户查詢。

以客户服務場景為例,用户查詢"我如何更新我的付款方式?"可能會匹配到"要查看您的付款歷史記錄,請訪問您賬户的賬單部分"這樣的結果。儘管兩者在語義上相關,但檢索結果並未提供解決用户問題所需的實際信息。

基於對比學習的嵌入微調方法

嵌入微調通過在特定任務數據上進行額外訓練來調整預訓練模型的表示能力。這種方法特別適用於需要匹配不同長度文本(如簡短查詢與詳細文檔)或理解領域特定術語的場景。例如,在雲計算領域,"擴展"和"實例"等術語具有專門的技術含義,通用模型可能無法準確表示這些概念。

對比學習是實現嵌入微調的核心技術。該方法通過在相關文本對之間最小化嵌入距離,同時在不相關文本對之間最大化嵌入距離,訓練模型區分有用和無用的檢索結果。

微調實現的五個關鍵階段

嵌入微調的完整流程包括五個核心階段:數據準備、模型選擇、損失函數設計、模型訓練和性能評估。下面以AI職位匹配為例,詳細説明每個階段的具體實現。

階段一:構建正負樣本對數據集

數據準備是微調過程中最關鍵且最耗時的環節。本案例從Hugging Face數據集中提取了涵蓋數據科學家、數據工程師、AI工程師等關鍵職位的工作描述。

 from datasets import load_dataset  
 
 # 從HF hub加載數據
 ds = load_dataset("datastax/linkedin_job_listings")

為了生成更貼近實際搜索場景的查詢,使用OpenAI的批處理API通過GPT-4o-mini為每個職位描述生成對應的類人化搜索查詢。批處理API雖然需要24小時處理時間,但成本僅為即時處理的50%,整個數據生成過程僅花費0.12美元。

考慮到大多數文本嵌入模型的512標記限制,需要對職位描述進行預處理,移除與職位資格無關的冗餘信息。經過數據清洗和去重處理,最終獲得1012個有效的職位描述-查詢配對。

為了提升對比學習效果,進一步為每個正樣本對構建負樣本對。使用預訓練嵌入模型計算所有職位描述間的語義相似度,然後為每個正樣本對選擇最不相似的職位描述作為負樣本,確保負樣本的唯一性。

 from sentence_transformers import SentenceTransformer  
import numpy as np  

# 加載嵌入模型
model = SentenceTransformer("all-mpnet-base-v2")  

# 編碼所有工作描述
job_embeddings = model.encode(df['job_description_pos'].to_list())  

# 計算相似性
similarities = model.similarity(job_embeddings, job_embeddings)

# 將與正匹配最不相似的JD匹配為負匹配

# 獲取相似性的排序索引
similarities_argsorted = np.argsort(similarities.numpy(), axis=1)  

# 初始化列表來存儲負樣本配對
negative_pair_index_list = []  

for i in range(len(similarities)):  

    # 從當前行的最小相似性索引開始
    j = 0  
    index = int(similarities_argsorted[i][j])  

    # 確保索引是唯一的
    while index in negative_pair_index_list:  
        j += 1  # 移動到下一個最小索引
        index = int(similarities_argsorted[i][j])  # 獲取下一個最小索引

    negative_pair_index_list.append(index)  

# 將負樣本配對添加到df
df['job_description_neg'] =   
               df['job_description_pos'].iloc[negative_pair_index_list].values

按照80%訓練集、10%驗證集、10%測試集的比例劃分數據,並將處理後的數據集上傳至Hugging Face Hub,便於後續訪問和使用。

 # 打亂數據集
df = df.sample(frac=1, random_state=42).reset_index(drop=True)  

# 分割為訓練、驗證和測試集(例如,80%訓練,10%驗證,10%測試)
train_frac = 0.8  
valid_frac = 0.1  
test_frac = 0.1  

# 定義訓練和驗證大小
train_size = int(train_frac * len(df))  
valid_size = int(valid_frac * len(df))  

# 創建訓練、驗證和測試數據集
df_train = df[:train_size]  
df_valid = df[train_size:train_size + valid_size]  
df_test = df[train_size + valid_size:]

from datasets import DatasetDict, Dataset  

# 將pandas DataFrames轉換回Hugging Face Datasets
train_ds = Dataset.from_pandas(df_train)  
valid_ds = Dataset.from_pandas(df_valid)  
test_ds = Dataset.from_pandas(df_test)  

# 合併到DatasetDict中
dataset_dict = DatasetDict({  
    'train': train_ds,  
    'validation': valid_ds,  
    'test': test_ds  
})  

# 將數據推送到hub
 dataset_dict.push_to_hub("shawhin/ai-job-embedding-finetuning")

處理完成的數據集可通過簡單的API調用進行加載:

 from datasets import load_dataset  
 
 # 導入數據
 dataset = load_dataset("shawhin/ai-job-embedding-finetuning")

階段二:預訓練模型選擇與評估

在獲得訓練數據後,需要選擇合適的預訓練模型作為微調基礎。通過對比多個基礎模型和語義搜索專用模型的性能,選擇最優的候選模型。

評估過程使用三元組評估器(TripletEvaluator),該評估器接受查詢、正樣本職位描述、負樣本職位描述的三元組,並計算模型在驗證集上的準確率。

 from sentence_transformers import SentenceTransformer  
from sentence_transformers.evaluation import TripletEvaluator  

# 導入模型
model_name = "sentence-transformers/all-distilroberta-v1"  
model = SentenceTransformer(model_name)  

# 創建評估器
evaluator_valid = TripletEvaluator(  
    anchors=dataset["validation"]["query"],  
    positives=dataset["validation"]["job_description_pos"],  
    negatives=dataset["validation"]["job_description_neg"],  
    name="ai-job-validation",  
)  
evaluator_valid(model)  

 #>> {'ai-job-validation_cosine_accuracy': np.float64(0.8811881188118812)}

經過多模型對比分析,選擇"all-distilroberta-v1"作為微調基礎,該模型在驗證集上展現出最高的基準準確率。

階段三:損失函數配置

損失函數的選擇需要根據具體的數據格式和下游任務需求進行確定。Sentence Transformers文檔提供了詳細的損失函數選擇指南,針對不同數據格式推薦相應的損失函數。

本案例採用MultipleNegativesRankingLoss,該損失函數專門設計用於處理(錨點、正樣本、負樣本)三元組格式的數據。

 from sentence_transformers.losses import MultipleNegativesRankingLoss  
 
 loss = MultipleNegativesRankingLoss(model)

階段四:模型微調訓練

在完成數據準備、模型選擇和損失函數配置後,開始進行模型微調訓練。首先需要定義關鍵的訓練超參數。

對比學習的有效性在很大程度上取決於較大的批次大小和充分的訓練時間。為了保持實現的簡潔性,本案例採用了經過驗證的超參數配置。

 from sentence_transformers import SentenceTransformerTrainingArguments  

num_epochs = 1  
batch_size = 16  
lr = 2e-5  
finetuned_model_name = "distilroberta-ai-job-embeddings"  

train_args = SentenceTransformerTrainingArguments(  
    output_dir=f"models/{finetuned_model_name}",  
    num_train_epochs=num_epochs,  
    per_device_train_batch_size=batch_size,  
    per_device_eval_batch_size=batch_size,  
    learning_rate=lr,  
    warmup_ratio=0.1,  
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss受益於批次中沒有重複樣本
    eval_strategy="steps",  
    eval_steps=100,  
    logging_steps=100,  
 )

使用SentenceTransformerTrainer進行模型訓練,該訓練器提供了完整的微調流程管理功能。

 from sentence_transformers import SentenceTransformerTrainer  

trainer = SentenceTransformerTrainer(  
    model=model,  
    args=train_args,  
    train_dataset=dataset["train"],  
    eval_dataset=dataset["validation"],  
    loss=loss,  
    evaluator=evaluator_valid,  
)  
 trainer.train()

階段五:模型性能評估

微調完成後,需要對模型性能進行全面評估。使用與階段二相同的評估方法,分別在驗證集和測試集上測試模型表現。

評估結果顯示,微調後的模型在驗證集上達到99%的準確率,在測試集上實現100%的準確率,表明微調過程顯著提升了模型在特定任務上的性能。

為了便於後續部署和使用,可以將訓練好的模型上傳至Hugging Face Hub:

 # 將模型推送到HF hub
 model.push_to_hub(f"shawhin/{finetuned_model_name}")

微調後的模型可以直接用於實際推理任務:

 # 導入模型
model = SentenceTransformer("shawhin/distilroberta-ai-job-embeddings")  

# 新查詢
query = "data scientist 6 year experience, LLMs, credit risk, content marketing"  
query_embedding = model.encode(query)  

# 編碼JD
jd_embeddings = model.encode(dataset["test"]["job_description_pos"])  

# 計算相似性
 similarities = model.similarity(query_embedding, jd_embeddings)

總結

本文深入探討了基於對比學習的嵌入模型微調技術,並通過AI職位匹配的實際案例驗證了該方法的有效性。微調後的模型在測試集上實現了100%的準確率,充分證明了針對特定領域進行模型優化的必要性和可行性。

嵌入模型微調不僅解決了通用模型在專業領域表現不佳的痛點,更為構建高質量的語義搜索系統提供了切實可行的技術路徑。通過精心設計的對比學習框架,模型能夠更好地理解領域特定的語義關係,顯著提升檢索的精確性和相關性。

展望未來,嵌入技術將朝着更加智能化和多元化的方向發展。多模態嵌入模型正成為研究熱點,其能夠在統一向量空間中融合文本、圖像、音頻等多種數據類型。結合本文介紹的微調方法,多模態模型有望在跨模態檢索、內容理解等複雜場景中發揮更大價值,為下一代智能搜索和推薦系統奠定堅實基礎。

隨着計算資源的不斷優化和訓練技術的持續改進,嵌入模型微調將變得更加高效和普及,為各行各業的智能化轉型提供強有力的技術支撐。

本文代碼:

https://avoid.overfit.cn/post/4137e7d3a7bc499da0eaf889923dad0c

作者:Shaw Talebi

user avatar
0 位用戶收藏了這個故事!

發佈 評論

Some HTML is okay.