thunlp/KG-Infused-RAG

# KG-Infused RAG 本仓库是 [KG-Infused RAG: Augmenting Corpus-Based RAG with External Knowledge Graphs](https://arxiv.org/pdf/2506.09542) 的官方实现。 ## 🧭 概述 我们提出了 ``KG-Infused RAG``，这是一个集成预先存在的大规模 KG 以增强检索和生成的框架。其核心是，KG-Infused RAG 采用了 🧠``spreading activation``，这是认知心理学中的一个概念，指的是激活在语义网络中从中心概念向相关概念传播。 The framework consists of three modules that together enable interpretable, fact-grounded multi-source retrieval: 1) KG-Guided Spreading Activation, 2) KG-Based Query Expansion, and 3) KG-Augmented Answer Generation. ## 🛠️ 设置 ### 1. 环境 第 1 步：克隆本仓库 ``` git clone git@github.com:thunlp/KG-Infused-RAG.git cd KG-Infused-RAG ``` 第 2 步：创建环境并安装依赖 ``` conda create -n kg-infused-rag python=3.10 conda activate kg-infused-rag pip install -r requirements.txt pip install -e . ``` ### 2. 数据准备 #### 数据集 所有评估和训练数据均可在[此处](https://drive.google.com/drive/folders/1GfkVZhjHJidCprIyeUuIGdyVJgauC-c1?usp=drive_link)下载。请将数据放置在 `./data/datasets` 目录下。数据分为两部分： - 评估数据：包括原始测试集，以及来自语料库和知识图谱的相应初始检索结果。 - 训练数据：包含训练集上的原始采样输出，以及从中构建的 DPO 训练数据。 #### 语料库与知识图谱 下载语料库，解压文件并将提取的数据放置在 `./data/corpus` 目录下： ``` mkdir -p data/corpus cd data/corpus wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz gunzip psgs_w100.tsv.gz ``` 我们实验中使用的知识图谱可通过 🤗 [Hugging Face](https://huggingface.co/datasets/Alphonse7/Wikidata5M-KG) 和 [ModelScope](https://modelscope.cn/datasets/alphonse7/Wikidata5m-KG) 获取。下载该 KG（Wikidata5M-KG），解压文件并将提取的数据放置在 `./data/KG` 目录下： ``` mkdir -p data/KG tar -xzvf wikidata5m_kg.tar.gz -C data/KG tar -xzvf wikidata5m_raw_data.tar.gz -C data/KG ``` 为语料库中的段落和 Wikidata5M-KG 中的实体描述生成 embedding： ``` bash ./scripts/generate_embeddings_corpus.sh bash ./scripts/generate_embeddings_kg.sh ``` ### 3. 模型 - Retriever：`Contriever-MS MARCO` - Generator：来自 `Qwen2.5` 系列、`LLaMA3.1` 系列以及 `GPT` 系列的模型，包括 GPT-4o-mini 和 GPT-4o。 ## 🚀 实现与实验 ### （可选）初始检索 您可以为每个输入问题从语料库中预先检索 top-k 段落，并从 KG 中检索相关实体，以便在主 pipeline 中**节省时间**： ``` bash ./scripts/retrieval.sh ``` ### 1. 主 Pipeline ``` bash ./scripts/kg_infused_rag.sh ``` ### 2. 训练（DPO） #### 2.1 构建 DPO 数据 **a) 收集采样数据**：多次运行主 pipeline，并在目标生成阶段改变生成参数（例如 `temperature`、`top_p`），以获得多样化的响应。将原始输出保存在诸如 `./data/sampling_results_for_dpo_data_construction/{stage}/` 的文件夹下。 **b) 构建 DPO 对**：使用以下脚本将采样数据转换为 DPO 格式的训练样本： ``` cd train python generate_dpo_data.py ``` 我们在[此处](https://drive.google.com/drive/folders/1GfkVZhjHJidCprIyeUuIGdyVJgauC-c1?usp=drive_link)提供了示例原始采样数据以及构建的 DPO 对。 #### 2.2 运行 DPO 训练 ``` bash train.sh ``` ## 📄 引用 如果您发现我们的代码、数据、模型或论文有用，请引用该[论文](https://arxiv.org/pdf/2506.09542)： ``` @article{wu2025kg, title={KG-Infused RAG: Augmenting Corpus-Based RAG with External Knowledge Graphs}, author={Wu, Dingjun and Yan, Yukun and Liu, Zhenghao and Liu, Zhiyuan and Sun, Maosong}, journal={arXiv preprint arXiv:2506.09542}, year={2025} } ```

标签：IaC 扫描, RAG, TruffleHog, 人工智能, 大模型, 学术论文, 用户模式Hook绕过, 逆向工具