使用LlamaIndex对嵌入模型进行适配器微调

在本文中，我们将介绍如何在LlamaIndex中对任何黑盒嵌入模型（如sentence_transformers、OpenAI等）的适配器进行微调。通过这种方法，可以将嵌入表示转换为新的潜在空间，以优化针对特定数据和查询的检索性能。这些改进的检索性能可以进一步提升RAG（检索增强生成）系统的整体表现。

1. 安装必需的库

首先，我们需要安装一些必需的库：

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

2. 下载数据

我们将使用Lyft和Uber的PDF文档作为训练和验证数据集。

!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'

3. 加载与解析数据

我们将使用SimpleDirectoryReader和SentenceSplitter来加载和解析PDF文档。

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

TRAIN_FILES = ["./data/10k/lyft_2021.pdf"]
VAL_FILES = ["./data/10k/uber_2021.pdf"]

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

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

4. 生成合成查询

我们将使用LLM（如gpt-3.5-turbo）来生成每个文本块的相关问题。

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

train_dataset = generate_qa_embedding_pairs(train_nodes)
val_dataset = generate_qa_embedding_pairs(val_nodes)

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

5. 嵌入适配器微调

我们使用EmbeddingAdapterFinetuneEngine在现有嵌入模型之上进行微调。

from llama_index.finetuning import EmbeddingAdapterFinetuneEngine
from llama_index.core.embeddings import resolve_embed_model
import torch

base_embed_model = resolve_embed_model("local:BAAI/bge-small-en")

finetune_engine = EmbeddingAdapterFinetuneEngine(
    train_dataset,
    base_embed_model,
    model_output_path="model_output_test",
    epochs=4,
    verbose=True,
)

finetune_engine.finetune()

embed_model = finetune_engine.get_finetuned_model()

6. 评估微调模型

使用不同的排名指标（如Hit-rate和Mean Reciprocal Rank）来评估微调模型。

from llama_index.embeddings.openai import OpenAIEmbedding
from eval_utils import evaluate, display_results

ada = OpenAIEmbedding()
ada_val_results = evaluate(val_dataset, ada)
display_results(["ada"], [ada_val_results])

bge = "local:BAAI/bge-small-en"
bge_val_results = evaluate(val_dataset, bge)
display_results(["bge"], [bge_val_results])

ft_val_results = evaluate(val_dataset, embed_model)
display_results(["ft"], [ft_val_results])

display_results(["ada", "bge", "ft"], [ada_val_results, bge_val_results, ft_val_results])

可能遇到的错误

1. 网络连接错误：如果在下载数据集时遇到网络连接问题，可以尝试检查网络连接或使用其他数据源。
2. 依赖库兼容性问题：确保已安装的库版本匹配。如果遇到版本不兼容的问题，可以尝试升级或降级相关库版本。
3. 显存不足：运行大模型或训练复杂模型时可能会遇到显存不足的问题，可以尝试减少批量大小或使用更小的模型。

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

参考资料:

• LlamaIndex文档
• OpenAI API