特定领域嵌入模型微调：从理论到实践
在这篇文章中，我们将深入探讨针对特定领域（如医学、法律或金融）微调嵌入模型的过程。我们会专门为目标领域生成数据集，并利用它来训练模型，使其更好地理解所选领域内微妙的语言模式和概念。最终，你将拥有一个针对特定领域优化的更强大的嵌入模型，从而在 NLP 任务中实现更准确的检索和更出色的结果。
嵌入模型：理解其概念
嵌入是对文本或图像的强大数值表示，能够捕捉语义关系。可以把文本或音频想象成多维空间中的一个点，相似的单词或短语在这个空间中会比不相似的更靠近。
嵌入在许多 NLP 任务中都至关重要，例如：
- 语义相似度：判断两段图像或文本的相似程度。
- 文本分类：根据文本的含义将数据分组。
- 问答系统：找到回答问题最相关的文档。
- 检索增强生成（RAG）：结合用于检索的嵌入模型和用于文本生成的语言模型，以提高生成文本的质量和相关性。
套娃表示学习（Matryoshka Representation Learning）
套娃表示学习（MRL）是一种创建“可截断”嵌入向量的技术。想象一系列嵌套的娃娃，每个娃娃里面都包含一个更小的娃娃。MRL 以这样的方式嵌入文本：较早的维度（就像外层的娃娃）包含最重要的信息，后续的维度则增加细节。这使得在需要时你可以仅使用嵌入向量的一部分，从而降低存储和计算成本。
Bge-base-en 模型
由北京智源人工智能研究院（BAAI）开发的 BAAI/bge-base-en-v1.5 模型是一个强大的文本嵌入模型。它在各种 NLP 任务中表现出色，并且在 MTEB 和 C-MTEB 等基准测试中也取得了良好的成绩。对于计算资源有限的应用（比如本文案例）来说，bge-base-en 模型是一个不错的选择。
为何要微调嵌入模型？
针对特定领域微调嵌入模型对于优化 RAG 系统至关重要。这个过程确保模型对相似性的理解与目标领域的特定上下文和语言细微差别相契合。经过微调的嵌入模型能够更好地检索到与问题最相关的文档，最终使 RAG 系统给出更准确、更相关的回答。
数据集格式：微调的基础
我们可以使用多种数据集格式进行微调，以下是最常见的类型：
- 正样本对：一对相关的句子（例如，问题与答案）。
- 三元组：（锚点，正样本，负样本）三元组，其中锚点与正样本相似，与负样本不相似。
- 带有相似度分数的对：一对带有表示它们关系的相似度分数的句子。
- 带有类别的文本：带有相应类别标签的文本。
在本文中，我们将创建一个问题 - 答案对的数据集来微调我们的 bge-base-en-v1.5 模型。
损失函数：指导训练过程
损失函数对于训练嵌入模型至关重要。它们衡量模型预测与实际标签之间的差异，为模型调整权重提供信号。
不同的损失函数适用于不同的数据集格式：
- 三元组损失：用于（锚点，正样本，负样本）三元组，促使模型将相似的句子放置得更近，不相似的句子放置得更远。
- 对比损失：用于正负样本对，鼓励相似句子靠近，不相似句子远离。
- 余弦相似度损失：用于带有相似度分数的句子对，促使模型生成的嵌入具有与提供的分数相匹配的余弦相似度。
- 套娃损失：一种专门设计用于创建可截断的套娃嵌入的损失函数。
代码示例
安装依赖
首先要安装必要的库。我们将使用 datasets、sentence-transformers 和 google-generativeai 来处理数据集、嵌入模型和文本生成。
apt-get -qq install poppler-utils tesseract-ocr
pip install datasets sentence-transformers google-generativeai
pip install -q --user --upgrade pillow
pip install -q unstructured["all-docs"] pi_heif
pip install -q --upgrade unstructured
pip install --upgrade nltk
我们还会安装 unstructured 用于 PDF 解析，nltk 用于文本处理。
PDF 解析和文本提取
我们使用 unstructured 库从 PDF 文件中提取文本和表格。
import nltk
import os
from unstructured.partition.pdf import partition_pdf
from collections import Counter
nltk.download('punkt')
nltk.download('averaged_perceptron_tagger')
nltk.download('punkt_tab')
def process_pdfs_in_folder(folder_path):
total_text = []
pdf_files = [f for f in os.listdir(folder_path) if f.endswith('.pdf')]
for pdf_file in pdf_files:
pdf_path = os.path.join(folder_path, pdf_file)
print(f"Processing: {pdf_path}")
elements = partition_pdf(pdf_path, strategy="auto")
display(Counter(type(element) for element in elements))
text = "\n\n".join([str(el) for el in elements])
total_text.append(text)
return "\n\n".join(total_text)
folder_path = "data"
all_text = process_pdfs_in_folder(folder_path)
我们遍历指定文件夹中的每个 PDF，并将内容分割为文本、表格和图表，然后将文本元素组合成一个单一的文本表示。
自定义文本分块
接下来，我们使用 nltk 将提取的文本分割成易于管理的块，这对于使文本更适合语言模型处理是必不可少的。
import nltk
nltk.download('punkt')
def nltk_based_splitter(text: str, chunk_size: int, overlap: int) -> list:
from nltk.tokenize import sent_tokenize
sentences = sent_tokenize(text)
chunks = []
current_chunk = ""
for sentence in sentences:
if len(current_chunk) + len(sentence) <= chunk_size:
current_chunk += " " + sentence
else:
chunks.append(current_chunk.strip())
current_chunk = sentence
if current_chunk:
chunks.append(current_chunk.strip())
if overlap > 0:
overlapping_chunks = []
for i in range(len(chunks)):
if i > 0:
start_overlap = max(0, len(chunks[i - 1]) - overlap)
chunk_with_overlap = chunks[i - 1][start_overlap:] + " " + chunks[i]
overlapping_chunks.append(chunk_with_overlap[:chunk_size])
else:
overlapping_chunks.append(chunks[i][:chunk_size])
return overlapping_chunks
return chunks
chunks = nltk_based_splitter(text=all_text,
chunk_size=2048,
overlap=0)
数据集生成器
在这部分，我们定义两个函数：
prompt函数为 Google Gemini 创建一个提示，基于提供的文本块请求一个问题 - 答案对。
import google.generativeai as genai
import pandas as pd
GOOGLE_API_KEY = "xxxxxxxxxxxx"
def prompt(text_chunk):
return f"""
Based on the following text, generate one Question and its corresponding Answer.
Please format the output as follows:
Question: [Your question]
Answer: [Your answer]
Text: {text_chunk}
"""
generate_with_gemini函数与 Gemini 模型交互，并使用创建的提示生成一个 QA 对。
def generate_with_gemini(text_chunk:str, temperature:float, model_name:str):
genai.configure(api_key=GOOGLE_API_KEY)
generation_config = {"temperature": temperature}
gen_model = genai.GenerativeModel(model_name, generation_config=generation_config)
response = gen_model.generate_content(prompt(text_chunk))
try:
question, answer = response.text.split("Answer:", 1)
question = question.replace("Question:", "").strip()
answer = answer.strip()
except ValueError:
question, answer = "N/A", "N/A"
return question, answer
运行问答生成
使用 process_text_chunks 函数，我们使用 Gemini 模型为每个文本块生成 QA 对。
def process_text_chunks(text_chunks:list, temperature:int, model_name=str):
results = []
for chunk in text_chunks:
question, answer = generate_with_gemini(chunk, temperature, model_name)
results.append({"Text Chunk": chunk, "Question": question, "Answer": answer})
df = pd.DataFrame(results)
return df
df_results = process_text_chunks(text_chunks=chunks,
temperature=0.7,
model_name="gemini-1.5-flash")
df_results.to_csv("generated_qa_pairs.csv", index=False)
这些结果随后被存储在一个 Pandas DataFrame 中。
加载数据集
接下来，我们从 CSV 文件中将生成的 QA 对加载到 HuggingFace 数据集中，并确保数据格式适合微调。
from datasets import load_dataset
dataset = load_dataset('csv', data_files='generated_qa_pairs.csv')
def process_example(example, idx):
return {
"id": idx,
"anchor": example["Question"],
"positive": example["Answer"]
}
dataset = dataset.map(process_example,
with_indices=True,
remove_columns=["Text Chunk", "Question", "Answer"])
加载模型
我们从 HuggingFace 加载 BAAI/bge-base-en-v1.5 模型，并确保选择合适的执行设备（CPU 或 GPU）。
import torch
from sentence-transformers import SentenceTransformer
from sentence-transformers.evaluation import (
InformationRetrievalEvaluator,
SequentialEvaluator
)
from sentence-transformers.util import cos_sim
from datasets import load_dataset, concatenate_datasets
from sentence-transformers.losses import MatryoshkaLoss, MultipleNegativesRankingLoss
model_id = "BAAI/bge-base-en-v1.5"
model = SentenceTransformer(
model_id, device="cuda" if torch.cuda.is_available() else "cpu"
)
定义损失函数
在这里，我们配置套娃损失函数，指定用于截断嵌入的维度。
matryoshka_dimensions = [768, 512, 256, 128, 64]
inner_train_loss = MultipleNegativesRankingLoss(model)
train_loss = MatryoshkaLoss(
model, inner_train_loss, matryoshka_dims=matryoshka_dimensions
)
内部损失函数 MultipleNegativesRankingLoss 有助于模型生成适合检索任务的嵌入。
定义训练参数
我们使用 SentenceTransformerTrainingArguments 来定义训练参数，包括输出目录、训练轮数、批次大小、学习率和评估策略等。
from sentence-transformers import SentenceTransformerTrainingArguments
from sentence-transformers.training_args import BatchSamplers
args = SentenceTransformerTrainingArguments(
output_dir="bge-finetuned",
num_train_epochs=1,
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
per_device_eval_batch_size=16,
warmup_ratio=0.1,
learning_rate=2e-5,
lr_scheduler_type="cosine",
optim="adamw_torch_fused",
tf32=True,
bf16=True,
batch_sampler=BatchSamplers.NO_DUPLICATES,
eval_strategy="epoch",
save_strategy="epoch",
logging_steps=10,
save_total_limit=3,
load_best_model_at_end=True,
metric_for_best_model="eval_dim_128_cosine_ndcg@10"
)
注意：如果在 Tesla T4 上训练时遇到错误，可以尝试注释掉 tf32=True 和 bf16=True 这两行来禁用 TF32 和 BF16 精度。
创建评估器
我们创建一个评估器来衡量模型在训练过程中的性能。评估器使用 InformationRetrievalEvaluator 针对套娃损失中的每个维度评估模型的检索性能。
corpus = dict(
zip(dataset['train']['id'],
dataset['train']['positive'])
)
queries = dict(
zip(dataset['train']['id'],
dataset['train']['anchor'])
)
relevant_docs = {}
for q_id in queries:
relevant_docs[q_id] = [q_id]
matryoshka_evaluators = []
for dim in matryoshka_dimensions:
ir_evaluator = InformationRetrievalEvaluator(
queries=queries,
corpus=corpus,
relevant_docs=relevant_docs,
name=f"dim_{dim}",
truncate_dim=dim,
score_functions={"cosine": cos_sim}
)
matryoshka_evaluators.append(ir_evaluator)
evaluator = SequentialEvaluator(matryoshka_evaluators)
微调前评估模型
在微调之前，我们评估基础模型以获得基线性能。
results = evaluator(model)
for dim in matryoshka_dimensions:
key = f"dim_{dim}_cosine_ndcg@10"
print(f"{key}: {results[key]}")
定义训练器
我们创建一个 SentenceTransformerTrainer 对象，指定模型、训练参数、数据集、损失函数和评估器。
from sentence-transformers import SentenceTransformerTrainer
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=dataset.select_columns(
["positive", "anchor"]
),
loss=train_loss,
evaluator=evaluator
)
开始微调
trainer.train() 方法启动微调过程，使用提供的数据和损失函数更新模型的权重。
trainer.train()
trainer.save_model()
训练完成后，我们将表现最佳的模型保存到指定的输出目录。
微调后评估
最后，我们加载微调后的模型，并使用相同的评估器对其进行评估，以衡量微调后性能的提升。
from sentence-transformers import SentenceTransformer
fine_tuned_model = SentenceTransformer(
args.output_dir, device="cuda" if torch.cuda.is_available() else "cpu"
)
results = evaluator(fine_tuned_model)
for dim in matryoshka_dimensions:
key = f"dim_{dim}_cosine_ndcg@10"
print(f"{key}: {results[key]}")
通过针对特定领域微调嵌入模型，你的 NLP 应用程序能够更深入地理解该领域内的特定语言和概念，这可以在问答、文档检索和文本生成等任务中带来显著的改进。
祝你微调顺利！
希望这篇公众号文章能满足你的需求，如果还有其他问题或需要进一步修改，请随时告诉我。欢迎关注公众号柏企科技圈
推荐阅读
3. AI Agent 架构新变革：构建自己的 Plan-and-Execute Agent
评论