PaperAgent 2024年07月10日
源码解读 - 微软GraphRAG框架
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

微软最新开源的基于知识图谱构建的检索增强生成（RAG）系统GraphRAG，旨在利用大型语言模型（LLMs）从非结构化文本中提取结构化数据，构建具有标签的知识图谱，以支持数据集问题生成、摘要问答等多种应用场景。GraphRAG 的一大特色是利用图机器学习算法针对数据集进行语义聚合和层次化分析，因而可以回答一些相对高层级的抽象或总结性问题，这一点恰好是常规 RAG 系统的短板。

🤔 GraphRAG 框架的创新点在于利用图机器学习算法进行语义聚合和层次化分析，能够回答常规 RAG 系统无法处理的高层级抽象或总结性问题，例如“该数据集的主题是什么”。

💻 GraphRAG 的实现基于两个阶段：索引 (Indexing) 和查询 (Query)。在索引阶段，利用 LLM 提取节点、边和协变量，并使用社区检测技术将知识图谱划分成模块化的社区，然后用 LLM 总结每个社区。

🚀 在查询阶段，针对特定查询，汇总所有相关社区的摘要，生成全局性的答案。

🔍 GraphRAG 的源码解析揭示了其系统架构、关键概念和核心工作流。源码结构清晰，包含缓存、配置、输入、输出、模型、提示等目录，并使用 Datashaper 框架执行工作流处理。

🛠️ GraphRAG 的索引阶段涉及多个工作流，例如创建文本单元、提取实体、总结实体、构建实体图、创建社区、合并文本单元等，每个工作流都有相应的代码实现，并利用拓扑排序确保执行顺序。

奔跑的日月 2024-07-10 11:33 湖北

今天分享一篇PaperAgent-RAG专栏技术交流群小伙伴（@知乎 奔跑的日月）关于微软最新开源的GraphRAG框架源码解读文章（已授权转载）。

文章有5千字，建议收藏观看，知乎原文：

https://zhuanlan.zhihu.com/p/707759736

1. 引言

这几天微软开源了一个新的基于知识图谱构建的检索增强生成（RAG）系统, GraphRAG, 该框架旨在利用大型语言模型（LLMs）从非结构化文本中提取结构化数据, 构建具有标签的知识图谱，以支持数据集问题生成、摘要问答等多种应用场景。GraphRAG 的一大特色是利用图机器学习算法针对数据集进行语义聚合和层次化分析，因而可以回答一些相对高层级的抽象或总结性问题, 这一点恰好是常规 RAG 系统的短板。说实话之前一直有在关注这个框架, 所以这两天花了点时间研究了一下源码, 结合之前的一些技术文档，本文主要是记录 GraphRAG 源码方面的一些解读, 也希望借此进一步理解其系统架构、关键概念以及核心工作流等。

本次拉取的 GraphRAG 项目源码对应 commit id 为 a22003c302bf4ffeefec76a09533acaf114ae7bb, 更新日期为 2024.07.05。

2. 框架概述

讨论代码前, 我们先简单了解下GraphRAG项目的目标与定位. 在论文中, 作者很明确地提出了一个常规RAG无法处理的应用场景:

However, RAG fails on global questions directed at an entire text corpus, such as “What are the main themes in the dataset?”, since this is inherently a queryfocused summarization (QFS) task, rather than an explicit retrieval task.

也就是类似该数据集的主题是什么这种high level的总结性问题, 作者认为, 这种应用场景本质上一种聚焦于查询的总结性(QueryFocused Summarization, QFS)任务, 单纯只做数据检索是无法解决的. 相应的, 其解决思路也在论文中清楚地描述出来了:

In contrast with related work that exploits the structured retrieval and traversal affordances of graph indexes (subsection 4.2), we focus on a previously unexplored quality of graphs in this context: their inherent modularity (Newman, 2006) and the ability of community detection algorithms to partition graphs into modular communities of closely-related nodes (e.g., Louvain, Blondel et al., 2008; Leiden, Traag et al., 2019). LLM-generated summaries of these community descriptions provide complete coverage of the underlying graph index and the input documents it represents. Query-focused summarization of an entire corpus is then made possible using a map-reduce approach: first using each community summary to answer the query independently and in parallel, then summarizing all relevant partial answers into a final global answer.

利用社区检测算法（如Leiden算法）将整个知识图谱划分模块化的社区(包含相关性较高的节点), 然后大模型自下而上对社区进行摘要, 最终再采取map-reduce方式实现QFS: 每个社区先并行执行Query, 最终汇总成全局性的完整答案.

实现方式是什么(How)?

Figure 1, Source https://arxiv.org/pdf/2404.16130

论文中给出了解决问题的基本思路, 与其他RAG系统类似, GraphRAG整个Pipeline也可划分为索引(Indexing)与查询(Query)两个阶段。索引过程利用LLM提取出节点（如实体）、边（如关系）和协变量（如 claim），然后利用社区检测技术对整个知识图谱进行划分，再利用LLM进一步总结。最终针对特定的查询，可以汇总所有与之相关的社区摘要生成一个全局性的答案.

3. 源码解析

官方文档说实话写得已经很清楚了, 不过想要理解一些实现上的细节, 还得深入到源码当中. 接下来, 一块看下代码的具体实现. 项目源码结构树如下:

.├── cache├── config├── emit├── graph│   ├── embedding│   ├── extractors│   │   ├── claims│   │   ├── community_reports│   │   ├── graph│   │   └── summarize│   ├── utils│   └── visualization├── input├── llm├── progress├── reporting├── storage├── text_splitting├── utils├── verbs│   ├── covariates│   │   └── extract_covariates│   │       └── strategies│   │           └── graph_intelligence│   ├── entities│   │   ├── extraction│   │   │   └── strategies│   │   │       └── graph_intelligence│   │   └── summarize│   │       └── strategies│   │           └── graph_intelligence│   ├── graph│   │   ├── clustering│   │   │   └── strategies│   │   ├── embed│   │   │   └── strategies│   │   ├── layout│   │   │   └── methods│   │   ├── merge│   │   └── report│   │       └── strategies│   │           └── graph_intelligence│   ├── overrides│   └── text│       ├── chunk│       │   └── strategies│       ├── embed│       │   └── strategies│       ├── replace│       └── translate│           └── strategies└── workflows    └── v1

Demo

研究具体功能前, 先简单跑下官方demo, 上手也很简单, 直接参考 Get Started (microsoft.github.io) 即可.

高能预警: 虽然只是一个简单demo, 但是Token消耗可是一点都不含糊, 尽管早有预期, 并且提前删除了原始文档超过一半的内容, 不过我这边完整跑下来还是花了差不多3刀费用, 官方完整demo文档跑一遍, 预计得消耗5~10刀

这里实际运行时间还是比较慢的, 大模型实际上是来来回回的在过整个文档, 其中一些比较重要的事项如下:

├── cache│   ├── community_reporting│   │   ├── create_community_report-chat-v2-0d811a75c6decaf2b0dd7b9edff02389│   │   ├── create_community_report-chat-v2-1205bcb6546a4379cf7ee841498e5bd4│   │   ├── create_community_report-chat-v2-1445bd6d097492f734b06a09e579e639│   │   ├── ...│   ├── entity_extraction│   │   ├── chat-010c37f5f6dedff6bd4f1f550867e4ee│   │   ├── chat-017a1f05c2a23f74212fd9caa4fb7936│   │   ├── chat-09095013f2caa58755e8a2d87eb66fc1│   │   ├── ...│   ├── summarize_descriptions│   │   ├── summarize-chat-v2-00e335e395c5ae2355ef3185793b440d│   │   ├── summarize-chat-v2-01c2694ab82c62924080f85e8253bb0a│   │   ├── summarize-chat-v2-03acd7bc38cf2fb24b77f69b016a288a│   │   ├── ...│   └── text_embedding│       ├── embedding-07cb902a76a26b6f98ca44c17157f47f│       ├── embedding-3e0be6bffd1c1ac6a091f5264858a2a1│       ├── ...├── input│   └── book.txt├── output│   └── 20240705-142536│       ├── artifacts│       │   ├── create_base_documents.parquet│       │   ├── create_base_entity_graph.parquet│       │   ├── create_base_extracted_entities.parquet│       │   ├── create_base_text_units.parquet│       │   ├── create_final_communities.parquet│       │   ├── create_final_community_reports.parquet│       │   ├── create_final_documents.parquet│       │   ├── create_final_entities.parquet│       │   ├── create_final_nodes.parquet│       │   ├── create_final_relationships.parquet│       │   ├── create_final_text_units.parquet│       │   ├── create_summarized_entities.parquet│       │   ├── join_text_units_to_entity_ids.parquet│       │   ├── join_text_units_to_relationship_ids.parquet│       │   └── stats.json│       └── reports│           ├── indexing-engine.log│           └── logs.json├── prompts│   ├── claim_extraction.txt│   ├── community_report.txt│   ├── entity_extraction.txt│   └── summarize_descriptions.txt└── settings.yaml

这个文件中的很多文档都值得仔细研究, 后续将结合代码详细说明.

此外, console 中会打印很多运行日志, 其中比较重要的一条就是完整的workflows, 会涉及到完整pipeline的编排:

⠹ GraphRAG Indexer ├── Loading Input (InputFileType.text) - 1 files loaded (0 filtered) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 0:00:00├── create_base_text_units├── create_base_extracted_entities├── create_summarized_entities├── create_base_entity_graph├── create_final_entities├── create_final_nodes├── create_final_communities├── join_text_units_to_entity_ids├── create_final_relationships├── join_text_units_to_relationship_ids├── create_final_community_reports├── create_final_text_units├── create_base_documents└── create_final_documents  All workflows completed successfully.

Index

索引阶段整体看下来应该算是整个项目的核心, 整体流程还是比较复杂的. 执行indexing的语句如下:

python -m graphrag.index --init --root ./ragtest
python -m graphrag.index --root ./ragtest

Figure 2, Indexing 函数调用关系图

简单跟一下, 发现实际调用的是 graphrag/index/__main__.py文件中的主函数, 使用 argparse 解析输入参数, 实际调用的是 graphrag/index/cli.py 中的 index_cli 函数.

继续解读源码前, 先简单看下相关函数的调用链路, 如上图所示, 其中灰色标记的函数是我们需要重点关注的.

限于篇幅, 我们在此只讨论默认配置运行流程, 大致梳理清楚相关逻辑后, 可自行修改相关配置.

result = PipelineConfig(
    root_dir=settings.root_dir,
    input=_get_pipeline_input_config(settings),
    reporting=_get_reporting_config(settings),
    storage=_get_storage_config(settings),
    cache=_get_cache_config(settings),
    workflows=[
        *_document_workflows(settings, embedded_fields),
        *_text_unit_workflows(settings, covariates_enabled, embedded_fields),
        *_graph_workflows(settings, embedded_fields),
        *_community_workflows(settings, covariates_enabled, embedded_fields),
        *(_covariate_workflows(settings) if covariates_enabled else []),    ],
)

这段代码的基本逻辑就是根据不同的功能生成完整的workflow序列, 此处需要注意的是这里并不考虑workflow之间的依赖关系, 单纯基于workflows/v1目录下的各workflow的模板生成一系列的workflow.

这里有必要再单独说明一下run.py::run_pipeline(), 该函数用于真正的执行所有pipeline, 其核心逻辑包含以下两部分:

await dump_stats()
for workflow_to_run in workflows_to_run:
    # Try to flush out any intermediate dataframes
    gc.collect()
    workflow = workflow_to_run.workflow
    workflow_name: str = workflow.name
    last_workflow = workflow_name
    log.info("Running workflow: %s...", workflow_name)
    if is_resume_run and await storage.has(
        f"{workflow_to_run.workflow.name}.parquet"    ):
        log.info("Skipping %s because it already exists", workflow_name)
        continue
    stats.workflows[workflow_name] = {"overall": 0.0}    await inject_workflow_data_dependencies(workflow)
    workflow_start_time = time.time()
    result = await workflow.run(context, callbacks)    await write_workflow_stats(workflow, result, workflow_start_time)
    # Save the output from the workflow
    output = await emit_workflow_output(workflow)
    yield PipelineRunResult(workflow_name, output, None)
    output = None
    workflow.dispose()
    workflow = None
stats.total_runtime = time.time() - start_time
await dump_stats()

根据以上信息, 我们可以大致梳理出索引环节的完整工作流

Workflow

截至目前, 我们实际上还没有真正分析index阶段的业务逻辑, 只是搞清楚了GraphRAG内置的这套pipeline编排系统该如何工作. 这里以index/workflows/v1/create_final_entities.py为例, 一起看下具体的一个workflow是如何运行的.

from datashaper import TableContainer, VerbCallbacks, VerbInput, progress_iterable, verb
@verb(name="cluster_graph")
def cluster_graph(
    input: VerbInput,    callbacks: VerbCallbacks,
    strategy: dict[str, Any],
    column: str,
    to: str,
    level_to: str | None = None,
    **_kwargs,
) -> TableContainer:

可以看出, 实际上就是加了一个verb装饰器而已, 进一步跟进 strategy 的实现可以发现, 这里的leiden算法实际上也是源自另一个图算法库 graspologic-org/graspologic: Python package for graph statistics (github.com)

Query

查询阶段的pipeline相对而言要简单些, 执行query的语句如下:

# Global searchpython -m graphrag.query \--root ./ragtest \--method global \"What are the top themes in this story?"# Local searchpython -m graphrag.query \--root ./ragtest \--method local \
"Who is Scrooge, and what are his main relationships?"

这里需要注意的是有Global/Local search两种模式. 还是先生成函数调用关系图, 对整体结构能有个大致了解.

Figure 3, Query 函数调用关系图

graphrag/query/__main__.py中的主函数会依据参数不同, 分别路由至cli::run_local_search()以及cli::run_global_search()

Global Search

cli::run_global_search()主要调用了factories.py::get_global_search_engine()函数, 返回一个详细定义的GlobalSearch类, 进一步跟进去, 发现该类跟LocalSearch类相似, 都是基于工厂模式创建, 其核心方法为structured_search/global_search/search.py::GlobalSearch.asearch(), 具体使用了map-reduce方法, 首先使用大模型并行地为每个社区的summary生成答案, 然后再汇总所有答案生成最终结果, 此处建议参考map/reduce的相关prompts, 作者给出了非常详细的说明.

也正是因为这种map-reduce的机制, 导致global search对token的消耗量极大.

Local Search

与全局搜索类似, cli::run_local_search()函数主要也是调用了factories.py::get_local_search_engine(), 返回一个LocalSearch类, 这里的asearch()相对比较简单, 会直接根据上下文给出回复, 这种模式更接近于常规的RAG语义检索策略, 所消耗的Token也比较少.

与全局搜索不同的地方在于, Local 模式综合了nodes, community_reports, text_units, relationships, entities, covariates等多种源数据, 这一点在官方文档中也给出了非常详细的说明, 不再赘述.

4. 一些思考

Reference

推荐阅读

欢迎关注我的公众号“PaperAgent”，每天一篇大模型（LLM）文章来锻炼我们的思维，简单的例子，不简单的方法，提升自己。


Fish AI Reader

Fish AI Reader

AI辅助创作，多种专业模板，深度分析，高质量内容生成。从观点提取到深度思考，FishAI为您提供全方位的创作支持。新版本引入自定义参数，让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅，AI 时代的下一个智能信息助手，助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

GraphRAG RAG 知识图谱 大型语言模型 社区检测
相关文章
Organizing Knowledge With Knowledge Graphs: Industry Trends
Is Claude 3 Outperforming GPT-4?
Graphs and Language
Harmonizing AI: Crafting Personalized Song Suggestions
AI News Weekly - Issue #377: Next in AI : Pioneers' Predictions! - Mar 21st 2024
COLLAGE: A New Machine Learning Approach to Deal with Floating-Point Errors in Low-Precision to Make LLM Training Accurate and Efficient
Databricks Announces Major Updates to Its AI Suite to Boost AI Model Accuracy
Leveraging Linguistic Expertise in NLP: A Deep Dive into RELIES and Its Impact on Large Language Models
Japanese Researchers Release “Fugaku-LLM” Trained on the Fugaku Supercomputer
GraphRAG: Knowledge Graphs for AI Applications with Kirk Marple - #681