FusionRAG - Multi-Level Cache-Based Fusion Retrieval Enhanced Generation System

一个面向长文档（>100K / 200K tokens）问答与多跳检索的 Fusion RAG（Retrieval-Augmented Generation）复现与实验框架。

🌟 项目特色

• 图谱增强检索: 结合实体关系图和层次化文档树进行精确检索
• 多数据集支持: 支持NovelQA、InfiniteChoice、InfiniteQA等数据集
• 分层摘要: 构建层次化文档摘要树，提升长文档理解能力
• 多进程优化: 支持并行处理，提升构建和推理效率
• 缓存机制: 智能缓存策略，避免重复计算
• 多种推理模式: 支持标准推理、vLLM优化、缓存优化等模式

📁 项目结构

FusionRAG/
├── src/
│   ├── core/
│   │   ├── build_tree.py          # 构建文档摘要树（分层摘要）
│   │   ├── dataloader.py          # 数据加载器，支持多种数据集格式
│   │   ├── extract_graph.py       # 实体/关系抽取与图结构构建
│   │   ├── query.py               # 检索器，图结构检索与密集检索逻辑
│   │   └── utils.py               # 工具函数（分块、计时、分数计算等）
│   ├── processors/
│   │   ├── process_utils.py       # 多进程任务封装与显存清理
│   │   └── redataloader.py        # 改写数据集加载器
│   ├── evaluation/
│   │   ├── evaluate.py            # 评估脚本，计算EM/RL等指标
│   │   ├── evalute_GraphRAG.py    # GraphRAG系统评估
│   │   └── calculate.py           # 性能对比分析
│   ├── prompts/
│   │   ├── prompt_dict.py         # 各类prompt模板
│   │
│   └── main/
│       ├── main.py                # 主流程入口
│       ├── main_cacheready.py     # 缓存就绪模式
│       ├── main_vllm.py           # vLLM优化模式
│       └── main_cache.py          # 缓存优化模式
├── configs/
│   ├── config.yaml                # 主配置文件
│   ├── qa_config.yaml             # 问答任务配置
│   ├── choice_config.yaml         # 选择题任务配置
│   └── cacheconfig.yaml           # 缓存配置
├── requirements.txt               # 项目依赖
└── README.md                      # 项目说明

🚀 快速开始

环境要求

• Python 3.10+
• CUDA 12.8+ (GPU推理)
• 40GB+ GPU显存 (推荐)

安装依赖

pip install -r requirements.txt

下载模型

需要下载以下模型：

• LLM模型: Qwen2.5-7B-Instruct 或其他兼容模型
• NLP模型: spaCy的英文和中文模型

# 下载spaCy模型
python -m spacy download en_core_web_lg
python -m spacy download zh_core_web_lg

配置文件

编辑 configs/config.yaml 设置模型路径和参数：

dataset:
  dataset_name: InfiniteQALoader
  dataset_path: ./data/longbook_qa_eng.jsonl

llm:
  llm_name: Qwen2.5-7B-Instruct
  llm_path: /path/to/your/model
  llm_device: cuda:0

paths:
  log_path: ./log
  answer_path: ./answer
  cache_path: ./cache

cluster:
  length: 1200
  overlap: 100
  merge_num: 5

🎯 使用方法

1. 标准推理模式

python main.py --config configs/config.yaml

2. vLLM优化模式

python main_vllm.py --config configs/config.yaml

3. 缓存优化模式

使用LMCache进行KV缓存优化：

python main_cache.py --config configs/config.yaml

📊 数据集支持

支持的数据集类型

1. InfiniteChoice: 长文档多选题数据集
2. InfiniteQA: 长文档问答数据集

🔧 核心组件

1. 文档摘要树构建

• 分块处理长文档
• 层次化摘要生成
• 支持多种合并策略

2. 知识图谱抽取

• 基于spaCy的实体识别
• 实体共现关系挖掘
• 图结构存储和检索

3. 混合检索策略

• 图谱路径搜索
• 语义相似度检索
• 多策略融合

4. 推理优化

• 多进程并行处理
• GPU显存智能管理
• KV缓存优化

📈 评估指标

准确性指标

• EM (Exact Match): 精确匹配分数
• Rouge-L: 最长公共子序列F1分数

效率指标

• 构建时间 (树+图谱)
• 推理时间
• 显存使用量

运行评估

python evalute_GraphRAG.py

⚙️ 配置说明

主要参数

参数	说明	默认值
length	文档分块长度	1200
overlap	块间重叠长度	100
merge_num	合并块数量	5
max_chunk_setting	最大检索块数	25