针对特定领域搜索的文本embedding微调实战

Embedding模型将文本表示为具有语义意义的向量。尽管它们用途广泛（例如检索、分类），但通用embedding模型在特定领域任务上可能表现不佳。解决这问题的一种方法是通过微调。

在本文中，云朵君将讨论这项技术背后的关键思想，并分享一个微调embedding的具体示例，该示例用于将查询与 AI 招聘信息进行匹配。

文本embedding模型的一个常见用途是检索增强生成 (RAG)。在这种情况下，给定一个基于 LLM 的系统输入（例如客户问题），系统会自动从知识库中检索相关上下文（例如常见问题解答），并将其传递给 LLM。

Embedding 通过以下3个步骤实现检索过程。

1. 为知识库中的所有项目计算向量表示（即embedding）
2. 输入文本被转换成向量表示（使用与步骤 1 相同的embedding模型）
3. 计算输入文本向量与知识库中每一项之间的相似度，并返回最相似的项目

使用embedding的三步搜索过程。图片由作者提供。

此过程（即语义搜索）提供了一种简单灵活的方法来搜索任意文本项。然而，存在一个问题。

相似性搜索问题

尽管语义搜索很流行，但它有一个核心问题。也就是说，仅仅因为查询和知识库项相似（即它们相关的embedding向量之间的角度很小），并不一定意味着该项目有助于回答查询。

例如，考虑这个问题

“如何更新我的付款方式？ ”

相似性搜索的最佳结果可能是：

“要查看您的付款历史记录，请访问您帐户的“账单”部分。 ”

虽然它们在语义上相似，但结果并没有提供有用的信息来回答问题。

微调embedding

解决这个问题的一种方法是通过微调。我们通过额外的训练来调整embedding模型的行为。

例如，我们可能希望将客户问题与常见问题解答中的相应答案进行匹配。这不仅需要将一小段文本（问题）与一长段文本（答案）进行匹配，还可能涉及理解特定领域的术语。例如，在云计算领域，“扩展”和“实例”等术语具有非常具体的含义，通用模型可能无法恰当地表达这些含义。

在这种情况下，对embedding模型进行微调需要对问题及其相应的答案进行训练。实现这一目标的关键方法是对比学习，它教会模型区分有用和无用的结果最小化相关对embedding之间的距离，同时最大化不相关对embedding之间的距离。

对比学习

对齐不同嵌入空间的标准方法是对比学习（CL）。CL 的一个关键直觉是用相似的方式表示同一信息的不同视图。

这包括学习最大化正对之间相似度并最小化负对之间相似度的表示。在图像-文本模型中，正对可能是带有合适标题的图像，而负对可能是带有不相关标题的图像（如下所示）。

CL 的两个关键方面有助于其有效性

1. 由于可以从数据的固有结构（例如，来自网络图像的元数据）中整理出正对和负对，因此 CL 训练数据不需要手动标记，从而可以进行更大规模的训练和更强大的表示。
2. 它通过一种特殊的损失函数同时最大化正对相似性并最小化负对相似性，如 CLIP所示。

如何进行微调？

我们可以将微调过程分为5 个关键步骤。

1. 收集正负样本对
2. 选择预训练模型
3. 选择损失函数
4. 微调模型
5. 评估模型

我们用具体的例子来演示和讨论每个步骤。

示例：对 AI 招聘信息进行微调

在这里，我将逐步讲解如何微调embedding模型，以便将求职者与职位描述进行匹配。

1）收集正负样本对

第一步是准备训练数据。这是整个过程中最重要（也是最耗时）的部分。

首先，Hugging Face 数据集中^[1]提取了各种关键职位（例如数据科学家、数据工程师、AI 工程师）的职位描述。

from datasets import load_dataset
# load data from HF hub
ds = load_dataset("datastax/linkedin_job_listings")
# convert to pandas df
df = ds['train'].to_pandas()
# keep only title and description
df = df[['title', 'description']]
df.shape

接下来，通过 GPT-4o-mini 生成与每个关键词对应的human-like搜索查询。

然后，我删除了与职位资格无关的描述部分。这一步很重要，因为大多数文本embedding模型无法处理超过 512 个标记。

为了生成正对，我将“清理后的”JDs与GPT-4o-mini中的合成（human-like）查询进行匹配。然后，删除所有重复的行，最终得到1012个JDs。

接下来为每个示例挑选出负样本对。我使用预先训练的embedding模型来计算数据集中所有已清理的JD之间的相似度。然后，对于每个正样本对，我挑选出与负样本最不相似的JD（同时确保没有两行包含相同的负样本）。

从sentence_transformers导入SentenceTransformer
导入numpy作为np 
# 加载embedding模型
model = SentenceTransformer("all-mpnet-base-v2") 
# 编码所有职位描述
job_embeddings = model.encode(df['job_description_pos'].to_list()) 
# 计算相似度
similarities = model.similarity(job_embeddings, job_embeddings)
# 将与正匹配最不相似的最小 JDs 匹配为负匹配
# 获取相似度的排序索引
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

最后，我将数据分成训练-验证-测试集，并将**它们上传到 HuggingFace **，这样就可以通过这一行函数调用来访问它们。

# 打乱数据集
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 数据集
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("embedding-finetuning")

我们可以用一行代码导入结果数据集。

from datasets import load_dataset
# 导入数据
dataset = load_dataset("embedding-finetuning")

2）选择一个预先训练的模型

有了（终于）拿到训练数据后，我们接下来选择一个预训练模型进行微调。我通过比较各种基础模型和语义搜索模型来做到这一点。

为此，我创建了一个评估器，它采用我们的示例（query、正 JD、负 JD）三元组并计算准确率。以下是验证集的准确率示例。

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 ”，因为它在验证集上具有最高的准确度（在任何微调之前）。

3）选择损失函数

接下来，我们需要选择一个损失函数。这取决于你的数据和后续任务。sentence_transformer文档中有一个很好的汇总表，列出了适合各种损失函数的数据格式^[2]。

在这里，我使用了MultipleNegativesRankingLoss，因为它符合我们的（锚点、正值、负值）三元组格式。

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

4）微调模型

数据、模型和损失函数准备就绪后，我们现在可以对模型进行微调。为此，我们首先定义各种训练参数。

关键点在于，对比学习受益于更大的批量大小和更长的训练时间。为了简单起见，我只使用了部分超参数。

from sentence_transformers import SentenceTransformerTrainingArguments
num_epochs = 1
batch_size = 16
lr = 2e-5
finetuned_model_name = "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()

5）评估模型

最后，我们可以用类似于步骤 2 中评估预训练模型的方式来评估微调模型。结果显示验证集的准确率为 99%，测试集的准确率为 100%。

作为可选步骤，我们可以将模型推送到 Hugging Face Hub，以便轻松导入进行推理。

# 将模型推送到 HF 集线器
model.push_to_hub(f"{finetuned_model_name}")
# 导入模型
model = SentenceTransformer("ai-job-embeddings") 
# 新查询
query = "数据科学家 6 年经验，法学硕士，信用风险，内容营销"
query_embedding = model.encode(query) 
# 编码 JD
jd_embeddings = model.encode(dataset["test"]["job_description_pos"]) 
# 计算相似度
similarities = model.similarity(query_embedding, jd_embeddings)

写在最后

微调是指通过额外的训练，使预训练模型适应特定任务。本文介绍了如何将这一理念应用于文本embedding模型，以提高其将（类似人类的）求职查询与职位描述进行匹配的准确性。

另一个微调应用是多模态embedding模型，它在单个向量空间中表示多种模态（例如文本和图像）。在我的下一篇文章中，我将讨论如何使用类似的过程来使多模态embedding适应特定的用例。