OpenMatch/RAG-DDR: 미분 가능한 데이터 보상(Differentiable Data Rewards)을 사용한 검색 증강

ICLR 2025 논문의 소스 코드입니다:

RAG-DDR: Optimizing Retrieval-Augmented Generation Using Differentiable Data Rewards

아래 링크를 클릭하여 저희의 논문, 체크포인트 및 공개 모델을 확인해 주세요:

만약 이 작업이 유용하다면, 저희 논문을 인용하고 별점 🌟을 주시면 감사하겠습니다.

@article{li2024rag,
title={RAG-DDR: Optimizing Retrieval-Augmented Generation Using Differentiable Data Rewards},
author={Li, Xinze and Mei, Sen and Liu, Zhenghao and Yan, Yukun and Wang, Shuo and Yu, Shi and Zeng, Zheni and Chen, Hao and Yu, Ge and Liu, Zhiyuan and others},
... 

이 환경에서 Pip 또는 Conda를 사용하여 다음 패키지를 설치하세요

Python==3.10.14
torch==2.2.1
transformers==4.40.0
... 

RAG-DDR의 lora 체크포인트를 여기서 직접 다운로드하고 병합하거나, 아래 흐름을 따라 RAG-DDR을 학습할 수 있습니다.

❗️참고: 체크포인트 파일에서 Gen_model_Llama3_8b와 Kr_model_for_Llama3_8b는 DDR 학습으로 정렬된 생성(Generation) 및 지식 정제(Knowledge Refinement) 모듈이며, 여기서 Llama3-8b가 생성 모듈로 사용됩니다. Gen_model_Minicpm_2.4b와 Kr_model_for_Minicpm_2.4b는 DDR 학습으로 정렬된 생성 모듈 및 지식 정제 모듈이며, 여기서 Minicpm-2.4b가 생성 모듈로 사용됩니다.

(1) git clone을 사용하여 이 프로젝트를 다운로드하세요:

git clone https://github.com/OpenMatch/RAG-DDR.git
cd RAG-DDR

(2) original train/dev dataset을 구성하세요:

RAG-DDR의 학습/개발 데이터셋을 구성하려면, 논문의 Table 5를 따라 해당 데이터셋을 수집하고 데이터를 다음 jsonl 형식으로 처리할 수 있습니다. 또한, 자체적으로 더 많은 데이터셋을 수집하여 DDR 학습에 사용할 수도 있습니다.

{
'id': str # 데이터의 id. 수집한 데이터에는 '0'부터 'n-1'까지 새로운 id를 재할당해야 하며, 'n'은 전체 데이터 개수입니다.
'question': str # 데이터의 질문.
... 
```}<tool_call|>```json```mermaid```</pre>{

그 후, 위에서 구성된 데이터를 위해 MS MARCO 2.1에서 관련 문서를 검색하려면 bge-large-en-v1.5를 사용할 수 있습니다. MS MARCO 2.1 코퍼스는 여기(❗️참고: 파일명은 `msmarco_v2.1_doc_segmented.tar`입니다)에서 다운로드할 수 있습니다. 구체적으로, 해당 데이터에 대한 관련 문서를 검색하려면 `data` 폴더로 이동하면 됩니다. 모든 문서를 임베딩으로 인코딩하여 `--output_dir`에 저장해야 하므로 bge-large를 사용해야 합니다:

cd data
bash get_embedding.sh

다음으로, 원래의 훈련/개발 데이터셋에 적합한 trec 파일을 가져와야 하며, 이는 `--query_path`에 저장되어 있고 `--trec_save_path`에 저장해야 합니다:

bash retrieve.sh

마지막으로, `--trec_save_path`를 기반으로 각 데이터에 대한 관련 문서를 검색하고 `original train/dev dataset`을 검색된 문서와 함께 `--output_path`에 저장할 수 있습니다:

bash construct_psg_data.sh

(3) `original testing dataset` 구성

`original testing dataset`을 구성하려면, 여기에서 KILT의 테스트 데이터셋을 다운로드하고 논문의 표 5에 따라 KILT에서 해당 테스트 데이터셋을 선택할 수 있습니다. 그 후, 위의 데이터 처리 단계를 기반으로 각 테스트 데이터셋에 대한 관련 문서를 검색하여 `original testing dataset`을 얻을 수 있습니다.

(4) 구성된 데이터를 직접 다운로드:

저희가 처리한 `original train/dev dataset`과 `original testing dataset`도 제공합니다. 훈련 및 테스트 데이터를 구성한 후, RAG-DDR 모델의 훈련을 시작할 수 있습니다. 이전에 구성한 `original train/dev dataset`을 사용하여 DPO(Direct Preference Optimization) 훈련/개발 데이터셋을 구성함으로써 RAG-DDR에서 생성 모델(Generation Model)을 훈련해야 합니다.

(1) 첫 번째 단계: 기본 생성 모델(vanilla Generation Model)로 Minicpm-2.4B-sft와 lama3-8B-Instruct 모델을, 지식 정제 모델(Knowledge Refinement Model)로 lama3-8B-Instruct를 다운로드해야 합니다.

(2) 두 번째 단계: `original train/dev dataset`을 전달해야 합니다.

into의 기본 지식 정제 모듈(Knowledge Refinement Module)로 전달하여 쿼리와 관련 없는 검색된 외부 문서를 필터링하고, 정제된 문서와 함께 데이터셋을 `--output_path`에 저장합니다.

. (❗️참고: 생성 모듈(Generation module)의 후속 생성을 위해 상위 5개 문서만 유지하는 데 지식 정제 기능을 사용해야 하므로, 추론 시간을 줄이려면 `--need_n`을 5로 설정하기만 하면 됩니다.

cd scripts
bash kr_inference.sh

(3) 세 번째 단계: 기본 생성 모델(Generation Model)과 두 번째 단계에서 얻은 정제된 데이터셋을 기반으로 DPO 샘플링 데이터를 생성합니다.

bash gen_forward.sh

(4) 네 번째 단계: 세 번째 단계에서 생성된 데이터 파일들은 `merge_response_file.py`를 사용하여 병합하여 DPO 샘플링 데이터를 얻습니다. 그런 다음 이 DPO 샘플을 기반으로 생성 모델(Generation Model)을 훈련하기 위한 DPO train/dev 데이터셋을 생성합니다.

bash get_gen_data.sh

(5) 다섯 번째 단계: 네 번째 단계에서 얻은 DPO train/dev 데이터셋을 사용하여 생성 모델(Generation Model)을 훈련합니다. 생성 모델로는 Minicpm-2.4B-sft 또는 lama3-8B-Instruct 모델을 선택할 수 있습니다.

bash train_gen.sh

(6) 여섯 번째 단계: 다섯 번째 단계에서 LoRA를 사용하여 훈련된 생성 모델(Generation Model)의 가중치를 결합합니다.

bash merge_lora.sh

지식 정제 모델(Knowledge Refinement Model)을 훈련한 후, 지식 정제 모델을 훈련할 수 있습니다.

(1) 첫 번째 단계: `scripts` 폴더로 이동하여, `original train/dev dataset`과 이전 단계에서 훈련된 생성 모델(Generation Model)을 기반으로 DPO 샘플링 데이터를 생성합니다.

cd scripts
bash kr_gen_forward.sh

(2) 두 번째 단계:
첫 번째 단계에서 생성된 데이터 파일들은 `merge_response_file.py`를 사용하여 병합하여 DPO 샘플링 데이터를 얻습니다. 그런 다음 이 DPO 샘플을 기반으로 지식 정제 모델(Knowledge Refinement Model)을 훈련하기 위한 DPO train/dev 데이터셋을 생성합니다.

bash get_kr_data.sh

(3) 세 번째 단계: 두 번째 단계에서 얻은 DPO train/dev 데이터셋을 사용하여 지식 정제 모델(Knowledge Refinement Model)을 훈련합니다.

bash train_kr.sh

RAG-DDR 모델을 훈련한 후에는 위에서 구성한 테스트 데이터로 RAG-DDR의 성능을 테스트할 수 있습니다. 검색된 문서를 입력하는 테스트 데이터셋을 RAG-DDR 모델에 공급하고, 지식 모듈(knowledge module)과 생성 모듈(generation module)을 순차적으로 거쳐 최종 답변을 얻어야 합니다.

(1) 첫 번째 단계: `원본 테스트 데이터셋`을 지식 정제 모듈(Knowledge Refinement Module)로 전달하여 쿼리와 관련 없는 검색된 외부 문서를 필터링하고, 정제된 문서가 포함된 데이터셋을 `--output_path`에 저장해야 합니다.

cd scripts
bash kr_inference.sh

(2) 두 번째 단계: 정제된 데이터셋을 얻은 후에는 이를 생성 모듈(Generation Module)에 공급하여 응답을 생성하고 RAG-DDR의 효과를 평가할 수 있습니다.

bash gen_inference.sh

다른 작업을 위해서는 `--task`, `--metric`, 그리고 `--max_new_token`에 대해 다른 작업 식별자(Task Identifiers), 평가 지표 식별자(metrics identifiers), 및 생성 최대 토큰 수를 설정해야 합니다:

| TASK | Task Identifiers | metrics identifiers | max tokens |
|---|---|---|---|
| NQ | nq | accuracy | 32 |
| ... | 

질문, 제안 및 버그 보고가 있으시면 다음 이메일 주소로 보내주십시오:

2490257@stu.neu.edu.cn