KG_RAG - 解锁医学知识图谱的力量和KG-RAG框架概述

KG-RAG 项目介绍

什么是 KG-RAG？

KG-RAG 是知识图谱增强生成 (Knowledge Graph-based Retrieval Augmented Generation) 的简称。它是一种与特定任务无关的框架，通过结合知识图谱 (KG) 的显性知识与大型语言模型 (LLM) 的隐性知识来增强文本生成。具体来说，KG-RAG 使用一个名为 SPOKE 的大规模生物医学知识图谱来提供生物医学背景。SPOKE 汇集了来自不同领域的 40 多个生物医学知识库，涵盖基因、蛋白质、药物、化合物、疾病及其关联。SPOKE 拥有超过2700万个节点和5300万条边。

KG-RAG 的核心功能在于从 SPOKE KG 中提取“提示感知上下文”，即足以回应用户提示的最小上下文。因此，该框架通过植入一个优化的领域特定的“提示感知上下文”来增强通用 LLM。

KG-RAG 的应用示例

一个典型的例子是 FDA 网站发布了一则关于药物“setmelanotide”的新闻，该药被 FDA 批准用于管理 Bardet-Biedl 综合征患者的体重。当使用 KG-RAG 提问 GPT-4 时，系统能够提供有关该药物的准确信息。

如何运行 KG-RAG

目前，KG-RAG 专为与疾病相关的提示而设计，我们正在积极改进其多样性。以下是运行 KG-RAG 的步骤：

步骤 1：克隆仓库

首先克隆该项目的代码库，所有使用的生物医学数据已上传至该仓库，无需单独下载。

步骤 2：创建虚拟环境

建议使用 Python 3.10.9 运行脚本。

conda create -n kg_rag python=3.10.9
conda activate kg_rag
cd KG_RAG

步骤 3：安装依赖

pip install -r requirements.txt

步骤 4：更新配置文件 config.yaml

确保更新 config.yaml 文件以包含运行脚本所需的所有信息。另一个文件 system_prompts.yaml 已经过配置，用于 KG-RAG 框架中的系统提示。

步骤 5：运行安装脚本

安装脚本将以交互方式运行，用于创建疾病向量数据库并可选下载 Llama 模型。

python -m kg_rag.run_setup

步骤 6：从终端运行 KG-RAG

您可以分别使用 GPT 和 Llama 模型运行 KG-RAG。

使用 GPT

# GPT_API_TYPE='azure'
python -m kg_rag.rag_based_generation.GPT.text_generation -g <选择的GPT型号 - "gpt-4" 或 "gpt-35-turbo">
# GPT_API_TYPE='openai'
python -m kg_rag.rag_based_generation.GPT.text_generation -g <选择的GPT型号 - "gpt-4" 或 "gpt-3.5-turbo">

使用 GPT 互动模式

用户可以逐步互动运行流程

# GPT_API_TYPE='azure'
python -m kg_rag.rag_based_generation.GPT.text_generation -i True -g <选择的GPT型号 - "gpt-4" 或 "gpt-35-turbo">
# GPT_API_TYPE='openai'
python -m kg_rag.rag_based_generation.GPT.text_generation -i True -g <选择的GPT型号 - "gpt-4" 或 "gpt-3.5-turbo">

使用 Llama

如果您未在安装步骤下载 Llama 模型，首次运行时可能需要一些时间下载。

python -m kg_rag.rag_based_generation.Llama.text_generation -m <方法-1 或 方法-2，若不指定则用"方法-1">

使用 Llama 互动模式

用户可逐步进行交互式运行

python -m kg_rag.rag_based_generation.Llama.text_generation -i True -m <方法-1 或 方法-2，若不指定则用"方法-1">

KG-RAG 命令行参数

以下是在命令行调用时可以使用的一些参数：

参数	默认值	含义	可选项	备注
-g	gpt-35-turbo	选择 GPT 模型	GPT models provided by OpenAI	仅适用于 GPT 模型
-i	False	是否启用互动模式（逐步显示过程）	True 或 False	可用于 GPT 和 Llama 模型
-e	False	是否显示图中关联的证据	True 或 False	可用于 GPT 和 Llama 模型
-m	method-1	选择分词器方法	method-1 或 method-2. method-1 使用 'AutoTokenizer' ，method-2 使用 'LlamaTokenizer' 且设置 'legacy' 为 False	仅适用于 Llama 模型

BiomixQA：基准数据集

BiomixQA 是一个经过策划的生物医学问答数据集，用于验证 KG-RAG 框架在不同 LLMs 上的表现。它包括：

多项选择题 (MCQ)
真/假题

该数据集的多样性和覆盖范围使其特别适合支持生物医学自然语言处理、知识图推理和问答系统研究与开发。该数据集可在 Hugging Face 上获取，加载简单：

from datasets import load_dataset

# 加载多项选择题数据
mcq_data = load_dataset("kg-rag/BiomixQA", "mcq")

# 加载真/假题数据
tf_data = load_dataset("kg-rag/BiomixQA", "true_false")

引用

如果您使用了 KG-RAG 项目，请在文献中引用以下文章：

@article{soman2023biomedical,
  title={Biomedical knowledge graph-enhanced prompt generation for large language models},
  author={Soman, Karthik and Rose, Peter W and Morris, John H and Akbas, Rabia E and Smith, Brett and Peetoom, Braian and Villouta-Reyes, Catalina and Cerono, Gabriel and Shi, Yongmei and Rizk-Jackson, Angela and others},
  journal={arXiv preprint arXiv:2311.17330},
  year={2023}
}