Llama Index中的属性图索引：深入解析_insert_nodes方法
在现代数据科学和人工智能领域，属性图（Property Graph）已成为处理复杂信息的重要工具。属性图通过结构化的方式表示实体及其关系，使得信息的检索和理解变得更加高效。本文将深入探讨Llama Index中的PropertyGraphIndex类的_insert_nodes方法，帮助程序员全面理解其工作原理及实际应用。
前置知识
在开始之前，确保你具备以下基础知识：
- Python基础：熟悉Python编程。
- OpenAI API密钥：你需要一个OpenAI API密钥来使用
OpenAI模型。 - Llama Index：使用
pip install llama-index安装Llama Index库。
环境设置
首先，让我们通过安装所需的包并配置OpenAI API密钥来设置环境。
# 安装Llama Index
%pip install llama-index
# 设置OpenAI API密钥
import os
os.environ["OPENAI_API_KEY"] = "sk-..."
# 配置日志
import logging
import sys
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
PropertyGraphIndex的_insert_nodes方法
_insert_nodes方法负责将节点插入到属性图索引结构中。它接受一个节点序列作为输入，并执行一系列操作来处理这些节点，包括提取三元组、嵌入节点、过滤重复节点等。
代码解析
def _insert_nodes(self, nodes: Sequence[BaseNode]) -> Sequence[BaseNode]:
"""Insert nodes to the index struct."""
if len(nodes) == 0:
return nodes
# run transformations on nodes to extract triplets
if self._use_async:
nodes = asyncio.run(
arun_transformations(
nodes, self._kg_extractors, show_progress=self._show_progress
)
)
else:
nodes = run_transformations(
nodes, self._kg_extractors, show_progress=self._show_progress
)
# ensure all nodes have nodes and/or relations in metadata
assert all(
node.metadata.get(KG_NODES_KEY) is not None
or node.metadata.get(KG_RELATIONS_KEY) is not None
for node in nodes
)
kg_nodes_to_insert: List[LabelledNode] = []
kg_rels_to_insert: List[Relation] = []
for node in nodes:
# remove nodes and relations from metadata
kg_nodes = node.metadata.pop(KG_NODES_KEY, [])
kg_rels = node.metadata.pop(KG_RELATIONS_KEY, [])
# add source id to properties
for kg_node in kg_nodes:
kg_node.properties[TRIPLET_SOURCE_KEY] = node.id_
for kg_rel in kg_rels:
kg_rel.properties[TRIPLET_SOURCE_KEY] = node.id_
# add nodes and relations to insert lists
kg_nodes_to_insert.extend(kg_nodes)
kg_rels_to_insert.extend(kg_rels)
# filter out duplicate kg nodes
kg_node_ids = {node.id for node in kg_nodes_to_insert}
existing_kg_nodes = self.property_graph_store.get(ids=list(kg_node_ids))
existing_kg_node_ids = {node.id for node in existing_kg_nodes}
kg_nodes_to_insert = [
node for node in kg_nodes_to_insert if node.id not in existing_kg_node_ids
]
# filter out duplicate llama nodes
existing_nodes = self.property_graph_store.get_llama_nodes(
[node.id_ for node in nodes]
)
existing_node_hashes = {node.hash for node in existing_nodes}
nodes = [node for node in nodes if node.hash not in existing_node_hashes]
# embed nodes (if needed)
if self._embed_kg_nodes:
# embed llama-index nodes
node_texts = [
node.get_content(metadata_mode=MetadataMode.EMBED) for node in nodes
]
if self._use_async:
embeddings = asyncio.run(
self._embed_model.aget_text_embedding_batch(
node_texts, show_progress=self._show_progress
)
)
else:
embeddings = self._embed_model.get_text_embedding_batch(
node_texts, show_progress=self._show_progress
)
for node, embedding in zip(nodes, embeddings):
node.embedding = embedding
# embed kg nodes
kg_node_texts = [str(kg_node) for kg_node in kg_nodes_to_insert]
if self._use_async:
kg_embeddings = asyncio.run(
self._embed_model.aget_text_embedding_batch(
kg_node_texts, show_progress=self._show_progress
)
)
else:
kg_embeddings = self._embed_model.get_text_embedding_batch(
kg_node_texts,
show_progress=self._show_progress,
)
for kg_node, embedding in zip(kg_nodes_to_insert, kg_embeddings):
kg_node.embedding = embedding
# if graph store doesn't support vectors, or the vector index was provided, use it
if self.vector_store is not None and len(kg_nodes_to_insert) > 0:
self._insert_nodes_to_vector_index(kg_nodes_to_insert)
if len(nodes) > 0:
self.property_graph_store.upsert_llama_nodes(nodes)
if len(kg_nodes_to_insert) > 0:
self.property_graph_store.upsert_nodes(kg_nodes_to_insert)
# important: upsert relations after nodes
if len(kg_rels_to_insert) > 0:
self.property_graph_store.upsert_relations(kg_rels_to_insert)
# refresh schema if needed
if self.property_graph_store.supports_structured_queries:
self.property_graph_store.get_schema(refresh=True)
return nodes
设计思路
-
检查节点数量：
- 如果输入的节点列表为空，直接返回原节点列表。
-
运行转换以提取三元组：
- 如果启用了异步模式，使用
arun_transformations方法异步运行转换。 - 否则，使用
run_transformations方法同步运行转换。
- 如果启用了异步模式，使用
-
确保节点包含元数据：
- 断言检查所有节点是否包含
KG_NODES_KEY或KG_RELATIONS_KEY元数据。
- 断言检查所有节点是否包含
-
处理节点和关系：
- 从节点的元数据中提取三元组节点和关系。
- 将源节点ID添加到三元组节点的属性中。
- 将三元组节点和关系添加到待插入列表中。
-
过滤重复的三元组节点：
- 获取现有三元组节点的ID，并过滤掉重复的三元组节点。
-
过滤重复的Llama节点：
- 获取现有Llama节点的哈希值，并过滤掉重复的Llama节点。
-
嵌入节点（如果需要）：
- 如果启用了节点嵌入，对Llama节点和三元组节点进行嵌入。
- 使用异步或同步方法获取节点文本的嵌入向量。
-
插入节点到向量存储（如果需要）：
- 如果向量存储已配置且有三元组节点待插入，调用
_insert_nodes_to_vector_index方法插入节点。
- 如果向量存储已配置且有三元组节点待插入，调用
-
插入Llama节点和三元组节点：
- 将Llama节点和三元组节点插入到属性图存储中。
-
插入关系：
- 将关系插入到属性图存储中。
-
刷新模式（如果需要）：
- 如果属性图存储支持结构化查询，刷新模式。
-
返回处理后的节点：
- 返回处理后的节点列表。
代码示例
from llama_index.core import PropertyGraphIndex, StorageContext
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.llms.openai import OpenAI
from llama_index.core.node_parser import SentenceSplitter
# 定义LLM
llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
# 创建属性图存储
property_graph_store = SimplePropertyGraphStore()
storage_context = StorageContext.from_defaults(graph_store=property_graph_store)
# 创建PropertyGraphIndex
pg_index = PropertyGraphIndex(
storage_context=storage_context,
llm=llm,
embed_kg_nodes=True,
show_progress=True,
)
# 创建节点解析器
node_parser = SentenceSplitter()
documents = [...] # 假设有一些文档
nodes = node_parser.get_nodes_from_documents(documents)
# 插入节点
inserted_nodes = pg_index._insert_nodes(nodes)
print(inserted_nodes)
总结
通过Llama Index的PropertyGraphIndex，我们可以轻松地将节点插入到属性图索引结构中。_insert_nodes方法的设计思路清晰，能够帮助程序员更好地理解和应用属性图技术。希望这篇博客能帮助你更好地理解和应用属性图技术。
扫一扫
3924
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
