利用LlamaIndex和本地PDF文檔，輕松打造知識圖譜GraphRAG

發布于 2024-8-7 14:56

瀏覽

0收藏

檢索增強生成（RAG）技術通過引入外部知識源，增強了大型語言模型的回答準確性和上下文契合度。盡管RAG在處理復雜異構信息時可能會忽略實體間的結構和聯系，例如，向量數據庫可能錯誤地將“員工”與“雇主”關聯得更緊密，而非“信息”。

知識圖譜的引入有效解決了這一局限。它采用節點和邊的三元組結構，如“雇主 — 提交 — 索賠”，清晰地表達了實體間的關系。這種結構化的方法讓知識圖譜在處理復雜數據搜索時更為精確和高效。

1 技術實現

1.1 安裝依賴項

!pip install -q pypdf
!pip install -q python-dotenv
!pip install -q pyvis
!pip install -q transformers einops accelerate langchain bitsandbytes sentence_transformers langchain-community langchain-core
!pip install -q llama-index
!pip install -q llama-index-llms-huggingface
!pip install -q llama-index-embeddings-langchain
!pip install -q llama-index-embeddings-huggingface

LlamaIndex：一個簡單、靈活的數據框架，用于將自定義數據源連接到LLMs
SimpleDirectoryReader：將本地文件數據加載到LlamaIndex的最簡單方式
KnowledgeGraphIndex：從非結構化文本自動構建知識圖譜
SimpleGraphStore：簡單的圖存儲索引
PyVis：一個Python庫，用于可視化和構建圖網絡

1.2 啟用診斷日志

日志可提供代碼執行情況的寶貴信息。

import os, logging, sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

1.3 連接Huggingface API

請更新至你的Hugging Face推理API端點。

from huggingface_hub import login

os.environ["HF_KEY"] = "Your Hugging Face access token goes here"
login(token=os.environ.get('HF_KEY'),add_to_git_credential=True)

1.4 加載PDF文檔

Innovate BC Innovator Skills Initiative
BC Arts Council Application Assistance

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader(input_dir="/content/", required_exts=".pdf").load_data()

2 構建知識圖譜索引

2.1 使用HuggingFace創建本地嵌入

HuggingFaceEmbeddings 類是 LangChain 庫的一部分，它封裝了 Hugging Face 提供的句子轉換器模型，用于創建文本嵌入。這個類支持調用 Hugging Face 平臺上所有可用的句子嵌入模型，以執行包括語義搜索、文檔聚類和問答等任務。在本次練習中，采用了 Hugging Face 的 multi-qa-MiniLM-L6-cos-v1 句子轉換器模型。

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

EMBEDDING_MODEL_NAME = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"

embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL_NAME, embed_batch_size=10)

2.2 從ServiceContext遷移到Settings

LlamaIndex v0.10.0 版本推出了全新的全局 Settings 對象，用以取代之前的 ServiceContext 配置方式。

這個 Settings 對象充當著全局配置的角色，采用了按需實例化的機制。例如，LLM 或嵌入模型等屬性，只有在相關底層模塊真正需要時，才會進行加載。這樣的設計使得資源的使用更加高效，響應更為迅速。

from llama_index.core import Settings

Settings.embed_model = embed_model
Settings.chunk_size = 256
Settings.chunk_overlap = 50

文檔在索引時會被劃分成有重疊的小塊，這個過程稱為“分塊”。通常塊的大小設為1024，重疊部分為20。鑒于我們的文檔較短，我們調整塊大小為256，并設置重疊為50，以優化處理。

2.3 定義自定義提示

from llama_index.core import PromptTemplate

system_prompt = """<|SYSTEM|># You are an AI-enabled admin assistant.
Your goal is to answer questions accurately using only the context provided.
"""

# 這將包裝 llama-index 內部的默認提示
query_wrapper_prompt = PromptTemplate("<|USER|>{query_str}<|ASSISTANT|>")

LLM_MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"

2.4 設置LLM

import torch
from llama_index.llms.huggingface import HuggingFaceLLM

llm = HuggingFaceLLM(
    context_window=4096,
    max_new_tokens=512,
    generate_kwargs={"temperature": 0.1, "do_sample": False},
    system_prompt=system_prompt,
    query_wrapper_prompt=query_wrapper_prompt,
    tokenizer_name=LLM_MODEL_NAME,
    model_name=LLM_MODEL_NAME,
    device_map="auto",
    # 如果使用CUDA以減少內存使用，請取消注釋此行
    model_kwargs={"torch_dtype": torch.float16 , "load_in_8bit":True}
)

Settings.llm = llm

2.5 構建知識圖譜索引

from llama_index.core.storage.storage_context import StorageContext
from llama_index.core import KnowledgeGraphIndex
from llama_index.core.graph_stores import SimpleGraphStore

# 設置存儲上下文
graph_store = SimpleGraphStore()
storage_context = StorageContext.from_defaults(graph_store=graph_store)

index = KnowledgeGraphIndex.from_documents(documents=documents,
                                           max_triplets_per_chunk=3,
                                           storage_cnotallow=storage_context,
                                           embed_model=embed_model,
                                           include_embeddings=True)

max_triplets_per_chunk：指的是每個文本塊中能夠提取的最大三元組數量。降低這個數值可以提升處理效率，因為它減少了需要處理的三元組數量。
include_embeddings：用于決定索引中是否包含嵌入項。默認設置為False，因為生成嵌入項在計算上可能相當耗費資源。

2.6 可視化知識圖譜

from pyvis.network import Network
g = index.get_networkx_graph()
net = Network(notebook=True, cdn_resources="in_line", directed=True)
net.from_nx(g)
net.save_graph("rag_graph.html")
from IPython.display import HTML, display
HTML(filename="rag_graph.html")

利用LlamaIndex和本地PDF文檔，輕松打造知識圖譜GraphRAG-AI.x社區

2.7 查詢

query_engine = index.as_query_engine(llm=llm, similarity_top_k=5)

done = False
while not done:
 print("*"*30)
 question = input("Enter your question: ")
 response = query_engine.query(question)
 print(response)
 done = input("End the chat? (y/n): ") == "y"