This is a final project of UCSB CS291A 2024 Fall.
To run this project, you need to set up following the original EfficientRAG.
You need to install PyTorch >= 2.1.0 first, and then install dependent Python libraries by running the command
pip install -r requirements.txtYou can also create a conda environment with python>=3.9
conda create -n <ENV_NAME> python=3.9 pip
conda activate <ENV_NAME>
pip install -r requirements.txt-
For convenience, we include a very small subset (500 examples for training and 100 examples for testing) from MuSiQue to enable running a small demo. The data have been preprocessed and stored under
data. To test the entire dataset, please construct training data by following the Preparation and Training Data Construction parts in EfficientRAG. -
Download the retriever model Contriever and base model DeBERTa, put them under
model_cache -
Deploy LLaMA-3-8B-Instruct with vLLM framework, and configure it in
src/language_models/llama.py
We will use the MuSiQue dataset as an example. You could train on 2WikiMQA in the same way.
Training Filter model
python src/efficient_rag/filter_training.py \
--dataset musique \
--save_path saved_models/filterTraining Labeler model
python src/efficient_rag/labeler_training.py \
--dataset musique \
--tags 2EfficientRAG retrieve procedure
python src/efficientrag_retrieve.py \
--dataset musique \
--retriever contriever \
--labels 2 \
--labeler_ckpt <<PATH_TO_LABELER_CKPT>> \
--filter_ckpt <<PATH_TO_FILTER_CKPT>> \
--topk 10 \TextRank Pruning
python src/textrank.py --fpath <<MODEL_INFERENCE_RESULT>> --top_k 10You can try different top_k to control the pruning proportion.
Use LLaMA-3-8B-Instruct as generator
python src/efficientrag_qa.py \
--fpath <<MODEL_INFERENCE_RESULT_AFTER_PRUNING>> \
--model llama-8B \
--dataset musiqueRetrieve results
python src/evaluation/retrieve.py --fpath <<QA_RESULT>>Correctness
python src/evaluation/correctness.py \
--fpath <<QA_RESULT>>
--model llama-8b-instructIf you find this paper or code useful, please cite by:
@inproceedings{zhuang2024efficientrag,
title={EfficientRAG: Efficient Retriever for Multi-Hop Question Answering},
author={Zhuang, Ziyuan and Zhang, Zhiyang and Cheng, Sitao and Yang, Fangkai and Liu, Jia and Huang, Shujian and Lin, Qingwei and Rajmohan, Saravan and Zhang, Dongmei and Zhang, Qi},
booktitle={Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing},
pages={3392--3411},
year={2024}
}