SELF-RAG

To run evaluations using an online retrieval model, you can use a subset of Wikipedia. This requires downloading a demo corpus and embeddings (approx. 9GB).

1. Clone the repository and run the download script:

   git clone git@github.com:AkariAsai/self-rag.git
   cd retrieval_lm
   bash download_demo_corpus.sh

2. Use the Retriever class to search for documents and feed them into the model prompt.

from passage_retrieval import Retriever

# Initialize retriever
retriever = Retriever({})
# Setup demo retriever with specific paths
retriever.setup_retriever_demo(
    "facebook/contriever-msmarco", 
    "enwiki_2020_intro_only/enwiki_2020_dec_intro_only.jsonl", 
    "enwiki_2020_intro_only/enwiki_dec_2020_contriever_intro/*",  
    n_docs=5, 
    save_or_load_index=False
)

# Search for documents
query_3 = "What is overfitting?"
retrieved_documents = retriever.search_document_demo(query_3, 5)

# Format prompts for the model using retrieved docs
# Note: format_prompt is defined in the quickstart section
prompts = [format_prompt(query_3, doc["title"] + "\n" + doc["text"]) for doc in retrieved_documents]
preds = model.generate(prompts, sampling_params)

# Display results
top_doc = retriever.search_document_demo(query_3, 1)[0]
print("Reference: {0}\nModel prediction: {1}".format(top_doc["title"] + "\n" + top_doc["text"], preds[0].outputs[0].text))

Creating training data for the Generator involves a multi-step pipeline. Note that this process is time-consuming for datasets larger than 10k instances. The high-level workflow is:

1. Process input data: Prepare initial data according to the required schema.
2. Critic Evaluation (Retrieval Tokens): Run the critic to judge retrieval necessity (both initial and multi-step).
3. Critic Evaluation (isUse): Run the critic to judge utility.
4. Run Contriever: Perform initial and continuous retrieval.
5. Create isRel and isSUP input data: Prepare prompts for relevance and support evaluation.
6. Critic Evaluation (isRel): Run the critic to judge relevance.
7. Critic Evaluation (isSup): Run the critic to judge support/groundedness.
8. Combine Data: Merge all outputs into the final training dataset.

If you prefer not to run this pipeline, you can download a pre-made training dataset of 150K instances here.

To fine-tune a model (e.g., Llama2-7B) as a Critic, use the train_special_tokens.py script. This process expands the vocabulary with reflection tokens. You can either generate training data using GPT-4 via scripts in data_creation/critic or download existing training data from the provided Google Drive link.

Ensure you have the training data file ready before running the command.

cd data_creation
torchrun --nproc_per_node=2 \
  --master_port=2568 train_special_tokens.py \
  --model_name_or_path meta-llama/Llama-2-7b-hf \
  --data_path PATH_TO_TRAIN_DATA_FILE \
  --bf16  True \
  --output_dir PATH_TO_CRITIC_MODEL \
  --num_train_epochs 3  \
  --per_device_train_batch_size 1 --per_device_eval_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --evaluation_strategy "no" \
  --save_strategy "steps" \
  --save_steps 300 \
  --save_total_limit 1 \
  --learning_rate 2e-5 \
  --weight_decay 0. \
  --warmup_ratio 0.01 \
  --lr_scheduler_type "cosine" \
  --logging_steps 10 \
  --fsdp "full_shard auto_wrap"

For long-form QA, use run_long_form_static.py with pre-retrieved passages. This avoids the high memory requirements of running a retriever (like DPR/Contriever) during inference.

Key Inference Parameters:

• w_rel (default 1.0): Emphasis on isRel (relevance) token probability during beam search.
• w_sup (default 1.0): Emphasis on isSup (support) token probability during beam search.
• w_use (default 0.5): Emphasis on isUse (utility/quality) token probability during beam search.
• threshold (default 0.2): Controls frequency of adaptive retrieval.
• max_depth (default 6): Maximum depth of search ($T$ in the paper).
• beam_width (default 2): Size of the beam in segment-level beam search.

python run_long_form_static.py \
  --model_name selfrag/selfrag_llama2_7b \
  --ndocs 5 --max_new_tokens 300 --threshold 0.2 \
  --use_grounding --use_utility --use_seqscore \
  --task asqa --input_file eval_data/asqa_eval_gtr_top100.json \
  --output_file YOUR_OUTPUT_FILE_NAME --max_depth 7 --mode always_retrieve

Contriever is used to retrieve passages for the queries.

1. Initial Retrieval

First, preprocess the data:

python create_retrieval_data.py \
    --input_files INPUT_FILE \
    --output_file INITIAL_RETRIEVAL_INPUT

Then, run Contriever (requires navigating to retrieval_lm):

cd retrieval_lm
python passage_retrieval.py \
    --model_name_or_path facebook/contriever-msmarco \
    --passages PATH_TO_CORPUS --passages_embeddings PATH_TO_EMBEDDINGS \
    --data INPUT_FILE \
    --output_dir INITIAL_RETRIEVAL_OUTPUT  --n_docs 10

2. Continuous Retrieval (t > 1)

Continuous retrieval is only performed for queries where initial retrieval necessity was true.

First, create the multi-retrieval input file using the initial retrieval tokens and the initial retrieval results:

python create_retrieval_data.py \
    --input_files INPUT_FILE \
    --output_file MULTI_RETRIEVAL_INPUT \
    --need_retrieval_files INITIAL_RETRIEVAL_TOKEN_OUTPUT \
    --multiple_sent --initial_retrieval_file INITIAL_RETRIEVAL_OUTPUT

Then, run the retrieval:

cd retrieval_lm
python passage_retrieval.py \
    --model_name_or_path facebook/contriever-msmarco \
    --passages PATH_TO_CORPUS --passages_embeddings PATH_TO_EMBEDDINGS \
    --data MULTI_RETRIEVAL_INPUT \
    --output_dir MULTI_RETRIEVAL_OUTPUT  --n_docs 10