TIGER-AI-Lab/LongRAG

이 레포지토리는 "LongRAG: Long-context LLMs를 활용한 검색 증강 생성(Retrieval-Augmented Generation) 개선"의 코드를 포함합니다. 현재 저희는 이 레포지토리를 다듬는 과정에 있습니다.

• 서론
• 설치
• 빠른 시작
• 코퍼스 준비 (선택 사항)
• Long Retriever
• Long Reader
• 라이선스
• 인용

전통적인 RAG 프레임워크에서 기본적인 검색 단위는 보통 짧습니다. 이러한 설계는 리트리버(retriever)가 '바늘' 단위를 찾기 위해 방대한 코퍼스를 탐색하도록 강제합니다. 반면, 리더(reader)는 짧게 검색된 단위에서 답변을 추출하기만 하면 됩니다. 이처럼 불균형하게 무거운 리트리버와 가벼운 리더의 설계는 최적화되지 않은 성능으로 이어질 수 있습니다. 저희는 'long retriever'와 'long reader'로 구성된 새로운 프레임워크 LongRAG를 제안합니다. 저희 프레임워크는 4K-token 검색 단위를 사용하며, 이는 이전보다 30배 더 긴 것입니다. 본 연구는 RAG와 long-context LLMs를 결합하는 미래 로드맵에 대한 통찰력을 제공합니다.

이 레포지토리를 클론하고 필요한 패키지를 설치하세요:

git clone https://github.com/TIGER-AI-Lab/LongRAG.git
cd LongRAG
pip install -r requirements.txt

'Long Reader' 섹션으로 이동하여 지침을 따르세요. 이렇게 하면 100개 예제에 대한 최종 예측 결과를 얻는 데 도움이 될 것입니다.
출력은 exp/ 디렉토리의 샘플 파일과 유사할 것입니다.

이것은 선택적인 단계입니다. 처리된 코퍼스를 직접 사용할 수 있습니다. 저희는 NQ와 HotpotQA에 대해 Hugging Face에 검색 코퍼스의 두 가지 버전을 공개했습니다.

from datasets import load_dataset
corpus_nq = load_dataset("TIGER-Lab/LongRAG", "nq_corpus")
corpus_hotpotqa = load_dataset("TIGER-Lab/LongRAG", "hotpot_qa_corpus")

만약 저희가 코퍼스를 어떻게 구성했는지 여전히 관심이 있다면, 여기서 읽기 시작할 수 있습니다.

• Wikipedia 원본 데이터 정리:*
  저희는 먼저 표준 프로세스를 따라 Wikipedia 원본 데이터를 정리합니다. WikiExtractor를 사용합니다.
  이는 Wikipedia 데이터베이스 백업 덤프에서 텍스트를 추출하고 정리하는 데 널리 사용되는 Python 스크립트입니다. 필요한 Python 환경을 사용하는지 확인해 주세요. 샘플 스크립트는 다음과 같습니다:

* 위키피디아 데이터 전처리 *
위키피디아 원본 데이터를 정리한 후, 다음 스크립트를 실행하여 더 많은 정보를 수집합니다.

`sh scripts/process_wiki_page.sh`

`dir_path`
: 이전 단계의 출력인 정리된 위키피디아 덤프 디렉터리 경로입니다.
`output_path_dir`
: 이 출력 디렉터리에는 각기 다른 위키피디아 페이지에 대한 딕셔너리를 나타내는 여러 pickle 파일이 포함됩니다.
`degree.pickle`
: 키는 위키피디아 페이지 제목이며, 값은 하이퍼링크의 개수입니다.
`abs_adj.pickle`
: 키는 위키피디아 페이지 제목이며, 값은 초록(abstract) 단락에 링크된 페이지입니다.
`full_adj.pickle`
: 키는 위키피디아 페이지 제목이며, 값은 전체 페이지에 링크된 페이지입니다.
`doc_size.pickle`
: 키는 위키피디아 페이지 제목이며, 값은 해당 페이지의 토큰 개수입니다.
`doc_dict.pickle`
: 키는 위키피디아 페이지 제목이며, 값은 페이지 텍스트입니다.
`corpus_title_path`
: NQ 데이터셋을 필터링하는 데 사용되는 키입니다. 원래 DPR 논문에서는 목록 페이지(list pages)나 용어 구분 페이지(disambiguation pages)와 같은 특정 위키피디아 페이지를 제거하여, 전체 위키피디아 페이지 수를 5백만 개에서 3백만 개로 줄였습니다. 공정한 비교를 위해 저희도 이러한 페이지들을 제외하기로 했습니다. (HotpotQA의 경우, 어떤 페이지도 제거하지 않았으므로 위키피디아 페이지 수는 5백만 개를 유지합니다.) DPR의 제목은 이 링크에서 다운로드할 수 있습니다.

저희는 처리된 위키피디아 데이터를 huggingface 리포지토리에 제공했습니다. 더 많은 정보가 필요하면 `nq_wiki`와 `hotpot_qa_wiki` 서브셋을 확인해 주세요. 이 두 데이터셋으로부터 쉽게 이러한 pickle 파일들을 파생시킬 수 있습니다.

* 검색 코퍼스 (Retrieval Corpus):* 여러 관련 문서를 그룹화하여 4K 토큰을 초과하는 긴 검색 단위(long retrieval units)를 구성할 수 있습니다. 이러한 설계는 코퍼스 크기(검색 단위의 개수)를 크게 줄일 수도 있습니다. 그러면 리트리버(retriever)의 작업이 훨씬 쉬워집니다. 또한, 긴 검색 단위는 모호성이나 혼란을 피하기 위해 정보 완전성을 향상시킬 것입니다.

`sh scripts/group_documents.sh`

`processed_wiki_dir`

: 위의 단계의 출력 디렉터리입니다.`mode`

:`abs`

은 HotpotQA 코퍼스용이며, `full`은 NQ 코퍼스용입니다.`output_dir`

: 출력 디렉터리입니다. 이 출력 디렉터리에는 검색 코퍼스를 위한 딕셔너리를 나타내는 여러 pickle 파일이 포함됩니다. 가장 중요한 것은 코퍼스 ID를 코퍼스 텍스트에 매핑하는 `group_text.pickle`입니다. 더 자세한 내용은 Hugging Face에서 공개된 코퍼스를 참조해 주십시오.

저희는 오픈 소스 밀집 검색 도구 키트인 Tevatron을 활용합니다. 모든 검색 실험에 사용됩니다. 저희가 사용한 기본 임베딩 모델은 bge-large-en-v1.5입니다. 샘플 스크립트를 제공했으니, 반드시 자체 데이터셋의 로컬 경로로 매개변수를 업데이트하십시오. 또한, 저희 스크립트는 시간 절약을 위해 4개의 GPU를 사용하여 코퍼스를 인코딩하므로, 사용 사례에 따라 이 부분을 업데이트해 주십시오.

`sh scripts/run_retrieve_tevatron.sh`

긴 컨텍스트 입력을 처리하는 강력한 능력을 가진 Gemini-1.5-Pro와 GPT-4o를 장기 리더(long reader)로 선택했습니다. (향후 긴 컨텍스트를 처리할 수 있는 다른 LLM도 테스트할 계획입니다.)

리더의 입력은 장기 리트리버에서 나온 모든 긴 검색 단위들의 연결본입니다. 저희는 Huggingface 레포지토리에 입력 파일을 제공했습니다.

mkdir -p exp/
sh scripts/run_eval_qa.sh

`test_data_name`

: 테스트 세트 이름으로, `nq`(NQ) 또는 `hotpot_qa`(HotpotQA)입니다.`test_data_split`

: 각 테스트 세트에는 세 가지 분할(splits)이 있습니다: `full`, `subset_1000`, `subset_100`. 빠른 시작이나 디버깅을 위해 `subset_100`으로 시작하고 `subset_1000`을 사용하는 것을 권장합니다.

결과가 비교적 안정적으로 나오도록 하기 위함입니다. `output_file_path`

: 출력 파일은 여기 `exp/` 디렉토리에 저장됩니다.
`reader_model`

: 우리가 사용하는 장문 컨텍스트 리더 모델로, 현재 코드는 `GPT-4o`, `GPT-4-Turbo`, `Gemini-1.5-Pro`, `Claude-3-Opus`를 지원합니다. 관련 API 키 및 API 설정은 코드에서 업데이트해야 한다는 점에 유의하십시오. 예를 들어, GPT-4 시리즈를 사용하는 경우 `utils/gpt_inference.py`에서 코드를 구성해야 하며, Gemini 시리즈를 사용하는 경우 `utils/gemini_inference.py`에서 코드를 구성해야 합니다. 앞으로 더 많은 모델을 계속 지원할 예정입니다.

출력 파일은 행마다 하나의 테스트 케이스를 포함합니다. `short_ans` 필드가 우리의 최종 예측값입니다.

{
"query_id": "383",
"question": "how many episodes of touching evil are there",
...

우리는 `exp/` 디렉토리에 몇 가지 샘플 출력 파일을 제공했습니다. 예를 들어, `exp/nq_gpt4o_100.json`에는 다음 파일 실행 결과가 포함되어 있습니다:

python eval/eval_qa.py
--test_data_name "nq"
--test_data_split "subset_100"
...

상위-1 검색 정확도는 88%이며, 정답 일치율은 64%입니다.

저희 작업에 사용된 각 서브셋의 라이선스를 확인해 주십시오.

| Dataset Name | License Type |
|---|---| 
| NQ | Apache License 2.0 |
| HotpotQA | CC BY-SA 4.0 License |

저희 프로젝트가 유용하다고 생각하시면 논문을 인용해 주시면 감사하겠습니다:

@article{jiang2024longrag,
title={LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs},
author={Ziyan Jiang, Xueguang Ma, Wenhu Chen},
...}