RAG实战2-如何使用LlamaIndex存储和读取向量

RAG实战2-如何使用LlamaIndex存储和读取embedding向量

本文是检索增强生成(Retrieval-augmented Generation,RAG)实战1-基于LlamaIndex构建第一个RAG应用的续集，在阅读本文之前请先阅读前篇。

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

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

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

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

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

存储

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

import logging
import sys
import torch
from llama_index.core import PromptTemplate, Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM

# 定义日志
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

# 定义system prompt
SYSTEM_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向量和向量索引数据：

import logging
import sys
import torch
from llama_index.core import PromptTemplate, Settings, StorageContext, load_index_from_storage
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM

# 定义日志
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

# 定义system prompt
SYSTEM_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，大模型每次的输出可能都是不一样的。另外，如果你在实验时获得的输出与文中的输出不一致，这也是正常的，这与多个因素有关。