LlamaIndex 教程：自定义属性图谱索引以提升 GraphRAG 准确率

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

这是 Neo4J 的客座文章

属性图索引是 LlamaIndex 的一个出色补充，也是对之前知识图谱集成的升级。首先，数据表示略有不同。在之前的集成中，图是用三元组表示的，但现在我们有了一个真正的属性图集成，其中节点具有标签，并且可选地具有节点属性。

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

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

该图展示了 LlamaIndex 中的属性图集成，从文档被传递给图构建器开始。这些构建器是负责提取结构化信息的模块化组件，提取的信息随后存储在知识图谱中。图可以使用各种或自定义模块来构建，突显了系统适应不同数据源或提取需求的灵活性。

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

在这篇博客文章中，你将学习如何：

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

代码可在 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 节点标签稍微有些模糊，它允许 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 上获取。