PRAG：提升问答系统性能的参数化检索增强生成工具

综合介绍

PRAG（Parametric Retrieval-Augmented Generation）是一个创新的检索增强生成工具，旨在通过将外部知识直接嵌入大语言模型（LLM）的参数空间，提升生成效果。该工具克服了传统上下文检索增强生成方法的局限性，减少了计算开销，并通过深度集成外部知识，增强了模型的推理和综合能力。PRAG 提供了端到端的实现，包括数据增强模块、参数训练模块和推理模块，适用于各种问答数据集的性能测试。

功能列表

• 数据增强模块：将文档转换为数据增强数据集。
• 参数训练模块：训练额外的 LoRA 参数以生成文档的参数化表示。
• 推理模块：合并相关文档的参数化表示，插入 LLM 中进行推理。
• 环境安装：提供详细的环境安装步骤和依赖项。
• 自我增强：支持直接使用预增强数据文件或自行处理数据增强。
• 检索准备：下载和准备 Wikipedia 数据集以进行检索。

使用帮助

环境安装

1. 创建并激活虚拟环境：

   conda create -n prag python=3.10.4
conda activate prag

2. 安装必要的依赖项：

   pip install torch==2.1.0
pip install -r requirements.txt

3. 修改 src/root_dir_path.py 中的 ROOT_DIR 变量为存储 PRAG 的文件夹地址。

数据增强

1. 使用预增强数据文件：

   tar -xzvf data_aug.tar.gz

2. 自行处理数据增强：
   • 下载 Wikipedia 数据集： bash
mkdir -p data/dpr
wget -O data/dpr/psgs_w100.tsv.gz https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz

   • 准备 BM25 检索： bash
# 具体步骤请参考项目文档


参数训练

1. 生成文档的参数化表示：

   # 具体步骤请参考项目文档

推理

1. 合并相关文档的参数化表示，插入 LLM 中进行推理：

   # 具体步骤请参考项目文档