RAG實戰2-如何使用LlamaIndex儲存和讀取向量

一蓑烟雨度平生發表於2024-03-06
AIIndex

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.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，大模型每次的輸出可能都是不一樣的。另外，如果你在實驗時獲得的輸出與文中的輸出不一致，這也是正常的，這與多個因素有關。

相關文章