157 `PropertyGraphIndex`类的`_insert_nodes_to_vector_index`方法

原创已于 2024-09-26 11:12:15 修改 · 341 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#LLM

于 2024-09-26 11:12:08 首次发布

llamindex文章专栏收录该内容

171 篇文章

订阅专栏

Llama Index中的属性图索引：深入解析`_insert_nodes_to_vector_index`方法

在现代数据科学和人工智能领域，属性图（Property Graph）已成为处理复杂信息的重要工具。属性图通过结构化的方式表示实体及其关系，使得信息的检索和理解变得更加高效。本文将深入探讨Llama Index中的PropertyGraphIndex类的_insert_nodes_to_vector_index方法，帮助程序员全面理解其工作原理及实际应用。

前置知识

在开始之前，确保你具备以下基础知识：

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_to_vector_index`方法

_insert_nodes_to_vector_index方法负责将向量节点插入到向量索引中。它接受一个标记节点列表作为输入，并执行一系列操作来处理这些节点，包括创建TextNode对象、添加元数据和嵌入向量等。

代码解析

def _insert_nodes_to_vector_index(self, nodes: List[LabelledNode]) -> None:
    """Insert vector nodes."""
    assert self.vector_store is not None

    llama_nodes: List[TextNode] = []
    for node in nodes:
        if node.embedding is not None:
            llama_nodes.append(
                TextNode(
                    text=str(node),
                    metadata={VECTOR_SOURCE_KEY: node.id, **node.properties},
                    embedding=[*node.embedding],
                )
            )
            if not self.vector_store.stores_text:
                llama_nodes[-1].id_ = node.id

        # clear the embedding to save memory, its not used now
        node.embedding = None

    self.vector_store.add(llama_nodes)

设计思路

断言检查：
- 在插入向量节点之前，进行断言检查，确保vector_store不为空。
创建TextNode对象：
- 遍历输入的标记节点列表，对于每个节点，如果其嵌入向量不为空，则创建一个TextNode对象。
- TextNode对象包含节点的文本表示、元数据（包括向量源ID和节点属性）以及嵌入向量。
处理向量存储的文本存储：
- 如果向量存储不存储文本，则将TextNode对象的ID设置为节点的ID。
清除嵌入向量：
- 为了节省内存，清除节点的嵌入向量。
插入向量节点：
- 将创建的TextNode对象列表插入到向量存储中。

代码示例

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)

# 插入节点到向量索引
pg_index._insert_nodes_to_vector_index(nodes)