QA对话目前是大语言模型的一大应用场景，在QA对话中，由于大语言模型信息的滞后性以及不包含业务知识的特点，我们经常需要外挂知识库来协助大模型解决一些问题。在外挂知识库的过程中，embedding模型的召回效果直接影响到大模型的回答效果，因此，在许多场景下，我们都需要微调我们的embedding模型来提高我们的召回效果。下面，我们就基于llama-index对BAAI/bge-base-zh-v1.5模型进行微调，关于该模型的介绍，可以参考https://huggingface.co/BAAI/bge-base-zh-v1.5。
平台介绍
对embedding模型进行微调的过程中需要使用GPU加速训练，由于家境贫寒，我这里就使用了Google colab提供的免费T4GPU进行微调测试。如果大家没办法使用这个，可以使用国内一些公司的GPU云平台，租便宜的GPU就行，微调这个模型所耗费的GPU资源不多。以下所有训练代码皆是在Jupter-notebook上编写并执行的。
依赖安装
安装一些依赖库，有些依赖需要制定版本，否则存在不兼容的问题。
!pip install langchain==0.0.300 llmx==0.0.15a0 openai==0.28.1 llama_index==0.8.23.post1 pypdf sentence-transformers 训练样本准备
我们当前的使用场景是QA问答场景，因此训练数据的格式最好也是问答的格式。我这里由于没有现成的问答样本（人工整理比较耗时），因此我就摘取了《明朝那些事儿》这个小说里面的部分章节，然后让GPT-3.5针对文章内容进行提问，从而形成问答对。代码如下
import json import openai import os from llama_index import SimpleDirectoryReader from llama_index.node_parser import SimpleNodeParser from llama_index.schema import MetadataMode from llama_index import ( VectorStoreIndex, SimpleDirectoryReader, ServiceContext, Response ) def load_corpus(docs, for_training=False, verbose=False): parser = SimpleNodeParser.from_defaults() if for_training: nodes = parser.get_nodes_from_documents(docs[:5], show_progress=verbose) else: nodes = parser.get_nodes_from_documents(docs[6:], show_progress=verbose) if verbose: print(f'Parsed {len(nodes)} nodes') return nodes SEC_FILE = ['embedding_test.txt'] # embedding_test.txt是我训练样本的文件名，即我摘取了部分小说章节并直接保存为了txt文件。 print(f"Loading files {SEC_FILE}") reader = SimpleDirectoryReader(input_files=SEC_FILE) docs = reader.load_data() print(f'Loaded {len(docs)} docs') docs_nodes = load_corpus(docs, for_training=True, verbose=True) len(docs_nodes) train_nodes = docs_nodes[:75] # 人工选择3分之2作为训练集 print(f'Loaded {len(train_nodes)} train docs') val_nodes = docs_nodes[76:] # 剩下三分之一作为验证集 print(f'Loaded {len(val_nodes)} val docs') 构造训练集和测试集
使用GPT3.5基于小说内容生成对应的问题，最后生成train_dataset.json作为训练集，val_dataset.json作为验证集。
from llama_index.finetuning import ( generate_qa_embedding_pairs, EmbeddingQAFinetuneDataset, ) from llama_index.llms import OpenAI os.environ["OPENAI_API_KEY"] = "sk-************" openai.api_key = os.environ["OPENAI_API_KEY"] openai.api_base = "https://************" prompt="""下方是上下文信息。 --------------------- {context_str} --------------------- 根据提供的上下文信息和没有先验知识的原则，仅基于以下查询生成问题。 你是一名教师/教授。你的任务是为即将到来的测验/考试设置{num_questions_per_chunk}个问题。这些问题应在文档中多样化，且仅限于所提供的上下文信息。 """ train_dataset = generate_qa_embedding_pairs(train_nodes, qa_generate_prompt_tmpl=prompt) val_dataset = generate_qa_embedding_pairs(val_nodes, qa_generate_prompt_tmpl=prompt) train_dataset.save_json("train_dataset.json") val_dataset.save_json("val_dataset.json") 微调Embedding模型
这里的微调都是使用的默认参数，在实际微调过程中，可根据实际情况进行调整。
from llama_index.finetuning import SentenceTransformersFinetuneEngine train_dataset = EmbeddingQAFinetuneDataset.from_json("train_dataset.json") val_dataset = EmbeddingQAFinetuneDataset.from_json("val_dataset.json") finetune_engine = SentenceTransformersFinetuneEngine( train_dataset, model_id="BAAI/bge-base-zh-v1.5", model_output_path="test_model", val_dataset=val_dataset, ) finetune_engine.finetune() #由于模型较小，且训练样本较少，微调过程非常快 embed_model = finetune_engine.get_finetuned_model() embed_model 评估微调后的模型
在评估阶段，我们对比了微调前、后的BAAI/bge-base-zh-v1.5模型以及OPENAI的ada002的Embedding模型的召回效果，代码如下：
from llama_index.embeddings import OpenAIEmbedding from llama_index import ServiceContext, VectorStoreIndex from llama_index.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 service_context = ServiceContext.from_defaults(embed_model=embed_model) nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()] index = VectorStoreIndex(nodes, service_context=service_context, 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 # assume 1 relevant doc eval_result = { "is_hit": is_hit, "retrieved": retrieved_ids, "expected": expected_id, "query": query_id, } eval_results.append(eval_result) return eval_results 注意，在执行下面的代码前，需要先在当前项目的目录下创建results文件夹，否则会导致程序执行失败。
from sentence_transformers.evaluation import InformationRetrievalEvaluator from sentence_transformers import SentenceTransformer def evaluate_st( dataset, model_id, name, ): corpus = dataset.corpus queries = dataset.queries relevant_docs = dataset.relevant_docs evaluator = InformationRetrievalEvaluator(queries, corpus, relevant_docs, name=name) model = SentenceTransformer(model_id) return evaluator(model, output_path="results/") OPENAI-ada002
ada = OpenAIEmbedding() ada_val_results = evaluate(val_dataset, ada) df_ada = pd.DataFrame(ada_val_results) hit_rate_ada = df_ada['is_hit'].mean() hit_rate_ada ada002模型的最终评测结果为0.9285714285714286
原始BAAI/bge-base-zh-v1.5
bge = "local:BAAI/bge-base-zh-v1.5" bge_val_results = evaluate(val_dataset, bge) df_bge = pd.DataFrame(bge_val_results) hit_rate_bge = df_bge['is_hit'].mean() hit_rate_bge 原始的bge-base-zh-v1.5模型的评测结果为0.7663744588744589
微调后的BAAI/bge-base-zh-v1.5
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() hit_rate_finetuned 微调后模型的最终评测结果为0.975。即微调后，我们的embedding模型在当前数据集的召回效果由0.766上升到0.975。注意，得分并不是越高越好，需考虑是否过拟合，可以在其他数据集上再评测下。
以上，即是一次简单的微调过程。感谢技术的发展和开源大佬们的贡献，使得人工智能的应用门槛越来越低。
参考资料：
https://colab.research.google.com/github/wenqiglantz/nvidia-sec-finetuning/blob/main/embedding-finetuning/finetune_embedding_nvidia_sec.ipynb codedocllamajsonopenaiparsetransformerctogpuprompttransformersapi上下文store上下文信息llm大模型tpu大语言模型数据集语言模型googlegpt知识库schema查询生成pandasgpu云文件夹url参考资料gpt-3.5多样化人工智能使用场景langchainhuggingfacegpt3lmsapp人工智能的应用训练数据人脸识别
