kg-gen：从任意文本生成知识图谱

📄 论文 | 🐍 软件包 | 🤖 MCP | 🔬 实验 | 👩🏻‍💻 数据集 | 🐦 X动态

欢迎！kg-gen 帮助您使用 AI 从任何纯文本中提取知识图谱。它可以处理小型和大型文本输入，也能处理对话格式的消息。

为什么要生成知识图谱？如果您想：

• 创建一个用于 RAG（检索增强生成）的图谱
• 为模型训练和测试创建图结构的合成数据
• 将任意文本结构化为图谱
• 分析源文本中概念之间的关系 那么 kg-gen 非常适合。

我们通过 LiteLLM 支持基于 API 和本地的模型提供商，包括 OpenAI、Ollama、Anthropic、Gemini、Deepseek 等。我们还使用 DSPy 进行结构化输出生成。

• 您可以通过运行 tests/ 中的脚本来试用。
• 运行我们的 KG 基准测试 MINE 的说明位于 MINE/。
• 阅读论文：KGGen：利用语言模型从纯文本中提取知识图谱

由您选择的模型驱动

传入一个 model 字符串以使用您选择的模型。模型调用通过 LiteLLM 路由，通常 LiteLLM 使用 {model_provider}/{model_name} 的格式。具体格式请参阅 https://docs.litellm.ai/docs/providers。

您可以传入的模型示例：

• openai/gpt-5
• gemini/gemini-2.5-flash
• ollama_chat/deepseek-r1:14b

您还可以通过 base_url 指定自定义的 API 基础 URL（示例在此）。

快速入门

安装模块：

pip install kg-gen

然后导入并使用 kg-gen。您可以采用两种格式提供文本输入：

1. 单个字符串
2. 消息对象列表（每个对象包含角色和内容）

以下是一些示例代码：

from kg_gen import KGGen

# 使用可选配置初始化 KGGen
kg = KGGen(
  model="openai/gpt-4o",  # 默认模型
  temperature=0.0,        # 默认温度
  api_key="YOUR_API_KEY"  # 如果已在环境变量中设置或使用本地模型，则可选
)

# 示例 1：带上下文的单个字符串
text_input = "Linda 是 Josh 的母亲。Ben 是 Josh 的兄弟。Andrew 是 Josh 的父亲。"
graph_1 = kg.generate(
  input_data=text_input,
  context="家庭关系"
)
# 输出：
# entities={'Linda', 'Ben', 'Andrew', 'Josh'} 
# edges={'is brother of', 'is father of', 'is mother of'} 
# relations={('Ben', 'is brother of', 'Josh'), 
#           ('Andrew', 'is father of', 'Josh'), 
#           ('Linda', 'is mother of', 'Josh')}

可视化 KG

KGGen.visualize(graph, output_path, open_in_browser=True)

更多示例——分块、聚类、传递消息数组

# 示例 2：带有分块和聚类的大文本
with open('large_text.txt', 'r') as f:
  large_text = f.read()
  
# 示例输入文本：
# """
# 神经网络是一种机器学习模型。深度学习是机器学习的一个子集，
# 它使用多层神经网络。监督学习需要训练数据来学习模式。
# 机器学习是一种人工智能技术，使计算机能够从数据中学习。
# 人工智能，也称为 AI，与更广泛的人工智能领域相关。
# 神经网络（NN）常用于机器学习应用。机器学习（ML）已经彻底改变了许多研究领域。
# ...
# """

graph_2 = kg.generate(
  input_data=large_text,
  chunk_size=5000,  # 按每 5000 字分块处理
  cluster=True      # 对相似实体和关系进行聚类
)
# 输出：
# entities={'neural networks', 'deep learning', 'machine learning', 'AI', 'artificial intelligence', 
#          'supervised learning', 'unsupervised learning', 'training data', ...} 
# edges={'is type of', 'requires', 'is subset of', 'uses', 'is related to', ...} 
# relations={('neural networks', 'is type of', 'machine learning'),
#           ('deep learning', 'is subset of', 'machine learning'),
#           ('supervised learning', 'requires', 'training data'),
#           ('machine learning', 'is type of', 'AI'),
#           ('AI', 'is related to', 'artificial intelligence'), ...}
# entity_clusters={
#   'artificial intelligence': {'AI', 'artificial intelligence'},
#   'machine learning': {'machine learning', 'ML'},
#   'neural networks': {'neural networks', 'neural nets', 'NN'}
#   ...
# }
# edge_clusters={
#   'is type of': {'is type of', 'is a type of', 'is a kind of'},
#   'is related to': {'is related to', 'is connected to', 'is associated with'
#  ...}
# }

# 示例 3：消息数组
messages = [
  {"role": "user", "content": "法国的首都是什么？"}, 
  {"role": "assistant", "content": "法国的首都是巴黎。"}
]
graph_3 = kg.generate(input_data=messages)
# 输出：
# entities={'Paris', 'France'} 
# edges={'has capital'} 
# relations={('France', 'has capital', 'Paris')}

# 示例 4：合并多个图谱
text1 = "Linda 是 Joe 的母亲。Ben 是 Joe 的兄弟。"

# 输入文本 2：Joseph 也被称为 Joe。
text2 = "Andrew 是 Joseph 的父亲。Judy 是 Andrew 的妹妹。Joseph 也被称为 Joe。"

graph4_a = kg.generate(input_data=text1)
graph4_b = kg.generate(input_data=text2)

# 合并图谱
combined_graph = kg.aggregate([graph4_a, graph4_b])

# 可选地对合并后的图谱进行聚类
clustered_graph = kg.cluster(
  combined_graph,
  context="家庭关系"
)
# 输出：
# entities={'Linda', 'Ben', 'Andrew', 'Joe', 'Joseph', 'Judy'} 
# edges={'is mother of', 'is father of', 'is brother of', 'is sister of'} 
# relations={('Linda', 'is mother of', 'Joe'),
#           ('Ben', 'is brother of', 'Joe'),
#           ('Andrew', 'is father of', 'Joe'),
#           ('Judy', 'is sister of', 'Andrew')}
# entity_clusters={
#   'Joe': {'Joe', 'Joseph'},
#   ...
# }
# edge_clusters={ ... }

从本仓库安装：

克隆此仓库，并使用 pip install -e '.[dev]' 安装依赖项。

您可以通过在根目录下运行 python tests/test_basic.py 来验证其是否正常工作。这也会在 tests/test_basic.html 中生成一个漂亮的可视化效果。

用于 AI 代理的 MCP 服务器

对于需要持久化记忆能力的 AI 代理：

# 安装并启动 MCP 服务器
pip install kg-gen
kggen mcp

# 与 Claude Desktop、自定义 MCP 客户端或其他 AI 应用程序一起使用

有关详细的设置和集成说明，请参阅 MCP 服务器文档。

功能

大文本分块处理

对于大文本，您可以指定 chunk_size 参数，以较小的块来处理文本：

graph = kg.generate(
  input_data=large_text,
  chunk_size=5000  # 按每 5000 个字符为一个块进行处理
)

聚类相似实体和关系

您可以在生成过程中或生成之后对相似的实体和关系进行聚类：

# 在生成过程中
graph = kg.generate(
  input_data=text,
  cluster=True,
  context="用于指导聚类的可选上下文"
)

# 或在生成之后
clustered_graph = kg.cluster(
  graph,
  context="用于指导聚类的可选上下文"
)

聚合多个图

您可以使用聚合方法将多个图合并在一起：

graph1 = kg.generate(input_data=text1)
graph2 = kg.generate(input_data=text2)
combined_graph = kg.aggregate([graph1, graph2])

消息数组处理

在处理消息数组时，kg-gen：

1. 会保留每条消息的角色信息
2. 维持消息的顺序和边界
3. 可以提取实体和关系：
   • 消息中提到的概念之间的关系
   • 发言者（角色）与概念之间的关系
   • 整个对话中多条消息之间的关系

例如，对于以下对话：

messages = [
  {"role": "user", "content": "法国的首都是什么？"},
  {"role": "assistant", "content": "法国的首都是巴黎。"}
]

生成的图可能包含以下实体：

• “user”
• “assistant”
• “France”
• “Paris”

以及以下关系：

• (user, “询问关于”, “France”)
• (assistant, “指出”, “Paris”)
• (Paris, “是……的首都”, “France”)

引用

如果您觉得 KGGen 很有用，请考虑引用：

@misc{mo2025kggenextractingknowledgegraphs,
      title={KGGen: 使用语言模型从纯文本中提取知识图谱}, 
      author={Belinda Mo 和 Kyssen Yu 和 Joshua Kazdan 和 Joan Cabezas 和 Proud Mpala 和 Lisa Yu 和 Chris Cundy 和 Charilaos Kanatsoulis 和 Sanmi Koyejo},
      year={2025},
      eprint={2502.09956},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2502.09956}, 
}

许可证

MIT 许可证。

kg-gen 快速上手指南

kg-gen 是一个基于 AI 的开源工具，能够从任意纯文本（包括长文档和对话记录）中提取结构化知识图谱。它支持通过 LiteLLM 调用各类大模型（如 OpenAI、Ollama、DeepSeek 等），并利用 DSPy 确保输出结构的规范性，适用于 RAG 增强、数据合成及关系分析等场景。

环境准备

• 系统要求：支持 Windows、macOS 和 Linux。
• Python 版本：建议 Python 3.8 及以上版本。
• 前置依赖：
  • 需具备可用的大模型 API Key（如 OpenAI、Anthropic 等）或本地模型服务（如 Ollama）。
  • 网络环境需能访问 PyPI 源及对应的模型服务商接口。

提示：国内开发者若使用 pip 安装较慢，可临时切换至清华或阿里镜像源：

export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

通过 PyPI 直接安装最新稳定版：

pip install kg-gen

若需从源码安装以获取开发功能：

git clone https://github.com/stair-lab/kg-gen.git
cd kg-gen
pip install -e '.[dev]'

验证安装是否成功（将生成测试图谱及可视化文件）：

python tests/test_basic.py

基本使用

1. 初始化与简单提取

导入库并配置模型即可开始提取。以下示例展示如何从一段简单的家庭关系描述中生成图谱：

from kg_gen import KGGen

# 初始化 KGGen，指定模型和 API Key
kg = KGGen(
  model="openai/gpt-4o",      # 支持格式：provider/model_name
  temperature=0.0,
  api_key="YOUR_API_KEY"      # 若已设置环境变量可省略
)

# 输入文本
text_input = "Linda is Josh's mother. Ben is Josh's brother. Andrew is Josh's father."

# 生成知识图谱
graph = kg.generate(
  input_data=text_input,
  context="Family relationships"
)

# 查看结果
print(graph.entities)   # 实体集合
print(graph.relations)  # 关系三元组

2. 可视化图谱

工具内置了可视化功能，可直接将生成的图谱保存为 HTML 并在浏览器打开：

KGGen.visualize(graph, output_path="my_graph.html", open_in_browser=True)

3. 进阶功能简述

• 处理长文本：设置 chunk_size 参数自动分块处理，并开启 cluster=True 进行实体聚类。
• 对话数据：直接传入包含 role 和 content 的消息列表（Messages Array）。
• 图谱合并：使用 kg.aggregate() 合并多个独立生成的图谱。

# 长文本处理示例
graph_large = kg.generate(
  input_data=large_text,
  chunk_size=5000,
  cluster=True
)

# 对话数据处理示例
messages = [
  {"role": "user", "content": "What is the capital of France?"}, 
  {"role": "assistant", "content": "The capital of France is Paris."}
]
graph_chat = kg.generate(input_data=messages)