微调特定于域的搜索的文本嵌入：附Python代码详解

阅读时长：20分钟

发布时间：2025-02-02

近日热文：全网最全的神经网络数学原理（代码和公式）直观解释
欢迎关注知乎和公众号的专栏内容
LLM架构专栏
 知乎LLM专栏
 知乎【柏企】
公众号【柏企科技说】【柏企阅文】
嵌入模型将文本表示为具有语义意义的向量。尽管它们可以很容易地用于无数的用例（例如检索、分类），但通用嵌入模型在特定领域的任务上可能表现不佳。克服此限制的一种方法是微调。在本文中，我将讨论这项技术背后的关键思想，并分享一个微调嵌入以将查询与AI招聘信息匹配的具体示例。

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

嵌入通过一个3步过程实现检索过程：

向量表示形式（即嵌入向量）是针对知识库中的所有项目计算的。
将输入文本转换为矢量表示形式（使用与步骤1中相同的嵌入模型）。
计算输入文本向量与知识库中每个项目之间的相似性，并返回最相似的项目。

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

相似性搜索的问题

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

例如，请考虑以下问题：“如何更新我的付款方式？”相似性搜索的顶部结果可能是：“要查看您的付款历史记录，请访问您帐户的Billing部分。”虽然它们在语义上相似，但结果并未提供有用的信息来回答问题。

微调嵌入

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

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

在这种情况下，微调嵌入模型涉及根据问题对及其适当的答案对其进行训练。做到这一点的关键方法是对比学习，它通过最小化相关对嵌入之间的距离同时最大化不相关的对嵌入之间的距离，教会模型区分有用和无用的结果。

如何微调？

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

收集正（和负）对
选择预训练模型
选择损失函数
微调模型
评估模型

我不会抽象地讨论每个步骤，而是使用一个具体的例子来演示（和讨论）每个步骤。

示例：微调AI Job Post上的嵌入

在这里，我将演练嵌入模型的微调，以将求职者与职位描述相匹配。正如我们将看到的，大多数步骤都是由sentence transformers库简化的。

1. 收集正对

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

首先，我从这个Hugging Face数据集中提取了各种关键头衔（例如数据科学家、数据工程师、AI工程师）的职位描述。

from datasets import load_dataset
ds = load_dataset("datastax/linkedin_job_listings")

接下来，我使用OpenAI的Batch API通过GPT-4o-mini生成与每个JD对应的类人搜索查询。Batch API需要24小时才能运行，但比立即完成要便宜50%。整个工作花了我0.12美元。（参见示例笔记本）

然后，我删除了与工作资格无关的描述的各个部分。这是一个重要的步骤，因为大多数文本嵌入模型无法处理超过512个标记。

为了生成正对，我将“清理”的JD与来自GPT-4o-mini的合成（类人）查询相匹配。然后，删除所有重复的行，从而产生1012个JD。

虽然我们可以到此为止，但我更进一步，为每个示例挑选出负对。我使用预先训练的嵌入模型来计算数据集中所有清理的JD之间的相似性。然后，对于每个正对，我挑选出与负数示例最不相似的JD（同时确保没有两行具有相同的负数示例）。

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)
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['job_description_neg'] = df['job_description_pos'].iloc[negative_pair_index_list].values

最后，我将数据拆分为train-validation-test集，并将它们上传到HuggingFace中心，这样就可以通过这个单行函数调用来访问它。

df = df.sample(frac=1, random_state=42).reset_index(drop=True)
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
train_ds = Dataset.from_pandas(df_train)
valid_ds = Dataset.from_pandas(df_valid)
test_ds = Dataset.from_pandas(df_test)
dataset_dict = DatasetDict({
    'train': train_ds,
    'validation': valid_ds,
    'test': test_ds
})
dataset_dict.push_to_hub("shawhin/ai-job-embedding-finetuning")

我们可以使用一行代码导入生成的数据集。

from datasets import load_dataset
dataset = load_dataset("shawhin/ai-job-embedding-finetuning")

2. 选择预训练模型

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

为此，我创建了一个计算器，它采用我们的示例（查询、正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)

在比较了几个模型之后，我选择了“all-distilroberta-v1”，因为它在验证集上具有最高的准确性（在任何微调之前）。

3. 选择一个损失函数

接下来，我们需要选择一个损失函数。这将取决于您的数据和下游任务。sentence transformers doc中有一个很好的摘要表，其中列出了适当损失函数的各种数据格式。

在这里，我使用了MultipleNegativesRankingLoss，因为它与我们的（anchor, positive, negative）三元组格式匹配。

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 = "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,
    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，以便轻松导入以进行推理。

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_embeddings = model.encode(dataset["test"]["job_description_pos"])
similarities = model.similarity(query_embedding, jd_embeddings)

后续我们会持续带来更多相关技术的深度解析和实践案例，敬请关注公众号 柏企科技圈 和 柏企阅文

本文由mdnice多平台发布

微调特定于域的搜索的文本嵌入：附Python代码详解