我们用于 Atlas 检索和预训练的预处理维基百科转储文件可按如下方式下载：

python preprocessing/download_corpus.py --corpus {语料下载键} --output_directory ${DATA_DIR} 

上述命令将下载一个语料并将其解压到 ${DATA_DIR}/{语料下载键}。

可用的语料如下表所示：

段落文件采用 jsonl 格式，每行序列化为一个 JSON 对象。默认情况下，每个段落应按以下格式组织：

{
    "id": "0", # 段落应具有唯一 ID
    "title": "兰花", # 应指定该段落来自的页面标题（如果没有合适的标题，可以为空字符串）
    "text": "兰花与其他植物很容易区分，因为它们具有一些非常明显的衍生特征或共源性状。其中包括：花的两侧对称性（合轴对称）、许多倒置的花朵、几乎总是高度特化的唇瓣、雄蕊和雌蕊融合在一起，以及极其微小的种子。", # 段落的主要文本
    "section": "描述" # 可选字段，表示段落所属的小节标题；如果非空，则默认会将此字段附加到标题后，形成 {title}: {section}
    ... # 您还可以添加其他字段以方便分析，但这些字段实际上不会被使用
}

如果您按照上述格式创建自己的段落文件，那么与 Atlas 配合使用应该会非常简单。

目前，我们无法开源论文中使用的 Common Crawl 索引。

我们正在开源基础、大、XL 和 XXL 尺寸的预训练 Atlas 模型。这些模型同时包含预训练的检索器和阅读器权重。

此外，我们还开源了性能最强的全微调 Natural Questions Atlas 模型，供希望进行最先进问答推理（或在其他问答任务上进行微调）的用户使用。

模型可按如下方式下载：

python preprocessing/download_model.py --model {模型下载键} --output_directory ${DATA_DIR} 

这将会把请求的模型下载到 ${DATA_DIR}/{模型下载键}，随后您可以通过将 ${DATA_DIR}/{模型下载键} 传递给 --model_path 参数，在脚本中使用这些模型。

下表详细列出了可用的模型：

如果未提供索引，Atlas 会自动构建一个索引。这种方式虽然方便，但耗时较长，尤其是在 GPU 工作节点较少或索引规模较大的情况下。

因此，我们提供了针对预训练 Atlas 检查点以及经 Natural Questions 数据微调的 Atlas 检查点的 wiki-dec2018 语料的预计算索引供下载。

这些索引可按如下方式下载：

python preprocessing/download_index.py --index {索引下载键} --output_directory ${DATA_DIR} 

上述脚本将下载请求的预训练索引，并将其保存到 ${DATA_DIR}/{索引下载键}。随后，您可以通过将这些索引传递给 --load_index_path 参数，在训练或评估中使用它们。关于索引的保存和加载的更多细节，请参阅 检索与索引详情。

可供下载的索引如下表所示：

这是最基础的任务，可能并不是您的最佳选择，尤其是当您的任务与已实现的其他任务非常相似时。 通过向 train.py 或 evaluate.py 传递 --task base 参数即可指定此任务。 该任务的训练/验证/测试数据应由 jsonl 文件组成，需以空格分隔的形式传入 train.py 或 evaluate.py，例如 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等。 输入文件应包含 query 字段，用于存储输入查询字符串，以及 target 字段，用于存储输出目标字符串，例如：

{"query": "输入到 Atlas 的内容", "target": "希望 Atlas 生成的内容"}

评估循环会计算评估损失，以及 Atlas 生成的输出与目标完全匹配的验证数据样本所占的比例。 如果您向脚本传递 --write_results 参数，Atlas 在验证数据上的预测结果将以如下格式写入保存检查点的目录：

{"query": "输入到 Atlas 的内容", "answers": ["希望 Atlas 生成的内容"], "generation": "Atlas 对该查询的预测结果", "passages": ["检索到的段落列表"]}

掩码语言建模任务实现了由 T5 提出的掩码语言建模预训练任务。这也是我们在论文中用来预训练主模型 Atlas 的任务。 通过向 train.py 传递 --task mlm 参数即可指定此任务。 该任务的训练/验证/测试数据应由 jsonl 文件组成，需以 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等形式传入 train.py。 这些文件应由具有以下格式的 JSON 对象组成：

{
  "text": "需要施加噪声并训练去噪的文本片段",
  "id": "该文本片段的唯一标识符"
  ... # 您还可以保留其他字段以便于分析，但这些字段实际上并不会被使用
}

其设计意图是，您可以将用于检索语料库的文件（通过 --passages 传递）直接用作训练数据。 该任务会应用 T5 的噪声函数处理 text 字段，从而自动创建输入和目标生成内容。 MLM 任务还会阻止 Atlas 检索正在尝试去噪的那篇文档。它通过过滤掉检索结果中与正在去噪的实例具有相同 id 字段的文档来实现这一点。如果去噪训练数据与 Atlas 正在检索的文档来自同一语料库，这一功能就显得尤为重要。 该任务具有以下特定于任务的参数：

  --mlm_noise_density MLM_NOISE_DENSITY
      输入文本中应被掩码跨度覆盖的比例（默认：0.15）
  --mlm_mean_noise_span_length MLM_MEAN_NOISE_SPAN_LENGTH
      MLM 掩码跨度的平均长度（默认：3）
  --min_words_per_lm_instance MIN_WORDS_PER_LM_INSTANCE
      如果实例中的词数少于此值，则会跳过该实例，不参与 MLM/LM/章节生成任务（默认：无）

如果您传递 --write_results 参数，Atlas 会将其填空预测结果写入文件。 在评估过程中，Atlas 会记录以下 MLM 任务的评估指标：

• eval_loss：生成的 MLM 填空跨度的评估损失
• accuracy：完美去噪的填空跨度所占比例
• f1：正确去噪的填空跨度的 token F1 分数
• rouge_1：生成的填空跨度相对于黄金参考掩码跨度的 ROUGE-1 分数
• rouge_2：生成的填空跨度相对于黄金参考掩码跨度的 ROUGE-2 分数
• rouge_L：生成的填空跨度相对于黄金参考掩码跨度的 ROUGE-L 分数

通过向 train.py 传递 --task lm 参数，Atlas 可以被训练为执行从左到右的语言建模任务。 该任务的训练/验证/测试数据应由 jsonl 文件组成，需以 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等形式传入 train.py。 这些文件应由具有以下格式的 JSON 对象组成：

{
  "text": "需要训练 Atlas 生成的文本片段",
  "id": "该文本片段的唯一标识符"
  ... # 您还可以保留其他字段以便于分析，但这些字段实际上并不会被使用
}

其设计意图是，您可以将用于检索语料库的文件（通过 --passages 传递）直接用作训练数据。 该任务会自动对 text 字段进行预处理，将其随机分为两部分：左侧作为条件上下文，右侧则作为 Atlas 模型将被训练生成的后续内容。 LM 任务还会阻止 Atlas 检索正在尝试生成的同一文档。它通过过滤掉检索结果中与正在生成的实例具有相同 id 字段的文档来实现这一点。如果去噪训练数据与 Atlas 正在检索的文档来自同一语料库，这一功能就显得尤为重要。

该任务具有以下特定于任务的参数：

  --min_words_per_lm_instance MIN_WORDS_PER_LM_INSTANCE
      具有少于 min_words_per_lm_instance 个词的实例将被跳过，不参与 MLM/LM/Section 生成（默认：无）
  --min_lm_context_ratio MIN_LM_CONTEXT_RATIO
      将文本分割为两个部分用于语言建模。左侧部分作为条件上下文，右侧部分用于生成。左侧部分必须大于右侧部分的 min_lm_context_ratio（默认：0.5）
  --max_lm_context_ratio MAX_LM_CONTEXT_RATIO
      将文本分割为两个部分用于语言建模。左侧部分作为条件上下文，右侧部分用于生成。左侧部分必须小于右侧部分的 max_lm_context_ratio（默认：0.5）

如果您传递 --write_results，Atlas 会将其语言模型预测结果写入文件。

在评估过程中，Atlas 会记录以下语言模型评估指标：

• eval_loss：参考数据中续写的评估读取损失
• accuracy：完全正确预测的续写比例
• f1：生成续写的 token F1 分数，表示正确生成的比例
• rouge_1：生成续写相对于黄金参考续写的 ROUGE-1 分数
• rouge_2：生成续写相对于黄金参考续写的 ROUGE-2 分数
• rouge_L：生成续写相对于黄金参考续写的 ROUGE-L 分数

通过向 train.py 传递 --task section，可以训练 Atlas 根据维基百科条目的标题和章节标题生成相应段落的文本。

此任务的训练/验证/测试数据应由 JSONL 文件组成，其格式应与维基百科转储中的 text-list-100-sec.jsonl 文件一致。这些文件可以通过遵循可下载的数据和模型中的说明获取，例如训练文件：enwiki-dec2018/text-list-100-sec.jsonl。这些文件应由每行一个 JSON 对象组成，格式如下：

{
  "id": "3793043", 
  "title": "百慕大三角",
  "section": "指南针偏差",
  "text": " 指南针问题是许多百慕大三角事件中经常提到的现象之一。尽管有人推测该地区可能存在异常的局部磁异常，但至今尚未发现此类异常。实际上，指南针会因与地磁极的关系而产生自然的偏差，这是航海家们几个世纪以来就已知的事实。"
}

该任务会自动将输入查询格式化为“{标题}, {章节}”——例如，在此示例中，输入到 Atlas 的内容将是 百慕大三角, 指南针偏差。输出将是示例中的 text 字段。 section 任务会防止 Atlas 生成与其检索到的同一段落相同的内容。它通过过滤掉检索结果中与正在生成的实例具有相同 id 字段的段落来实现这一点。

该任务具有以下特定于任务的参数：

  --min_words_per_lm_instance MIN_WORDS_PER_LM_INSTANCE
      具有少于 min_words_per_lm_instance 个词的实例将被跳过，不参与 MLM/LM/Section 生成（默认：无）

如果您传递 --write_results，Atlas 会将其生成的维基百科章节文本预测结果写入文件。

在评估过程中，Atlas 会记录以下 section 任务的评估指标：

• eval_loss：参考数据中续写的评估读取损失
• accuracy：完全正确预测的续写比例
• f1：生成续写的 token F1 分数，表示正确生成的比例
• rouge_1：生成续写相对于黄金参考续写的 ROUGE-1 分数
• rouge_2：生成续写相对于黄金参考续写的 ROUGE-2 分数
• rouge_L：生成续写相对于黄金参考续写的 ROUGE-L 分数

通过向 train.py 或 evaluate.py 传递 --task qa，可以训练 Atlas 回答开放域问答问题。在快速入门与代码库概览部分有一个 QA 的示例。

我们在论文中使用此任务处理 NaturalQuestions、TriviaQA 和 TempLama 数据集。

此任务的训练/验证/测试数据应由 JSONL 文件组成，这些文件应以 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等形式传递给 train.py。 每个文件应包含一行 JSON 实例，格式如下：

{
  "question": "百慕大三角在哪里",
  "answers": ["北大西洋西部"],
   ... # 您可以保留其他字段以便于分析，但这些字段不会实际用于训练
}

问题将根据任务特定参数 --qa_prompt_format 进行格式化，默认值为 question: {question} answer: <extra_id_0>。 例如，上述问题将自动格式化为输入到 Atlas 的查询：question: 百慕大三角在哪里 answer: <extra_id_0>。 监督目标来自 target 字段。如果该字段不存在，则监督目标将从 answers 字段中的可用答案中随机选择，并格式化为 <extra_id_0> {answer}。

如果您传递 --write_results，Atlas 会将其预测的答案写入文件。

在评估过程中，Atlas 会记录以下开放域问答的评估指标：

• eval_loss：评估答案的评估读取损失
• exact_match：生成答案的开放域问答精确匹配分数
• f1：生成答案的开放域问答 F1 分数

Natural Questions 和 TriviaQA

您可以通过运行以下命令下载 NaturalQuestions 和 TriviaQA 数据：

python preprocessing/prepare_qa.py --output_directory ${DATA_DIR} 

这将下载 train.jsonl、train.64-shot.jsonl（我们使用的少量样本训练数据集）、dev.jsonl 和 test.jsonl，并将其保存到 ${DATA_DIR}/data/nq_data 和 ${DATA_DIR}/data/triviaqa_data 中。

有关使用 NQ 维基百科索引进行少量样本和标准微调及评估的示例脚本可在 examples/nq 中找到。只需替换训练/验证/测试文件，即可将该脚本用于 TriviaQA。

TempLama

我们基于 TempLAMA 数据集定义了一个完形填空式问答任务，用于评估索引的忠实性和时间迁移能力。

您可以通过运行以下脚本下载 TempLAMA 数据并创建和格式化我们衍生的数据集：

python preprocessing/prepare_templama.py --output_directory ${DATA_DIR} 

这将创建文件 temp_lama.train.2017.jsonl、temp_lama.valid.2017.jsonl、temp_lama.test.2017.jsonl、temp_lama.train.2020.jsonl、temp_lama.valid.2020.jsonl、temp_lama.test.2020.jsonl，位于 ${DATA_DIR}/data/templama_data/ 目录下。 这些文件将包含完形填空题，答案会根据年份有所不同。

运行 TempLama 训练和评估的示例脚本可以在 examples/templama 中找到。（注意使用了 qa_prompt_format {question}，它会关闭 TriviaQA 和 NQ 所使用的自动 QA 提示格式化功能）

Atlas 可以通过在 train.py 或 evaluate.py 中添加 --task multiple_choice 参数来训练回答多项选择题。我们在 MMLU 的实验中使用了这一任务。 该任务的训练/验证/测试数据应由 jsonl 文件组成，这些文件应作为 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等参数传递给 train.py。 每个文件每行应包含一个 JSON 实例，格式如下：

{
  "question": "以下哪个是包含垂体的体腔？", 
  "options": {
    "A": "腹腔",
    "B": "颅腔",
    "C": "胸膜腔", 
    "D": "脊髓腔"
    ... # 你可以有更多的（或更少的）选项，只要它们的键是按字母顺序连续的大写字母，从 A 开始即可
  }, 
  "answer": "B",
  ... # 你还可以保留其他字段以便于分析，但这些字段并不会被实际使用
}

这些数据会被自动格式化为 Atlas 的输入查询，形式为 question: {question} answers: (A) {options['A']} (B) {options['B']} (C) {options['C']} (D) {options['D']} Answer: <extra_id_0>，目标生成格式为 <extra_id_0> {answer letter}。 上述示例会被格式化为：question: {以下哪个是包含垂体的体腔？ answers: (A) 腹腔 (B) 颅腔 (C) 胸膜腔 (D) 脊髓腔 Answer: <extra_id_0>，目标生成为 {extra_id_0} B。

多项选择问答任务有以下特定参数：

  --multiple_choice_num_options
      多项选择问答中选项的数量（MMLU 是 4 个）（默认值：4）
  --multiple_choice_train_permutations {single,cyclic,all}
      在进行多项选择训练时（例如 MMLU），是否启用答案顺序的排列组合。这有助于消除模型对任意答案顺序的偏好，从而提升效果。建议使用 'all' 模式。single：不进行排列；cyclic：循环排列；all：所有可能的答案排列组合。（默认值：single）
  --multiple_choice_eval_permutations {single,cyclic,all}
      在进行多项选择评估时（例如 MMLU），是否启用答案顺序的排列组合。这同样可以减少模型对答案顺序的偏好，从而提升效果。使用 'all' 模式效果最佳，但速度较慢。'cyclic' 是一个不错的折中方案。single：不进行排列；cyclic：循环排列；all：所有可能的答案排列组合。（默认值：single）

排列选项会自动复制输入数据，并对答案顺序进行排列（例如，“A”变为“颅腔”，“B”变为“胸膜腔”等）。 当监督数据量非常少（或零样本情况下），这种方法可以显著提升效果。 如果你使用了 --multiple_choice_eval_permutations 参数中的 cyclic 或 all 选项，代码会自动对不同排列组合的结果进行汇总，为你提供最终的评估结果。 关于排列去偏化的更多细节，请参阅论文 Atlas: Few-shot Learning with Retrieval Augmented Language Models 的附录。

如果你指定了 --write_results，Atlas 会将其预测的答案写入文件，格式如下：

{
  "question": "应用提示模板后的输入内容",
  "generation": "在对排列组合结果进行汇总后，概率最高的答案字母",
  "choice_probs": "每个答案选项的概率（归一化到总选项数）",
  "all_probs": "未汇总前的所有答案排列组合的概率",
  "permutations": ["针对每种答案排序组合的预测对象列表"]
}

MMLU

专门用于运行 MMLU 实验的 ReadMe 文档可在 这里 查看。 我们提供了一个用于下载和预处理 MMLU 数据的工具，并且在 examples/mmlu 中提供了针对我们所探索的各种实验设置的示例脚本。 这些内容在 MMLU 的专用 ReadMe 文档中都有详细说明。

Atlas 可以通过使用 --task fever 参数在 train.py 或 evaluate.py 中训练，使其能够根据语料库将文本陈述分类为“支持”、“反驳”或“信息不足”，例如用于 FEVER 任务。 你可以通过运行以下脚本来下载 FEVER 数据：

python preprocessing/prepare_fever.py --output_directory ${DATA_DIR} 

该任务的训练/验证/测试数据应由 jsonl 文件组成，这些文件应作为 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等参数传递给 train.py。 每个文件每行应包含一个 JSON 实例，格式如下：

{
  "claim": "需要评估的陈述", 
  "label": "要么是 'SUPPORTS'、'REFUTES'，要么是 'NOT ENOUGH INFO'",
   ... # 你可以保留其他字段以便于分析，但这些字段并不会被实际使用
}

Atlas 会自动处理这些实例，并将其格式化为输入 question: {claim} answer: <extra_id_0> 和输出 <extra_id_0> {true, false 或 maybe}。 如果你指定了 --write_results，Atlas 会将其预测的标签写入文件。 在评估过程中，Atlas 还会记录以下开放域问答的评估指标：

• accuracy: 模型正确分类的陈述数量。

Atlas 可以通过在 train.py 或 evaluate.py 中使用 --task kilt 参数来训练执行 KILT 任务。 KILT 数据可以从 这里 获取。

该任务的训练/验证/测试数据应由 JSONL 文件组成，这些文件应作为 --train_data train_file_1.jsonl train_file_2.jsonl 和 --eval_data eval_file_1.jsonl eval_file_2.jsonl 等参数传递给 train.py。 每个文件应每行包含一个 JSON 实例，格式如下（即代码库可以直接接受 KILT 格式）：

{'id': # 原始数据点的 ID，如果有的话；否则为唯一 ID
 'input': # 问题 / 主张 / 句子 / 等等
 'output': [ # 每个元素可能包含答案、证据来源或两者
    {
    'answer': # 文本形式的答案
    'provenance': [
        # 针对答案的 KILT 数据集中的证据集合
        {
            'wikipedia_id':  # *必须* 
            'title': 
            'section': 
            'start_paragraph_id': 
            'start_character': 
            'end_paragraph_id':
            'end_character': 
            'bleu_score': # 相对于原始证据的 BLEU 分数
            'meta': # 数据集/任务特定的元数据
        }
        ] 
      }
    ]
 'meta': # 数据集/任务特定的元数据
 }

Atlas 会自动根据 input 字段将这些实例处理为 Atlas 查询输入，并根据 answer 字段生成目标输出。

如果您传递 --write_results 参数，Atlas 会将其预测标签写入文件。

在评估过程中，Atlas 将记录以下开放域问答任务的评估指标：

• accuracy: 生成内容与参考答案完全匹配的频率
• exact_match: 在应用开放域问答标准化后，生成内容与参考答案完全匹配的频率
• f1: 生成内容与参考答案之间的 token 级别 F1 分数重叠

构建或刷新索引涉及为检索索引中的每条文档计算嵌入，然后将其保存在内存中，以便在后续检索时能够快速计算最大内积或最近邻。（假设我们有 W 个 GPU 工作节点，检索索引中总共有 N 条文档）

请注意，每个 GPU 工作节点负责 N/W 条文档的分片。文档嵌入的计算非常简单，过程如下：

• 暂停正在进行的模型训练
• 每个工作节点并行计算其分片中文档的嵌入，按批次迭代处理，并将结果保存在一个大型 PyTorch 张量中（FAISS 支持也可用，但稍后讨论）。
• 当所有工作节点都完成其分片的嵌入计算后，我们可以继续进行模型训练。

N 可能非常大（1000万到1亿），因此嵌入所有文档的速度会很慢。然而，由于可以在各个工作节点之间并行化，如果 GPU 数量足够，整个过程可以相对较快。尽管如此，对于大型索引，这个过程仍然可能耗时较长。我们提供了索引保存功能，可以将索引保存到磁盘，以避免频繁重建索引。此外，随着检索器的不断训练，缓存的嵌入可能会过时或“陈旧”，需要重新计算，这会带来额外的成本。有关减少或避免频繁刷新索引的方法，请参阅处理陈旧索引的策略。

Atlas 在训练循环中执行检索：即在前向传播过程中调用检索功能。下面我们将简要描述 Atlas 前向传播中的步骤，包括如何完成检索： （假设我们有 W 个 GPU 工作节点，检索索引中总共有 N 条文档，为简化起见，假设每 GPU 的训练批大小为 1，且任务为问答任务）。

• 每个工作节点会收到一个问题，并将其嵌入为查询向量。
• 然后执行全归约操作，使得每个工作节点都拥有全部 W 个查询向量的副本（即当前小批量中所有问题的查询向量）。
• 接着，每个工作节点在其负责的文档分片上对小批量中的所有查询向量执行 GPU 上的最大内积搜索。
• 每个工作节点会将其分片中每个查询的前 K 个结果通过归约操作发送回嵌入该查询的 GPU。
• 最终，每个工作节点会获得其查询对应的前 W × K 个结果，从中选出真正的前 K 个结果。
• 检索过程至此完成，随后可以继续进行标准的分布式数据并行前向传播（即运行模型前向传播、计算梯度、在各工作节点之间聚合梯度，并更新模型参数）。

{
  "question": "百慕大三角在哪里",
  "answers": ["北大西洋西部海域"],
  "passages": [
    {
      "text": "第一段落的内容",
      "title": "第一段落的标题",
      "id": "第一段落的唯一标识符"
      ... # 其他字段也可存在，但不会被使用
    },
    {
      "text": "第二段落的内容",
      "title": "第二段落的标题",
      "id": "第二段落的唯一标识符"
    },
    ... # 如有需要，可添加更多段落
  ]
}

用法: train.py/evaluate.py [-h] [--name NAME] [--checkpoint_dir CHECKPOINT_DIR] [--model_path MODEL_PATH] [--per_gpu_batch_size PER_GPU_BATCH_SIZE] [--per_gpu_embedder_batch_size PER_GPU_EMBEDDER_BATCH_SIZE] [--local_rank LOCAL_RANK]
                [--main_port MAIN_PORT] [--seed SEED] [--log_freq LOG_FREQ] [--eval_freq EVAL_FREQ] [--save_freq SAVE_FREQ] [--train_data TRAIN_DATA [TRAIN_DATA ...]] [--eval_data EVAL_DATA [EVAL_DATA ...]] [--write_results]
                [--dont_write_passages] [--load_index_path LOAD_INDEX_PATH] [--save_index_path SAVE_INDEX_PATH] [--save_index_n_shards SAVE_INDEX_N_SHARDS] [--index_mode {flat,faiss}] [--faiss_index_type {ivfflat,flat,ivfsq,sq,pq}]
                [--faiss_code_size FAISS_CODE_SIZE] --reader_model_type
                {t5-small,t5-base,t5-large,t5-3b,t5-11b,google/t5-v1_1-base,google/t5-v1_1-large,google/t5-v1_1-xl,google/t5-v1_1-xxl,google/t5-base-lm-adapt,google/t5-large-lm-adapt,google/t5-xl-lm-adapt,google/t5-xxl-lm-adapt}
                [--text_maxlength TEXT_MAXLENGTH] [--target_maxlength TARGET_MAXLENGTH] [--n_context N_CONTEXT] [--passages PASSAGES [PASSAGES ...]] [--max_passages MAX_PASSAGES] [--retriever_model_path RETRIEVER_MODEL_PATH]
                [--retrieve_only] [--train_retriever] [--use_file_passages] [--retriever_n_context RETRIEVER_N_CONTEXT] [--gold_score_mode {evalnormsum,loop,ppmean,emdr,pdist,adist}] [--closed_book]
                [--temperature_score TEMPERATURE_SCORE] [--temperature_gold TEMPERATURE_GOLD] [--compute_crossattention_stats] [--filtering_overretrieve_ratio FILTERING_OVERRETRIEVE_RATIO]
                [--freeze_retriever_steps FREEZE_RETRIEVER_STEPS] [--query_side_retriever_training] [--retrieve_with_rerank] [--n_to_rerank_with_retrieve_with_rerank N_TO_RERANK_WITH_RETRIEVE_WITH_RERANK]
                [--decoder_format DECODER_FORMAT] [--decoder_prompt_format DECODER_PROMPT_FORMAT] [--encoder_format ENCODER_FORMAT] [--retriever_format RETRIEVER_FORMAT] [--generation_max_length GENERATION_MAX_LENGTH]
                [--generation_min_length GENERATION_MIN_LENGTH] [--generation_length_penalty GENERATION_LENGTH_PENALTY] [--generation_num_beams GENERATION_NUM_BEAMS] [--task {base,mlm,lm,multiple_choice,kilt,section,fever,qa}]
                [--mlm_noise_density MLM_NOISE_DENSITY] [--mlm_mean_noise_span_length MLM_MEAN_NOISE_SPAN_LENGTH] [--min_words_per_lm_instance MIN_WORDS_PER_LM_INSTANCE] [--min_lm_context_ratio MIN_LM_CONTEXT_RATIO]
                [--max_lm_context_ratio MAX_LM_CONTEXT_RATIO] [--qa_prompt_format QA_PROMPT_FORMAT] [--multiple_choice_num_options MULTIPLE_CHOICE_NUM_OPTIONS] [--multiple_choice_train_permutations {single,cyclic,all}]
                [--multiple_choice_eval_permutations {single,cyclic,all}] [--warmup_steps WARMUP_STEPS] [--total_steps TOTAL_STEPS] [--scheduler_steps SCHEDULER_STEPS] [--accumulation_steps ACCUMULATION_STEPS] [--dropout DROPOUT]
                [--lr LR] [--lr_retriever LR_RETRIEVER] [--clip CLIP] [--scheduler {linear,cosine,fixed}] [--weight_decay WEIGHT_DECAY] [--save_optimizer] [--epsilon EPSILON] [--alpha ALPHA] [--beta2 BETA2]
                [--refresh_index REFRESH_INDEX] [--shuffle] [--precision {fp16,fp32,bf16}] [--shard_optim] [--shard_grads] [--use_gradient_checkpoint_reader] [--use_gradient_checkpoint_retriever]

可选参数：
  -h, --help            显示此帮助信息并退出
  --name NAME           实验名称，也用作目录名（默认值：experiment_name）
  --checkpoint_dir CHECKPOINT_DIR
                        模型保存在此目录下（默认值：./checkpoint/）
  --model_path MODEL_PATH
                        用于初始化的预训练模型路径（传入 'none' 表示从 T5 和 Contriever 初始化）（默认值：无）
  --per_gpu_batch_size PER_GPU_BATCH_SIZE
                        每个 GPU/CPU 的训练批次大小。（默认值：1）
  --per_gpu_embedder_batch_size PER_GPU_EMBEDDER_BATCH_SIZE
                        Embedder 每个 GPU 的批次大小。（默认值：512）
  --local_rank LOCAL_RANK
                        用于分布式训练：本地进程编号（默认值：-1）
  --main_port MAIN_PORT
                        主端口（用于多节点任务）（默认值：-1）
  --seed SEED           初始化时使用的随机种子（默认值：0）
  --log_freq LOG_FREQ   训练过程中每 <log_freq> 步记录一次训练统计信息（默认值：100）
  --eval_freq EVAL_FREQ
                        训练过程中每 <eval_freq> 步评估一次模型（默认值：500）
  --save_freq SAVE_FREQ
                        训练过程中每 <save_freq> 步保存一次模型（默认值：5000）
  --train_data TRAIN_DATA [TRAIN_DATA ...]
                        以空格分隔的 JSONL 格式训练数据集路径列表（默认值：空列表）
  --eval_data EVAL_DATA [EVAL_DATA ...]
                        以空格分隔的 JSONL 格式评估数据集路径列表（默认值：空列表）
  --write_results       将评估结果保存到文件中（默认值：False）
  --dont_write_passages
                        如果要写结果，段落可能会占用大量空间，使用此标志可以不将段落写入导出的结果中（默认值：False）
  --load_index_path LOAD_INDEX_PATH
                        用于加载索引、段落嵌入和段落的路径（默认值：None）
  --save_index_path SAVE_INDEX_PATH
                        用于保存索引和/或嵌入的路径（默认值：None）
  --save_index_n_shards SAVE_INDEX_N_SHARDS
                        将索引保存为多少个分片文件。必须是工作进程数的整数倍。（默认值：128）
  --index_mode {flat,faiss}
                        使用扁平的 PyTorch 索引或 Faiss 索引来检索 k 个最近邻（默认值：flat）
  --faiss_index_type {ivfflat,flat,ivfsq,sq,pq}
                        IVFFlat、IndexFlatIP、IVFScalarQuantizer、ScalarQuantizer 或带有 faiss-gpu 的 IndexPQ（默认值：flat）
  --faiss_code_size FAISS_CODE_SIZE
                        PQ 量化参数（默认值：None）
  --reader_model_type {t5-small,t5-base,t5-large,t5-3b,t5-11b,google/t5-v1_1-base,google/t5-v1_1-large,google/t5-v1_1-xl,google/t5-v1_1-xxl,google/t5-base-lm-adapt,google/t5-large-lm-adapt,google/t5-xl-lm-adapt,google/t5-xxl-lm-adapt}
                        阅读器 FID 模型的 T5 架构，例如 google/t5-xl-lm-adapt（默认值：None）
  --text_maxlength TEXT_MAXLENGTH
                        输入文本片段（问题+段落拼接后）的最大 token 数。超过此长度的输入将被截断。（默认值：200）
  --target_maxlength TARGET_MAXLENGTH
                        训练模型时目标输出的最大 token 长度。超过此长度的目标将被截断。如果设置为 -1，则不进行截断（默认值：None）
  --n_context N_CONTEXT
                        传递给阅读器的 top k 段落数量（默认值：1）
  --passages PASSAGES [PASSAGES ...]
                        包含要索引和检索的段落的 JSONL 文件路径列表。如果使用 --load_index_path 加载已保存的索引，则此参数无效（默认值：None）
  --max_passages MAX_PASSAGES
                        要索引的最大段落数量。设置为 -1 表示读取段落文件中的所有段落（默认值：-1）
  --retriever_model_path RETRIEVER_MODEL_PATH
                        用于初始化的 Contriever 模型路径（如果传入 --model_path 参数，则覆盖此值）（默认值：facebook/contriever）
  --retrieve_only       传入此参数以防止加载阅读器，仅运行检索评估（默认值：False）
  --train_retriever     传入此参数以同时训练检索器和阅读器（默认值：False）
  --use_file_passages   使用训练或评估 JSONL 文件中 "passages" 字段中的段落，而不是通过检索获取段落（默认值：False）
  --retriever_n_context RETRIEVER_N_CONTEXT
                        用于训练检索器的 top k 段落数量（默认值：5）
  --gold_score_mode {evalnormsum,loop,ppmean,emdr,pdist,adist}
                        训练检索器的方法。`pdist` 是论文中对 `ppmean` 的称呼。`adist` 是论文中对 `evalnormsum` 的称呼（默认值：ppmean）
  --closed_book         不使用检索功能——退化为 T5 模型。如果设置了 n_context、n_context_retriever 和 encoder_format，则此选项会覆盖它们（默认值：False）
  --temperature_score TEMPERATURE_SCORE
                        检索器的 softmax 温度（默认值：0.01）
  --temperature_gold TEMPERATURE_GOLD
                        检索器蒸馏目标分布的 softmax 温度（默认值：0.01）
  --compute_crossattention_stats
  --filtering_overretrieve_ratio FILTERING_OVERRETRIEVE_RATIO
                        如果进行过滤，先按此比例超额检索 topK，然后再过滤掉不需要的结果。很有用；只有在处理不需过滤检索结果的任务时才设为 1（默认值：2）
  --freeze_retriever_steps FREEZE_RETRIEVER_STEPS
                        冻结检索器 n 步（默认值：-1）
  --query_side_retriever_training
                        传入此参数以启用查询端的检索器微调（解绑 Contriever 编码器中段落和查询编码器的参数，并冻结段落编码器。有助于避免索引刷新。（默认值：False）
  --retrieve_with_rerank
                        传入此参数以启用使用全新段落编码器进行重新排序的检索（默认值：False）
  --n_to_rerank_with_retrieve_with_rerank N_TO_RERANK_WITH_RETRIEVE_WITH_RERANK
                        当传入 --retrieve_with_rerank 时，需要重新排序的段落数量。数值越高越慢但越准确。推荐 64-128（默认值：128）
  --decoder_format DECODER_FORMAT
                        解码器的格式。模型将按照该格式进行训练，评估也将采用与 decoder_prompt_format 选项相反的格式（默认值：无）
  --decoder_prompt_format DECODER_PROMPT_FORMAT
                        解码器提示的格式，例如“{query} 的答案是什么？”（默认值：无）
  --encoder_format ENCODER_FORMAT
                        阅读器编码器预处理的格式字符串（默认值：{query} 标题：{title} 上下文：{text}）
  --retriever_format RETRIEVER_FORMAT
                        检索器编码器预处理的格式字符串（默认值：{title} {text}）
  --generation_max_length GENERATION_MAX_LENGTH
  --generation_min_length GENERATION_MIN_LENGTH
  --generation_length_penalty GENERATION_LENGTH_PENALTY
  --generation_num_beams GENERATION_NUM_BEAMS
  --task {base,mlm,lm,multiple_choice,kilt,section,fever,qa}
                        模型执行的任务。用于设置预处理、检索过滤、评估等。（默认值：无）
  --mlm_noise_density MLM_NOISE_DENSITY
                        输入文本中应被掩码跨度遮盖的比例（默认值：0.15）
  --mlm_mean_noise_span_length MLM_MEAN_NOISE_SPAN_LENGTH
                        MLM 掩码跨度的平均长度（默认值：3）
  --min_words_per_lm_instance MIN_WORDS_PER_LM_INSTANCE
                        对于 MLM/LM/Section Generation，如果实例中的单词数少于 min_words_per_lm_instance，则会跳过该实例（默认值：无）
  --min_lm_context_ratio MIN_LM_CONTEXT_RATIO
                        将文本分为两个部分进行语言建模。“左半部分作为条件上下文，右半部分用于生成。”左半部分必须大于右半部分的 min_lm_context_ratio。
                        （默认值：0.5）
  --max_lm_context_ratio MAX_LM_CONTEXT_RATIO
                        将文本分为两个部分进行语言建模。“左半部分作为条件上下文，右半部分用于生成。”左半部分必须小于右半部分的 max_lm_context_ratio。
                        （默认值：0.5）
  --qa_prompt_format QA_PROMPT_FORMAT
                        当使用 --task qa 时，如何将问题格式化为输入提示（默认值：问题：{question} 答案：<extra_id_0>）
  --multiple_choice_num_options MULTIPLE_CHOICE_NUM_OPTIONS
                        多项选择问答中有多少个选项（例如 MMLU 是 4 个）（默认值：4）
  --multiple_choice_train_permutations {single,cyclic,all}
                        在多项选择任务（如 MMLU）训练时，是否使用答案顺序的排列组合。这可以改善结果，消除模型对任意答案顺序的偏好。建议使用 all。
                        single：无排列；cyclic：循环排列；all：所有可能的答案排列组合。（默认值：single）
  --multiple_choice_eval_permutations {single,cyclic,all}
                        在多项选择任务（如 MMLU）评估时，是否使用答案顺序的排列组合。这可以改善结果，消除模型对任意答案顺序的偏好。最好使用 all，但速度很慢。
                        cyclic 是一个不错的折中方案。single：无排列；cyclic：循环排列；all：所有可能的答案排列组合。（默认值：single）
  --warmup_steps WARMUP_STEPS
                        学习率预热步数（默认值：1000）
  --total_steps TOTAL_STEPS
                        总训练步数（默认值：1000）
  --scheduler_steps SCHEDULER_STEPS
                        调度器的总步数。如果未指定，则 scheduler_total_step = total_step。（默认值：无）
  --accumulation_steps ACCUMULATION_STEPS
                        梯度累积（默认值：1）
  --dropout DROPOUT     掉落率（默认值：0.1）
  --lr LR               学习率（默认值：0.0001）
  --lr_retriever LR_RETRIEVER
                        检索器的学习率（默认值：1e-05）
  --clip CLIP           梯度裁剪（默认值：1.0）
  --scheduler {linear,cosine,fixed}
                        使用的学习率调度策略（默认值：cosine）
  --weight_decay WEIGHT_DECAY
                        训练中应用的权重衰减量（默认值：0.1）
  --save_optimizer      传入此标志以在保存的检查点中保存优化器状态（默认值：False）
  --epsilon EPSILON     AdamW 的 epsilon 值（默认值：1e-06）
  --alpha ALPHA         AdamW 的 alpha 值（默认值：1.0）
  --beta2 BETA2         AdamW 的 beta2 值（默认值：0.999）
  --refresh_index REFRESH_INDEX
                        索引刷新计划。格式：起始步-结束步：刷新频率，起始步-结束步：刷新频率。例如，--refresh_index 0-100:10,100-1000000:500 将在前 100 步中每 10 步刷新一次索引，然后从第 100 步到第 100 万步每 500 步刷新一次。对于固定计划，可以直接传入一个数字，例如 --refresh_index 100 将每 100 步刷新一次索引。传入 -1 表示永不刷新。（默认值：0-1000000:1000000）
  --shuffle             在训练时打乱数据（默认值：False）
  --precision {fp16,fp32,bf16}
                        数值精度——如果可用，建议使用 bf16；fp16 在训练中可能不稳定（默认值：fp32）
  --shard_optim         训练时的内存优化：使用分片数据并行将优化器状态分散到可用的 GPU 上，推荐用于大型模型（默认值：False）
  --shard_grads         训练时的内存优化：使用分片数据并行将梯度分散到可用的 GPU 上，推荐用于大型模型（默认值：False）
  --use_gradient_checkpoint_reader
                        在阅读器中使用梯度检查点技术（默认值：False）
  --use_gradient_checkpoint_retriever
                        在检索器中使用梯度检查点技术（默认值：False）