SELF-RAG 教程：通过自我反思学习检索、生成和批判

SELF-RAG 是一套创新的框架，它能够让现有的大型语言模型（LLM）学会检索、生成和批判，从而提高生成内容的准确性和质量，且不影响模型的多功能性。相比于传统的检索增强生成（RAG）方法，SELF-RAG 可以根据不同的查询需求动态检索，并通过“反思”机制对生成结果进行自我批判。

1. SELF-RAG 的核心思想

SELF-RAG 的核心在于通过反思机制提高生成内容的质量。传统的 RAG 框架通常在生成前只检索一次，而 SELF-RAG 可以根据需要动态检索多次，甚至跳过检索。此外，SELF-RAG 通过批判性反思标记来评估生成内容的各个方面，如相关性、支持性和可用性。

• 打比方：想象你在写一篇文章，传统 RAG 就像在写之前先找好所有参考资料，而 SELF-RAG 则像是在写作过程中，随时根据需要查找新资料，并且在每一步都对自己的写作进行批判性反思。
• 速记句：SELF-RAG 通过反思机制动态检索和批判生成内容。
  

2. 如何安装 SELF-RAG

安装 SELF-RAG 的依赖库非常简单。你可以通过 pip 安装所需的库，也可以使用 Conda 创建虚拟环境。

pip install -r requirements.txt

或使用以下命令创建 Conda 环境：

conda env create -f environment.yml

• 速记句：通过 pip 或 conda 快速安装 SELF-RAG。

3. 快速开始：如何使用 SELF-RAG 进行推理

你可以通过 HuggingFace Hub 下载 Self-RAG 模型，并使用 vllm（一个高效的推理框架）运行推理。下面是一个简单的推理示例：

from vllm import LLM, SamplingParams

model = LLM("selfrag/selfrag_llama2_7b", download_dir="/gscratch/h2lab/akari/model_cache", dtype="half")
sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=100, skip_special_tokens=False)

def format_prompt(input, paragraph=None):
  prompt = "### Instruction:\n{0}\n\n### Response:\n".format(input)
  if paragraph is not None:
    prompt += "[Retrieval]<paragraph>{0}</paragraph>".format(paragraph)
  return prompt

query_1 = "Leave odd one out: twitter, instagram, whatsapp."
query_2 = "Can you tell me the difference between llamas and alpacas?"
queries = [query_1, query_2]

preds = model.generate([format_prompt(query) for query in queries], sampling_params)
for pred in preds:
  print("Model prediction: {0}".format(pred.outputs[0].text))

• 速记句：使用 SELF-RAG 进行推理时可以根据查询动态检索或跳过检索。

4. 检索器设置

SELF-RAG 使用 Contriever 作为默认的检索组件。你可以下载预处理过的段落数据，并使用 Contriever-MSMARCO 模型进行检索。

下载段落数据：

cd retrieval_lm
wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz

下载生成的段落嵌入：

wget https://dl.fbaipublicfiles.com/contriever/embeddings/contriever-msmarco/wikipedia_embeddings.tar

• 速记句：SELF-RAG 的检索器基于 Contriever 模型，并可以使用预处理好的数据进行检索。

5. SELF-RAG 的训练流程

SELF-RAG 的训练过程分为两个主要部分：批判模型（Critic）的训练和生成模型（Generator）的训练。批判模型会学习如何使用反思标记评估生成内容，而生成模型则会根据这些批判信息生成更高质量的内容。

步骤 1：批判数据生成

我们使用 GPT-4 生成批判模型的训练数据，并为每种反思标记创建数据集。

步骤 2：批判模型训练

使用批判数据训练 Llama2-7B 模型：

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

步骤 3：生成数据生成

生成模型的数据是由批判模型和检索器共同生成的。

步骤 4：生成模型训练

使用 DeepSpeed 优化训练效率：

bash script_finetune_7b.sh

• 速记句：SELF-RAG 的训练包括批判模型和生成模型的协同训练。

6. SELF-RAG 的推理

SELF-RAG 在推理过程中可以根据生成的内容动态检索相关信息。在短文本生成任务中（如问答系统），通常只需检索一次。而在长文本生成任务中（如长篇问答），SELF-RAG 可以多次检索并生成支持性的答案。

• 打比方：短文本生成就像回答一个简单问题，而长文本生成则类似于写一篇需要多次查阅资料的论文。
• 速记句：SELF-RAG 在长文本生成中可以多次检索信息。

7. SELF-RAG 的基准测试

SELF-RAG 提供了多种基准测试，包括短文本生成（如 PubHealth、TriviaQA）和长文本生成（如 ASQA、FactScore）。你可以根据任务的不同选择合适的测试数据进行评估。

python run_short_form.py --model_name selfrag/selfrag_llama2_7b --input_file eval_data/popqa_longtail_w_gs.jsonl

• 速记句：SELF-RAG 提供了短文本和长文本生成的基准测试。

8. 常见问题

1. 如何训练新的预训练语言模型？
   你可以直接在 SELF-RAG 的训练脚本中修改 model_name_or_path 来适配新的模型。

2. Mistral-7B 版的 SELF-RAG 会发布吗？
   目前社区已经训练了基于 Mistral-7B 的 SELF-RAG 模型，你可以在 HuggingFace 上找到该模型。

• 速记句：通过修改训练脚本，SELF-RAG 可以适配新的预训练模型。

总结

SELF-RAG 是一种能够通过自我反思提升生成质量的新型框架。它结合了检索、多次生成和批判性反思，能够在不同的任务场景下提供高质量的回答。通过本教程，你应该能够理解 SELF-RAG 的核心思想，安装并运行它，设置检索器，训练批判和生成模型，并进行推理和基准测试。

---

参考文献

1. Akari Asai et al. (2024). Self-RAG: Learning to Retrieve, Generate and Critique through Self-reflection. ICLR 2024. Link
2. HuggingFace Hub. SELF-RAG 模型. Link