Hannibal046/xRAG: 단일 토큰을 이용한 검색 증강 생성 (RAG)용 극한 컨텍스트 압축

xRAG의 공식 저장소입니다: 단일 토큰을 이용한 검색 증강 생성 (Retrieval-augmented Generation, RAG)을 위한 극한 컨텍스트 압축 (Extreme Context Compression)

필요한 패키지는 Dockerfile을 참조하세요.

wandb와 accelerate를 설정하세요.

wandb login
accelerate config

HuggingFace

모델 (Model)	백본 (Backbone)	다운로드 (Download)
xRAG-7b	mistralai/Mistral-7B-Instruct-v0.2	🤗 Hugging Face
xRAG-MoE	mistralai/Mixtral-8x7B-Instruct-v0.1	🤗 Hugging Face

tutorial.ipynb에서 xRAG를 위한 튜토리얼을 제공합니다.

확인해 보세요!

• 사전 학습 (pretraining) 데이터 및 검색 (retrieval)을 위한 코퍼스 (corpus)로 enwiki-dec2021을 다운로드하세요.

• prepare_data.ipynb에서 지시어 튜닝 (instruction tuning) 데이터를 준비하세요.

• TriviaQA를 다운로드하세요.

• ColBERT-v2를 사용하여 검색 (retrieval)을 수행하세요.

훈련 스크립트는 scripts/에 있습니다.
예를 들어, SFR을 사용하여 Mistral-7b를 훈련하려면:

accelerate launch \
--mixed_precision bf16 \
--num_machines 1 \
...

평가 (evaluation) 코드는 src/eval에 있습니다.
예를 들어, TriviaQA에서 평가하려면:

검색 증강 (retrieval augmentation) 없이:

CUDA_VISIBLE_DEVICES=0 python -m src.eval.run_eval \
--data triviaqa \
--model_name_or_path Hannibal046/xrag-7b

검색 증강 (retrieval augmentation)을 사용하여:

CUDA_VISIBLE_DEVICES=0 python -m src.eval.run_eval \
--data triviaqa \
--model_name_or_path Hannibal046/xrag-7b \
...

xRAG를 사용하여:

CUDA_VISIBLE_DEVICES=0 python -m src.eval.run_eval \
--data triviaqa \
--model_name_or_path Hannibal046/xrag-7b \
...

xRAG의 벤치마크 (benchmark)를 위해 src/language_modeling/profiler.py에 코드를 제공합니다.

python -m src.language_modeling.profiler --instruction_length 54 --generation_length 30 --dataset triviaqa --use_xrag
python -m src.language_modeling.profiler --instruction_length 54 --generation_length 30 --dataset triviaqa