利用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"