RAG实战：手把手教你用LlamaIndex存取Embedding向量！

在之前我们介绍了如何使用LlamaIndex构建一个非常简单的RAG应用，初步了解了LlamaIndex构建RAG应用的大体流程。在运行前篇的程序时，我们会发现两个令人头痛的问题：

使用llama-index-llms-huggingface构建本地大模型时，会花费相当一部分时间。
在对文档进行切分，将切分后的片段转化为embedding向量，构建向量索引时，会花费大量的时间。

上面两个问题虽然不会影响程序的使用，但是严重影响了我们的调试。试想一下，如果每次修改几行代码就要等待几分钟启动程序，那确实有点折磨人。

在[搭建一个大模型API服务]中，我们介绍了如何使用SWIFT框架搭建一个大模型API服务，这很好地解决了第一个问题。我们可以将构建本地大模型替换为API服务，这样就不用每次启动程序时都重新加载一遍模型权重了。不过为了方便演示，本文仍使用本地构建的方式加载大模型。本文要解决的痛点是第二个问题。

对于第二个问题，很容易就能想到可以将构建好的embedding向量和向量索引存储在文件或数据库(如Milvus向量数据库)中，然后在需要时从文件或数据库中直接读取这些数据。

存储

下面的代码展示了如何使用LlamaIndex将embedding向量和向量索引存储到文件中。

importloggingimportsysimporttorchfromllama_index.coreimportPromptTemplate,Settings,SimpleDirectoryReader,VectorStoreIndexfromllama_index.core.node_parserimportSentenceSplitterfromllama_index.embeddings.huggingfaceimportHuggingFaceEmbeddingfromllama_index.llms.huggingfaceimportHuggingFaceLLM# 定义日志logging.basicConfig(stream=sys.stdout,level=logging.DEBUG)logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))# 定义system promptSYSTEM_PROMPT="""You are a helpful AI assistant."""query_wrapper_prompt=PromptTemplate("[INST]<<SYS>>\n"+SYSTEM_PROMPT+"<</SYS>>\n\n{query_str}[/INST] ")# 使用llama-index创建本地大模型llm=HuggingFaceLLM(context_window=4096,max_new_tokens=2048,generate_kwargs={"temperature":0.0,"do_sample":False},query_wrapper_prompt=query_wrapper_prompt,tokenizer_name='/yldm0226/models/Qwen1.5-14B-Chat',model_name='/yldm0226/models/Qwen1.5-14B-Chat',device_map="auto",model_kwargs={"torch_dtype":torch.float16},)Settings.llm=llm# 使用llama-index-embeddings-huggingface构建本地embedding模型Settings.embed_model=HuggingFaceEmbedding(model_name="/yldm0226/RAG/BAAI/bge-base-zh-v1.5")# 读取文档documents=SimpleDirectoryReader("document").load_data()# 对文档进行切分，将切分后的片段转化为embedding向量，构建向量索引index=VectorStoreIndex.from_documents(documents,transformations=[SentenceSplitter(chunk_size=256)])# 将embedding向量和向量索引存储到文件中index.storage_context.persist(persist_dir='doc_emb')# 构建查询引擎query_engine=index.as_query_engine(similarity_top_k=5)# 查询获得答案response=query_engine.query("不耐疲劳，口燥、咽干可能是哪些证候？")print(response)

关键代码为index.storage_context.persist(persist_dir='doc_emb')，其中persist_dir是存储路径。

运行上述代码，我们可以得到以下输出：

从提供的中医临床证候信息来看，口燥、咽干的症状可能与以下证候相关： 1. 津液不足证：由于津液生成不足或者体内燥热导致，表现为口眼喉鼻干燥，咽干是其中的一个症状。 2. 津亏热结证：津液亏虚加上热邪内结，也可能出现口燥和咽干。 3. 津液亏涸证：严重的津液亏损可能导致口唇干燥、咽部干燥，伴随其他严重脱水症状。 4. 燥干清窍证：气候干燥或体质原因引起的津液缺乏，口鼻咽喉干燥也是其特征。 5. 津伤化燥证：燥热内蕴或内热化燥损伤津液，也会出现口燥、频饮但不解渴的现象。 因此，这些证候都有可能与不耐疲劳和口燥、咽干的症状相符合，需要结合其他临床表现来确定具体的证候类型。建议在中医诊断中由专业医生根据全人情况判断。

我们找到刚才定义的persist_dir所在的路径，可以发现该路径下有以下几个文件：

default_vector_store.json：用于存储embedding向量。
docstore.json：用于存储文档切分出来的片段。
graph_store.json：用于存储知识图数据。
image__vector_store.json：用于存储图像数据。
index_store.json：用于存储向量索引。

在上述代码中，我们只用到了纯文本文档，所以生成出来的graph_store.json和image__vector_store.json中没有数据。

读取

在将embedding向量和向量索引存储到文件中后，我们就不需要重复地执行对文档进行切分，将切分后的片段转化为embedding向量，构建向量索引的操作了。以下代码演示了如何使用LlamaIndex读取结构化文件中的embedding向量和向量索引数据：

importloggingimportsysimporttorchfromllama_index.coreimportPromptTemplate,Settings,StorageContext,load_index_from_storagefromllama_index.embeddings.huggingfaceimportHuggingFaceEmbeddingfromllama_index.llms.huggingfaceimportHuggingFaceLLM# 定义日志logging.basicConfig(stream=sys.stdout,level=logging.DEBUG)logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))# 定义system promptSYSTEM_PROMPT="""You are a helpful AI assistant."""query_wrapper_prompt=PromptTemplate("[INST]<<SYS>>\n"+SYSTEM_PROMPT+"<</SYS>>\n\n{query_str}[/INST] ")# 使用llama-index创建本地大模型llm=HuggingFaceLLM(context_window=4096,max_new_tokens=2048,generate_kwargs={"temperature":0.0,"do_sample":False},query_wrapper_prompt=query_wrapper_prompt,tokenizer_name='/yldm0226/models/Qwen1.5-14B-Chat',model_name='/yldm0226/models/Qwen1.5-14B-Chat',device_map="auto",model_kwargs={"torch_dtype":torch.float16},)Settings.llm=llm# 使用llama-index-embeddings-huggingface构建本地embedding模型Settings.embed_model=HuggingFaceEmbedding(model_name="/yldm0226/RAG/BAAI/bge-base-zh-v1.5")# 从存储文件中读取embedding向量和向量索引storage_context=StorageContext.from_defaults(persist_dir="doc_emb")index=load_index_from_storage(storage_context)# 构建查询引擎query_engine=index.as_query_engine(similarity_top_k=5)# 查询获得答案response=query_engine.query("不耐疲劳，口燥、咽干可能是哪些证候？")print(response)

关键代码为storage_context = StorageContext.from_defaults(persist_dir="doc_emb")和index = load_index_from_storage(storage_context),StorageContext.from_defaults(persist_dir="doc_emb")表示从doc_emb目录中读取embedding向量和向量索引，load_index_from_storage(storage_context)表示根据存储的embedding向量和向量索引重新构建检索索引。

运行上述程序，可以得到以下输出：

从提供的中医临床证候信息来看，口燥、咽干的症状可能与以下证候相关： 1. 津液不足证：由于津液生成不足或者体内燥热导致，表现为口眼喉鼻干燥，咽干是其中的一个症状。 2. 津亏热结证：津液亏虚加上热邪内结，也可能出现口燥和咽干。 3. 津液亏涸证：严重的津液亏损可能导致口唇干燥、咽部干燥，伴随其他严重脱水症状。 4. 燥干清窍证：气候干燥或体质原因引起的津液缺乏，口鼻咽喉干燥也是其特征。 5. 津伤化燥证：燥热内蕴或内热化燥损伤津液，也会出现口燥、频饮但不解渴的现象。 因此，这些证候都有可能与不耐疲劳和口燥、咽干的症状相符合，需要结合其他临床表现来确定具体的证候类型。建议在中医诊断中由专业医生根据全人情况判断。

需要注意的是，为了输出的可复现性，我们将大模型的temperature设置为0，do_sample设置为False，所以两次得到的输出基本相同；如果将temperature设置为大于0的小数，do_sample设置为True，大模型每次的输出可能都是不一样的。另外，如果你在实验时获得的输出与文中的输出不一致，这也是正常的，这与多个因素有关。