NLP（一百零一）Embedding模型微調(diào)實踐

本文將會介紹如何使用Sentence Transformers和AutoTrain對開源的Embedding模型bge-base-zh-v1.5進(jìn)行微調(diào)，并驗證Embedding模型微調(diào)后的效果。

在RAG框架或者語義相似度計算任務(wù)時，Embedding模型是我們常常會打交道的模型。Sentence Transformers 是一個 Python 庫，用于使用和訓(xùn)練各種應(yīng)用的Embedding模型，例如檢索增強(qiáng)生成 (RAG)、語義搜索、語義文本相似度、釋義挖掘 (paraphrase mining) 等等。其 3.0 版本的更新是該工程自創(chuàng)建以來最大的一次，引入了一種新的訓(xùn)練方法。

本文將會以智源研究院（BAAI）開源的Embedding模型bge-base-zh-v1.5作為基準(zhǔn)模型，展示如何使用Sentence Transformers進(jìn)行評估，并使用多種訓(xùn)練框架對其進(jìn)行微調(diào)，驗證微調(diào)后的模型效果會有所提升。

評估指標(biāo)Baseline

我們在文章NLP（八十二）RAG框架中的Retrieve算法評估中，使用LlamaIndex框架對RAG流程中的各種Retrieve算法，包括Embedding模型召回，進(jìn)行了評估，評估指標(biāo)采用Hit Rate和MRR。本文將繼續(xù)使用這篇文章中給出的數(shù)據(jù)集進(jìn)行評估。

Sentence Transformers給出了便捷的Embedding模型召回效果的評估方法，我們在這里采用InformationRetrievalEvaluator評估器，評估指標(biāo)使用accuracy@5, accuracy@10, map@100, mrr@10, ndcg@10等。

示例評估代碼如下：

# -*- coding: utf-8 -*-
# @file: bge_base_zh_eval.py
import os
import json
import time
import torch
from pprint import pprint
from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from sentence_transformers.util import cos_sim

project_dir = os.path.dirname(os.path.abspath(__file__)).split('/src')[0]

# data process
# load dataset, get corpus, queries, relevant_docs
with open(os.path.join(project_dir, "data/doc_qa.json"), "r", encoding="utf-8") as f:
    content = json.loads(f.read())

corpus = content['corpus']
queries = content['queries']
relevant_docs = content['relevant_docs']

# # Load a model
# 替換成自己的模型完整路徑或使用huggingface modl id
model_name = "bge-base-zh-v1.5"
model_path = os.path.join(project_dir, f"models/{model_name}")
model = SentenceTransformer(model_path, device="cuda" if torch.cuda.is_available() else "cpu")
print("Model loaded")

s_time = time.time()

# # Evaluate the model
evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name=f"{os.path.basename(model_path)}",
    score_functions={"cosine": cos_sim}
)

# Evaluate the model
result = evaluator(model)
pprint(result)
print(f"Time cost: {time.time() - s_time:.2f}s")

我們在評估器中傳入queries, corpus, relevant_docs字典，加載完模型后即可進(jìn)行評估。

評估結(jié)果在下文中給出，作為baseline（基準(zhǔn)）指標(biāo)。

微調(diào)數(shù)據(jù)合成

在介紹Embedding模型微調(diào)前，我們將使用LlamaIndex框架來構(gòu)建微調(diào)數(shù)據(jù)集。數(shù)據(jù)集的構(gòu)建方法已在文章NLP（八十六）RAG框架Retrieve階段的Embedding模型微調(diào)中給出，這里再次介紹。

在LlamaIndex框架中，可方便地使用generate_qa_embedding_pairs方法，利用Prompt工程對文本生成相關(guān)問題并進(jìn)行關(guān)聯(lián)。

Embedding模型的微調(diào)數(shù)據(jù)合成腳本如下：

# -*- coding: utf-8 -*-
# @file: make_ft_corpus.py
import os
from llama_index.legacy.finetuning import (
    generate_qa_embedding_pairs
)
from llama_index.llms.openai import OpenAI
from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from dotenv import load_dotenv

load_dotenv()

project_dir = os.path.dirname(os.path.abspath(__file__)).split('/src')[0]

TRAIN_FILES = [os.path.join(project_dir, "data/ft_train.txt")]
VAL_FILES = [os.path.join(project_dir, "data/ft_test.txt")]

TRAIN_CORPUS_FPATH = os.path.join(project_dir, "data/ft_train_corpus.json")
VAL_CORPUS_FPATH = os.path.join(project_dir, "data/ft_val_corpus.json")


def load_corpus(files, verbose=False):
    if verbose:
        print(f"Loading files {files}")

    reader = SimpleDirectoryReader(input_files=files)
    docs = reader.load_data()
    if verbose:
        print(f"Loaded {len(docs)} docs")

    parser = SentenceSplitter(chunk_size=250, chunk_overlap=0)
    nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)

    if verbose:
        print(f"Parsed {len(nodes)} nodes")

    return nodes


train_nodes = load_corpus(TRAIN_FILES, verbose=True)
val_nodes = load_corpus(VAL_FILES, verbose=True)

llm = OpenAI(model="gpt-3.5-turbo", api_key=os.getenv("OPENAI_API_KEY"))

qa_generate_prompt_tmpl = """\
Context information is below.

---------------------
{context_str}
---------------------

Given the context information and not prior knowledge.
generate only questions based on the below query.

You are a Professor. Your task is to setup \
{num_questions_per_chunk} questions for an upcoming \
quiz/examination in Chinese. The questions should be diverse in nature \
across the document in Chinese. The questions should not contain options, not start with Q1/ Q2. \
Restrict the questions to the context information provided.
"""

train_dataset = generate_qa_embedding_pairs(nodes=train_nodes, llm=llm, num_questions_per_chunk=1, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl)
val_dataset = generate_qa_embedding_pairs(nodes=val_nodes, llm=llm, num_questions_per_chunk=1, qa_generate_prompt_tmpl=qa_generate_prompt_tmpl)

train_dataset.save_json(TRAIN_CORPUS_FPATH)
val_dataset.save_json(VAL_CORPUS_FPATH)

輸出結(jié)果如下：

Output:

Loading files ['/Users/admin/PycharmProjects/embedding_model_exp/data/ft_train.txt']
Loaded 1 docs
Parsing nodes: 100%|██████████| 1/1 [00:00<00:00, 23.54it/s]
Parsing nodes:   0%|          | 0/1 [00:00<?, ?it/s]Parsed 137 nodes
Loading files ['/Users/admin/PycharmProjects/embedding_model_exp/data/ft_test.txt']
Loaded 1 docs
Parsing nodes: 100%|██████████| 1/1 [00:00<00:00, 45.84it/s]
  0%|          | 0/137 [00:00<?, ?it/s]Parsed 111 nodes
100%|██████████| 137/137 [03:34<00:00,  1.57s/it]
100%|██████████| 111/111 [01:55<00:00,  1.04s/it]

這樣，我們就能得到微調(diào)數(shù)據(jù)集了，保存為ft_train_corpus.json和ft_val_corpus.json。

Embedding模型微調(diào)

接下來，我們將會對bge-base-zh-v1.5模型進(jìn)行微調(diào)，微調(diào)的目的是讓模型更適配我們自己的數(shù)據(jù)集，從而取得更好的召回效果。

以下筆者將介紹三種模型訓(xùn)練的框架，幫助讀者更好地實現(xiàn)Embedding模型微調(diào)。

使用 `sentence-transformers v3`

這里，我們使用的sentence-transformers模塊的版本為V3.0.0。

利用該模塊，我們不難實現(xiàn)Embedding模型微調(diào)，微調(diào)代碼如下：

# -*- coding: utf-8 -*-
# @file: ft_sentence_transformers_trainer.py
import os
import json
import time
import torch
from datasets import Dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from sentence_transformers.util import cos_sim
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers import SentenceTransformerTrainingArguments
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers import SentenceTransformerTrainer

start_time = time.time()
project_dir = os.path.dirname(os.path.abspath(__file__)).split('/src')[0]

# load eval dataset
with open(os.path.join(project_dir, "data/ft_val_dataset.json"), "r", encoding="utf-8") as f:
    eval_content = json.loads(f.read())

corpus, queries, relevant_docs = eval_content['corpus'], eval_content['queries'], eval_content['relevant_docs']
# load train dataset
with open(os.path.join(project_dir, "data/ft_train_dataset.json"), "r", encoding="utf-8") as f:
    train_content = json.loads(f.read())

train_anchor, train_positive = [], []
for query_id, context_id in train_content['relevant_docs'].items():
    train_anchor.append(train_content['queries'][query_id])
    train_positive.append(train_content['corpus'][context_id[0]])

train_dataset = Dataset.from_dict({"positive": train_positive, "anchor": train_anchor})

print(train_dataset)
print(train_dataset[0:5])

# Load a model
model_name = 'bge-base-zh-v1.5'
# 替換成自己的模型完整路徑或使用huggingface modl id
model_path = os.path.join(project_dir, f"models/{model_name}")
model = SentenceTransformer(model_path, device="cuda:0" if torch.cuda.is_available() else "cpu")
print("Model loaded")

# # Evaluate the model
evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name=f"{model_name}",
    score_functions={"cosine": cos_sim}
)
train_loss = MultipleNegativesRankingLoss(model)

# define training arguments
args = SentenceTransformerTrainingArguments(
    output_dir=f"ft_{model_name}",  # output directory and hugging face model ID
    num_train_epochs=5,  # number of epochs
    per_device_train_batch_size=2,  # train batch size
    gradient_accumulation_steps=2,  # for a global batch size of 512
    per_device_eval_batch_size=4,  # evaluation batch size
    warmup_ratio=0.1,  # warmup ratio
    learning_rate=2e-5,  # learning rate, 2e-5 is a good value
    lr_scheduler_type="cosine",  # use constant learning rate scheduler
    optim="adamw_torch_fused",  # use fused adamw optimizer
    tf32=True,  # use tf32 precision
    bf16=True,  # use bf16 precision
    batch_sampler=BatchSamplers.NO_DUPLICATES,
    eval_strategy="epoch",  # evaluate after each epoch
    save_strategy="epoch",  # save after each epoch
    logging_steps=10,  # log every 10 steps
    save_total_limit=3,  # save only the last 3 models
    load_best_model_at_end=True,  # load the best model when training ends
    metric_for_best_model=f"eval_{model_name}_cosine_ndcg@10",  # Optimizing for the best ndcg@10 score
)

# train the model
trainer = SentenceTransformerTrainer(
    model=model,    # the model to train
    args=args,      # training arguments
    train_dataset=train_dataset.select_columns(
        ["positive", "anchor"]
    ),  # training dataset
    loss=train_loss,
    evaluator=evaluator
)

trainer.train()
trainer.save_model()
print(f"cost time: {time.time() - start_time:.2f}s")

筆者在1張NVIDIA A800-SXM4-80GB型號的GPU上進(jìn)行訓(xùn)練，耗時約63.10秒。同時，我們會將微調(diào)后的Embedding模型保存在GPU上。

使用 `AutoTrain`

AutoTrain是一種自動訓(xùn)練和部署最先進(jìn)的機(jī)器學(xué)習(xí)模型的方法，與 Hugging Face 生態(tài)系統(tǒng)無縫集成。它提供了一種自動方式來訓(xùn)練和部署最先進(jìn)的機(jī)器學(xué)習(xí)模型。該應(yīng)用程序支持廣泛的機(jī)器學(xué)習(xí)任務(wù)，包括文本分類、文本回歸、實體識別、摘要、問答、翻譯和表格任務(wù)。

你不需要寫任何代碼，就能實現(xiàn)各類機(jī)器學(xué)習(xí)任務(wù)的訓(xùn)練。

這里，我們以Embedding模型訓(xùn)練為例，來介紹如何在AutoTrain實現(xiàn)模型訓(xùn)練。

首先，我們將上述微調(diào)數(shù)據(jù)集進(jìn)行數(shù)據(jù)處理，使其格式符合datasets模塊的格式，并上傳至HuggingFace Hub，項目名稱為jclian91/embedding_exp_semiconductor。

模型訓(xùn)練的配置文件（config.yml）如下：

task: sentence-transformers:pair
base_model: /workspace/code/embedding_model_exp/models/bge-base-zh-v1.5
project_name: autotrain-pair
log: tensorboard
backend: local

data:
  path: jclian91/embedding_exp_semiconductor
  train_split: train
  valid_split: dev
  column_mapping:
    sentence1_column: anchor
    sentence2_column: positive

params:
  max_seq_length: 512
  epochs: 5
  batch_size: 4
  lr: 2e-5
  optimizer: adamw_torch_fused
  scheduler: cosine
  gradient_accumulation: 2
  mixed_precision: fp16

訓(xùn)練命令為CUDA_VISIBLE_DEVICES=0 autotrain --config config.yml。

全程我們沒有寫任何代碼，只是配置了模型訓(xùn)練的相關(guān)參數(shù)，就完成了Embedding模型的微調(diào)，輕松 + 愉快。

使用 `LlamaIndex`

使用LlamaIndex框架來進(jìn)行Embedding模型微調(diào)，筆者已經(jīng)在文章NLP（八十六）RAG框架Retrieve階段的Embedding模型微調(diào)中介紹過了，這次不再詳述，但使用起來也是非常方便的。

評估指標(biāo)對比

我們對基準(zhǔn)模型、不同訓(xùn)練方式的Embedding模型進(jìn)行指標(biāo)對比，結(jié)果如下表所示：

• ft_bge-base-zh-v1.5和autotrain-bge-base-zh-v15都是對基準(zhǔn)模型bge-base-zh-v1.5進(jìn)行微調(diào)得到的模型，前者使用sentence-transformers微調(diào)，后者使用AutoTrain微調(diào)。

• 評估腳本為 src/baseline_eval/bge_base_zh_eval.py，使用Mac CPU測試，CPU型號為 Apple M2 Pro 。

• 不同的訓(xùn)練框架并不會太影響模型微調(diào)后的表現(xiàn)（除去部分模型參數(shù)的影響），決定Embedding模型效果的是數(shù)據(jù)集和模型本身，還有模型參數(shù)等因素。

總結(jié)

本文重點(diǎn)介紹了如何使用Sentence Transformers和AutoTrain對開源的Embedding模型bge-base-zh-v1.5進(jìn)行微調(diào)，并驗證Embedding模型微調(diào)后的效果。

Sentence Transformers是一個寶庫，它介紹了關(guān)于Embedding模型方方面面的內(nèi)容，是了解、深入Embedding模型必不可少的工具。后續(xù)筆者將會介紹Embedding模型量化、俄羅斯套娃嵌入模型（Matryoshka Representation Learning, MRL）等相關(guān)方面的內(nèi)容。

本文給出的Python代碼將會是一個更全面的Embedding模型相關(guān)實驗的部分內(nèi)容，后續(xù)項目將會陸續(xù)補(bǔ)充。該項目已開源至Github，網(wǎng)址為: https://github.com/percent4/embedding_model_exp .

參考文獻(xiàn)

1. Training and Finetuning Embedding Models with Sentence Transformers v3: https://huggingface.co/blog/train-sentence-transformers

2. Fine-tune Embedding models for Retrieval Augmented Generation (RAG): https://www.philschmid.de/fine-tune-embedding-model-for-rag

3. 俄羅斯套娃 (Matryoshka) 嵌入模型概述: https://huggingface.co/blog/zh/matryoshka

4. Finetune Embeddings: https://docs.llamaindex.ai/en/stable/examples/finetuning/embeddings/finetune_embedding/

5. NLP（八十六）RAG框架Retrieve階段的Embedding模型微調(diào): https://mp.weixin.qq.com/s?__biz=MzU2NTYyMDk5MQ==&mid=2247486333&idx=1&sn=29d00d472647bc5d6e336bec22c88139&chksm=fcb9b2edcbce3bfb42ea149d96fb1296b10a79a60db7ad2da01b85ab223394191205426bc025&token=1376257911&lang=zh_CN#rd

6. How to Fine-Tune Custom Embedding Models Using AutoTrain: https://huggingface.co/blog/abhishek/finetune-custom-embeddings-autotrain

7. Upload a dataset to the Hub: https://huggingface.co/docs/datasets/v1.16.0/upload_dataset.html

8. NLP（八十二）RAG框架中的Retrieve算法評估: https://mp.weixin.qq.com/s?__biz=MzU2NTYyMDk5MQ==&mid=2247486199&idx=1&sn=f24175b05bdf5bc6dd42efed4d5acae8&chksm=fcb9b367cbce3a711fabd1a56bb5b9d803aba2f42964b4e1f9a4dc6e2174f0952ddb9e1d4c55&token=402005631&lang=zh_CN#rd