Tomaz Bratanic • 2024-06-11

在 LlamaIndex 中定制属性图索引

了解如何实现实体去重和自定义检索方法以提高 GraphRAG 准确性

这是 Neo4j 的一篇客座文章

属性图索引是 LlamaIndex 的一个极好补充，也是对之前知识图谱集成的升级。首先，数据表示方式略有不同。在之前的集成中，图谱用三元组表示，但现在我们有了适当的属性图集成，其中节点具有标签和可选的节点属性。

每个节点都被分配一个标签来指示其类型，例如 Person、Organization、Project 或 Department。节点和关系也可以存储节点属性，用于其他相关详细信息，例如出生日期或项目的开始和结束日期，如本例所示。

其次，属性图索引被设计成模块化的，因此您可以使用一个或多个（自定义）知识图谱构建器以及检索器，使其成为构建您的第一个知识图谱或根据您的特定需求定制实现方式的绝佳工具。

图片说明了 LlamaIndex 中的属性图集成，首先是将文档传递给图谱构建器。这些构建器是模块化组件，负责提取结构化信息，然后将其存储在知识图谱中。图谱可以使用各种或自定义模块构建，这突显了系统适应不同数据源或提取需求的灵活性。

然后，图谱检索器访问知识图谱以检索数据。此阶段也是模块化的，允许使用多个检索器或旨在查询图中特定类型数据或关系的自定义解决方案。最后，检索到的数据由 LLM 用于生成答案，代表该过程的输出或从中获得的洞察。这种流程强调了一个高度适应性和可扩展性的系统，其中每个组件都可以独立修改或替换，以增强整体功能或根据特定需求进行调整。

在这篇博文中，您将学习如何

使用模式引导式提取构建知识图谱
使用文本嵌入和词语相似度技术的组合执行实体去重
设计一个自定义图谱检索器
最后，您将使用自定义检索器实现问答流程

代码可在 GitHub 上获取。

环境设置

在这篇博文中，我们将使用 Neo4j 作为底层的图数据库。最简单的入门方法是在 Neo4j Aura 上获取一个免费实例，它提供 Neo4j 数据库的云实例。或者，您也可以通过下载 Neo4j Desktop 应用程序并创建一个本地数据库实例来设置本地 Neo4j 数据库实例。

from llama_index.graph_stores.neo4j import Neo4jPGStore

username="neo4j"
password="stump-inlet-student"
url="bolt://52.201.215.224:7687"

graph_store = Neo4jPGStore(
    username=username,
    password=password,
    url=url,
)

此外，您还需要一个有效的 OpenAI API 密钥。

import os

os.environ["OPENAI_API_KEY"] = "sk-"

数据集

在这篇博文中，我们将使用一个从 Diffbot 获取的新闻文章样本数据集，我已将其发布在GitHub 上以便于访问。

由于属性图索引处理文档，我们需要将新闻中的文本封装为 LlamaIndex 文档。

import pandas as pd
from llama_index.core import Document

news = pd.read_csv(
  "https://raw.githubusercontent.com/tomasonjo/blog-datasets/main/news_articles.csv")
documents = [Document(text=f"{row['title']}: {row['text']}") for i, row in news.iterrows()]

图谱构建

如前所述，LlamaIndex 提供了多种开箱即用的图谱构建器。在本例中，我们将使用 SchemaLLMPathExtractor，它允许我们定义想要从文档中提取的图结构模式。

我们首先定义我们希望 LLM 提取的节点和关系的类型。

entities = Literal["PERSON", "LOCATION", "ORGANIZATION", "PRODUCT", "EVENT"]
relations = Literal[
    "SUPPLIER_OF",
    "COMPETITOR",
    "PARTNERSHIP",
    "ACQUISITION",
    "WORKS_AT",
    "SUBSIDIARY",
    "BOARD_MEMBER",
    "CEO",
    "PROVIDES",
    "HAS_EVENT",
    "IN_LOCATION",
]

如您所见，我们的图谱提取主要围绕人和组织。接下来，我们将指定与每个节点标签相关的关系。

# define which entities can have which relations
validation_schema = {
    "Person": ["WORKS_AT", "BOARD_MEMBER", "CEO", "HAS_EVENT"],
    "Organization": [
        "SUPPLIER_OF",
        "COMPETITOR",
        "PARTNERSHIP",
        "ACQUISITION",
        "WORKS_AT",
        "SUBSIDIARY",
        "BOARD_MEMBER",
        "CEO",
        "PROVIDES",
        "HAS_EVENT",
        "IN_LOCATION",
    ],
    "Product": ["PROVIDES"],
    "Event": ["HAS_EVENT", "IN_LOCATION"],
    "Location": ["HAPPENED_AT", "IN_LOCATION"],
}

例如，一个人可以有以下关系

工作于 (WORKS_AT)
董事会成员 (BOARD_MEMBER)
首席执行官 (CEO)
参与事件 (HAS_EVENT)

除了 EVENT 节点标签，该模式相当具体，EVENT 节点标签略微模糊，允许 LLM 捕获各种类型的信息。

现在我们已经定义了图谱模式，我们可以将其输入到 SchemaLLMPathExtractor 中并用它来构建图谱。

from llama_index.core import PropertyGraphIndex

kg_extractor = SchemaLLMPathExtractor(
    llm=llm,
    possible_entities=entities,
    possible_relations=relations,
    kg_validation_schema=validation_schema,
    # if false, allows for values outside of the schema
    # useful for using the schema as a suggestion
    strict=True,
)

NUMBER_OF_ARTICLES = 250

index = PropertyGraphIndex.from_documents(
    documents[:NUMBER_OF_ARTICLES],
    kg_extractors=[kg_extractor],
    llm=llm,
    embed_model=embed_model,
    property_graph_store=graph_store,
    show_progress=True,
)

此代码从 250 篇新闻文章中提取图谱信息，但您可以根据需要调整数量。总共有 2500 篇文章。

请注意，使用 GPT-4o 提取 250 篇文章大约需要 7 分钟。但是，您可以通过使用 num_workers 参数进行并行化来加速此过程。

我们可以可视化一个小图谱来检查存储的内容。

构建的图谱包含文本块（蓝色），其中包含文本和嵌入。如果在文本块中提到了实体，则文本块与实体之间存在 MENTIONS 关系。此外，实体之间也可以存在关系。

实体去重

实体去重或消歧是图谱构建中重要但经常被忽略的一步。本质上，这是一个清洗步骤，您尝试匹配代表单个实体的多个节点，并将它们合并为一个节点，以获得更好的图结构完整性。

例如，在我们构建的图谱中，我发现了一些可以合并的示例。

我们将使用文本嵌入相似度和词语距离的组合来查找潜在的重复项。我们首先在图谱中的实体上定义向量索引。

graph_store.structured_query("""
CREATE VECTOR INDEX entity IF NOT EXISTS
FOR (m:`__Entity__`)
ON m.embedding
OPTIONS {indexConfig: {
 `vector.dimensions`: 1536,
 `vector.similarity_function`: 'cosine'
}}
""")

接下来的 Cypher 查询用于查找重复项，它相当复杂，我和 Michael Hunger、Eric Monk 花了几个小时才完善它。

similarity_threshold = 0.9
word_edit_distance = 5
data = graph_store.structured_query("""
MATCH (e:__Entity__)
CALL {
  WITH e
  CALL db.index.vector.queryNodes('entity', 10, e.embedding)
  YIELD node, score
  WITH node, score
  WHERE score > toFLoat($cutoff)
      AND (toLower(node.name) CONTAINS toLower(e.name) OR toLower(e.name) CONTAINS toLower(node.name)
           OR apoc.text.distance(toLower(node.name), toLower(e.name)) < $distance)
      AND labels(e) = labels(node)
  WITH node, score
  ORDER BY node.name
  RETURN collect(node) AS nodes
}
WITH distinct nodes
WHERE size(nodes) > 1
WITH collect([n in nodes | n.name]) AS results
UNWIND range(0, size(results)-1, 1) as index
WITH results, index, results[index] as result
WITH apoc.coll.sort(reduce(acc = result, index2 IN range(0, size(results)-1, 1) |
        CASE WHEN index <> index2 AND
            size(apoc.coll.intersection(acc, results[index2])) > 0
            THEN apoc.coll.union(acc, results[index2])
            ELSE acc
        END
)) as combinedResult
WITH distinct(combinedResult) as combinedResult
// extra filtering
WITH collect(combinedResult) as allCombinedResults
UNWIND range(0, size(allCombinedResults)-1, 1) as combinedResultIndex
WITH allCombinedResults[combinedResultIndex] as combinedResult, combinedResultIndex, allCombinedResults
WHERE NOT any(x IN range(0,size(allCombinedResults)-1,1) 
    WHERE x <> combinedResultIndex
    AND apoc.coll.containsAll(allCombinedResults[x], combinedResult)
)
RETURN combinedResult  
""", param_map={'cutoff': similarity_threshold, 'distance': word_edit_distance})
for row in data:
    print(row)

不深入太多细节，我们使用文本嵌入和词语距离的组合来查找图谱中的潜在重复项。您可以调整 similarity_threshold 和 word_distance 来找到最佳组合，以便在不过多产生误报的情况下检测尽可能多的重复项。不幸的是，实体消歧是一个难题，没有完美的解决方案。使用这种方法，我们取得了相当不错的结果，但也存在一些误报。

['1963 AFL Draft', '1963 NFL Draft']
['June 14, 2023', 'June 15 2023']
['BTC Halving', 'BTC Halving 2016', 'BTC Halving 2020', 'BTC Halving 2024', 'Bitcoin Halving', 'Bitcoin Halving 2024']

您可以自行调整参数，并在合并重复节点之前添加一些手动例外。

实现自定义检索器

很好，我们已经基于新闻数据集构建了一个知识图谱。现在，让我们看看我们的检索器选项。目前，有四种现有的检索器可用

LLMSynonymRetriever：接受查询，并尝试生成关键词和同义词来检索节点（以及连接到这些节点的路径）。
VectorContextRetriever：根据节点的向量相似度检索节点，然后获取连接到这些节点的路径
TextToCypherRetriever：使用图数据库模式、您的查询和提示模板来生成和执行 Cypher 查询
CypherTemplateRetriever：与其让 LLM 自由生成任何 Cypher 语句，我们可以提供一个 Cypher 模板，并让 LLM 填充参数。

此外，实现自定义检索器也很简单，这正是我们将在这里做的。我们的自定义检索器将首先识别输入查询中的实体，然后分别为每个识别的实体执行 VectorContextRetriever。

首先，我们将定义实体提取模型和提示。

from pydantic import BaseModel
from typing import Optional, List


class Entities(BaseModel):
    """List of named entities in the text such as names of people, organizations, concepts, and locations"""
    names: Optional[List[str]]


prompt_template_entities = """
Extract all named entities such as names of people, organizations, concepts, and locations
from the following text:
{text}
"""

现在我们可以继续实现自定义检索器。

from typing import Any, Optional

from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.retrievers import CustomPGRetriever, VectorContextRetriever
from llama_index.core.vector_stores.types import VectorStore
from llama_index.program.openai import OpenAIPydanticProgram


class MyCustomRetriever(CustomPGRetriever):
    """Custom retriever with entity detection."""
    def init(
        self,
        ## vector context retriever params
        embed_model: Optional[BaseEmbedding] = None,
        vector_store: Optional[VectorStore] = None,
        similarity_top_k: int = 4,
        path_depth: int = 1,
        include_text: bool = True,
        **kwargs: Any,
    ) -> None:
        """Uses any kwargs passed in from class constructor."""
        self.entity_extraction = OpenAIPydanticProgram.from_defaults(
            output_cls=Entities, prompt_template_str=prompt_template_entities
        )
        self.vector_retriever = VectorContextRetriever(
            self.graph_store,
            include_text=self.include_text,
            embed_model=embed_model,
            similarity_top_k=similarity_top_k,
            path_depth=path_depth,
        )

    def custom_retrieve(self, query_str: str) -> str:
        """Define custom retriever with entity detection.

        Could return `str`, `TextNode`, `NodeWithScore`, or a list of those.
        """
        entities = self.entity_extraction(text=query_str).names
        result_nodes = []
        if entities:
            print(f"Detected entities: {entities}")
            for entity in entities:
                result_nodes.extend(self.vector_retriever.retrieve(entity))
        else:
            result_nodes.extend(self.vector_retriever.retrieve(query_str))
        final_text = "\n\n".join(
            [n.get_content(metadata_mode="llm") for n in result_nodes]
        )
        return final_text

MyCustomRetriever 类只有两个方法。您可以使用 init 方法实例化将在检索器中使用的任何函数或类。在此示例中，我们实例化了实体检测 OpenAI 程序以及向量上下文检索器。

custom_retrieve 方法在检索期间被调用。在我们的自定义检索器实现中，我们首先识别文本中的任何相关实体。如果找到任何实体，我们将迭代并为每个实体执行向量上下文检索器。另一方面，如果没有识别到实体，我们将整个输入传递给向量上下文检索器。

如您所见，您可以通过集成现有检索器或从头开始，轻松地为您的用例定制检索器，因为您可以通过使用图数据库的 structured_query 方法轻松执行 Cypher 语句。

问答流程

让我们使用自定义检索器来回答一个示例问题来结束本文。我们需要将检索器传递给 RetrieverQueryEngine。

from llama_index.core.query_engine import RetrieverQueryEngine

custom_sub_retriever = MyCustomRetriever(
    index.property_graph_store,
    include_text=True,
    vector_store=index.vector_store,
    embed_model=embed_model
)

query_engine = RetrieverQueryEngine.from_args(
    index.as_retriever(sub_retrievers=[custom_sub_retriever]), llm=llm
)

让我们来测试一下！

response = query_engine.query(
    "What do you know about Maliek Collins or Darragh O’Brien?"
)
print(str(response))
# Detected entities: ['Maliek Collins', "Darragh O'Brien"]
# Maliek Collins is a defensive tackle who has played for the Dallas Cowboys, Las Vegas Raiders, and Houston Texans. Recently, he signed a two-year contract extension with the Houston Texans worth $23 million, including a $20 million guarantee. This new deal represents a raise from his previous contract, where he earned $17 million with $8.5 million guaranteed. Collins is expected to be a key piece in the Texans' defensive line and fit well into their 4-3 alignment.
# Darragh O’Brien is the Minister for Housing and has been involved in the State’s industrial relations process and the Government. He was recently involved in a debate in the Dáil regarding the pay and working conditions of retained firefighters, which led to a heated exchange and almost resulted in the suspension of the session. O’Brien expressed confidence that the dispute could be resolved and encouraged unions to re-engage with the industrial relations process.

总结

在这篇博文中，我们探讨了在 LlamaIndex 中定制属性图索引的细节，重点介绍了实现实体去重和设计自定义检索方法以提高 GraphRAG 准确性。属性图索引提供了一种模块化和灵活的方法，可以利用各种图构建器和检索器来根据您的特定需求定制实现。无论您是构建您的第一个知识图谱还是针对独特的数据集进行优化，这些可定制的组件都提供了强大的工具包。我们邀请您测试属性图索引集成，看看它们如何提升您的知识图谱项目。

一如既往，代码可在 GitHub 上获取。

使用 LlamaIndex 工作流构建知识图谱代理
2025-01-15
使用 LlamaIndex 和 Memgraph 构建知识图谱
2024-11-21
介绍属性图索引：一种使用 LLM 构建知识图谱的强大新方法
2024-05-29
使用 LlamaIndex 和 Neo4j 构建多模态 RAG 管道
2023-12-18