RAG-DDR: 미분 가능한 데이터 보상을 사용한 검색 증강 생성 최적화

저희 논문의 소스 코드는 다음과 같습니다:

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 # 데이터의 질문.
...

그 후, 위에서 구성한 데이터를 기반으로 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 getembedding.sh

다음으로, 원래의 train/dev 데이터셋에 적합한 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) train/dev 데이터셋을 구성함으로써 RAG-DDR의 Generation Model을 훈련해야 합니다.

(1) 첫 번째 단계: 기본 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) 첫 번째 단계: original testing dataset을 지식 정제 모듈(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