FedRAG 数据准备指南：基于 FAISS 的医疗语料下载与索引构建（Flower 联邦 RAG 实践）

FedRAG 数据准备指南：基于 FAISS 的医疗语料下载与索引构建（Flower 联邦 RAG 实践）

【免费下载链接】flowerFlower: A Friendly Federated AI Framework项目地址: https://gitcode.com/GitHub_Trending/flo/flower

本文围绕 Flower 联邦框架中 FedRAG 示例的examples/fedrag/data目录，系统讲解四大医疗/通用语料的下载、分块（chunking）与 FAISS 检索索引的构建全流程。读者将掌握prepare.sh/prepare.py的使用方法、各语料的规模与磁盘占用、retriever.yaml的嵌入模型配置，以及IndexIVFFlat + METRIC_L2索引的底层工作原理，从而能自主复现 FedRAG 的数据准备阶段。

FedRAG 中数据准备目录的角色

在 FedRAG 示例 的整体流水线中，客户端依赖本地的文档库完成 top-k 检索。而examples/fedrag/data目录正是这个“本地文档库”的构建入口：它负责下载语料、执行必要的分块操作，并为更高效的文档检索生成 FAISS 索引。

该目录的核心文件包括：

• prepare.sh：一键下载与建索引的 Bash 入口脚本；
• prepare.py：执行下载、建索引、抽样查询及 QA 基准数据下载的 Python 脚本；
• download.py：负责各语料的下载与分块；
• statpearls.py：StatPearls 语料的 XML 解析与分块脚本；
• README.md：本数据准备指南文档。

下载后的语料统一存放于examples/fedrag/data/corpus/<dataset_name>/目录，每个语料目录下包含原始数据与chunk/分块目录，索引构建产物faiss.index与all_doc_ids.npy也生成于此。

支持的语料与规模概览

当前版本支持四个语料，全部衍生自 MedRAG 工具包。下表完整列出了各语料的体量（数据来自 data/README.md）：

语料	大小	文档数	分块数（Snippets）	领域
PubMed	~70 GB	2390 万	2390 万	生物医学
StatPearls	~2 GB	9.3k	30.12 万	临床
Textbooks	~209 MB	18	12.58 万	医学
Wikipedia	~44 GB	650 万	2990 万	通用

[!NOTE] 对每个语料而言，其对应的索引所需磁盘空间可能与被索引文档本身相当。即索引的磁盘开销与语料大小基本是 1:1 的关系，在规划存储时需要预留双倍空间。

默认情况下，示例使用体量较小的Textbooks与StatPearls两个语料来演示 FedRAG 流水线——因为PubMed与Wikipedia的文档数量极为庞大，下载与建索引会消耗大量时间。

一键执行：prepare.sh 与 prepare.py

默认行为

在examples/fedrag/data目录下直接运行：

./prepare.sh

默认情况下脚本会下载StatPearls与Textbooks两个语料，并使用每个语料前 100 个分块（即前 100 个 chunk 文件）来创建索引。这样做的目的是快速生成索引、引导示例跑通，而非构建完整索引。

全量下载与全量建索引

若需要下载全部四个语料，并使用每个语料的全部文件建索引，运行：

./prepare.sh --datasets "pubmed" "statpearls" "textbooks" "wikipedia" --index_num_chunks 0

其中：

• --datasets：空格分隔的语料名列表，合法取值为pubmed、statpearls、textbooks、wikipedia；
• --index_num_chunks：建索引时考虑的分块文件数量，设置为0时表示使用全部文件。

从 prepare.sh 的源码可以看到其参数解析逻辑：--datasets会清空默认值并逐个收集后续参数直至遇到下一个--开头参数；--index_num_chunks则直接读取下一个参数值；脚本最终将参数透传给python ./prepare.py执行。脚本开头设置了set -e，任一环节失败即终止，避免在下载或建索引出错时继续执行后续步骤。

prepare.py 的执行流程

prepare.py 是实际执行者，其流程为：

1. 解析--datasets（默认textbooks）与--index_num_chunks（默认0）；
2. 对--datasets去重后，逐个调用DownloadCorpora.download(corpus=...)下载语料；
3. 调用retriever.build_faiss_index(dataset_name=..., batch_size=32, num_chunks=...)构建 FAISS 索引；
4. 每个语料索引构建完成后，用样例查询"What are the complications of a cardiovascular disease?"执行一次knn=2的检索验证；
5. 最后下载 MIRAGE QA 基准数据到data/mirage.json。

注意--index_num_chunks 0在 prepare.py 中被转换为None，从而在后续建索引时不做截断、使用全部文件。

语料下载机制解析

语料下载由 download.py 中的DownloadCorpora.download()方法完成，不同语料的来源与方式不同：

• PubMed / Textbooks / Wikipedia：通过git clone从 Hugging Face 的 MedRAG 数据集仓库拉取（如https://huggingface.co/datasets/MedRAG/<corpus>）。clone 时先以GIT_LFS_SKIP_SMUDGE=1跳过 LFS 大文件，再在目标目录内执行git lfs pull拉取全部大文件，避免一次性下载所有二进制内容；
• StatPearls：根据 StatPearls 的隐私政策，其内容不允许被再分发，因此不通过 Hugging Face 获取，而是用wget直接从 NCBI Bookshelf 下载statpearls_NBK430685.tar.gz，解压后调用 statpearls.py 进行分块处理；
• 若目标目录已存在，则跳过下载（幂等处理）。

StatPearls 语料为 NXML 格式，statpearls.py 使用xml.etree.ElementTree解析每个.nxml文件：提取文章标题、章节标题与子标题，拼接为title -- section -- subtitle形式的前缀；对正文段落（p标签）与列表（list标签）进行智能合并（短段落与相邻内容合并，单条列表项超过 1000 字符时拆分），最终为每个分块生成形如{id, title, content, contents}的 JSON 记录，写入corpus/statpearls/chunk/<filename>.jsonl，每个文件一行一个分块。这里生成的id形如文件名_序号，后续检索时会依据该格式还原文档名与分块序号。

FAISS 索引器：构建与查询

索引构建与查询都由 retriever.py 中的Retriever类完成，FAISS 库负责索引创建与相似度检索。

嵌入模型

默认使用sentence-transformers/all-MiniLM-L6-v2模型，将句子与段落映射到384 维的稠密向量空间。模型与维度均可在 retriever.yaml 中修改：

embedding_model: "sentence-transformers/all-MiniLM-L6-v2" embedding_dimension: 384

更换嵌入模型时，必须同步更新embedding_dimension为对应模型的实际输出维度，否则索引构建阶段 retriever.py 的维度过滤逻辑（embedding.shape == (self.emb_dim,)）会过滤掉全部向量。从源码看，设备放置由sentence_transformers.util.get_device_name()决定，并交由SentenceTransformer构造函数处理。

构建索引：build_faiss_index()

build_faiss_index(dataset_name, batch_size=32, num_chunks=None)的构建流程为：

1. 读取corpus/<dataset_name>/chunk/下所有.jsonl文件，num_chunks非空时只取前num_chunks个文件（便于开发调试）；
2. 按batch_size（默认 32）批量编码文本为嵌入向量（convert_to_numpy=True）；
3. 过滤维度不符的向量，统一转为float32（FAISS 的硬性要求）；
4. 以IndexFlatL2作为量化器（quantizer），簇数nlist = int(sqrt(嵌入总数))；
5. 构建faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)，执行train()与add()；
6. 保存faiss.index与all_doc_ids.npy（文档 ID 与嵌入一一对应，用于检索后映射回原文）。

每次调用会先删除旧的faiss.index与all_doc_ids.npy，保证索引可重建。注意：IndexIVFFlat需要先train()再add()，如果语料极少或嵌入维度异常，训练阶段可能因数据不足而失败，这也是默认只截取前 100 个 chunk 依然足以跑通的原因之一。

查询索引：query_faiss_index()

query_faiss_index(dataset_name, query, knn=8)是配套的检索辅助函数，前提是索引已生成（可通过Retriever.index_exists()检查faiss.index与all_doc_ids.npy是否存在）。流程为：

1. 加载faiss.index与all_doc_ids.npy；
2. 对查询文本生成嵌入向量；
3. 调用index.search()返回 top-k 的得分与索引位置；
4. 依据文档 ID 的文件名_序号格式，从chunk/目录还原出对应分块的title与content，按{rank, score, title, content}结构返回有序结果。

检索度量语义

当前实现使用IndexIVFFlat与faiss.METRIC_L2。由于 L2 距离衡量的是不相似度（dissimilarity），因此检索得分越低越好——这一定义贯穿客户端检索与服务端合并排序（见 README.md 的 Document Retrieval and Merge 一节）。

QA 基准数据集：MIRAGE

除语料与索引外，数据准备阶段还会下载 QA 基准数据集。所有 QA 基准数据来自 MIRAGE 基准，由 mirage_qa.py 中的MirageQA.download()从 MIRAGE 仓库的benchmark.json流式下载并保存为data/mirage.json。

FedRAG 流水线支持 PubMedQA、BioASQ、MMLU、MedQA、MedMCQA 五个 QA 基准（默认评估pubmedqa与bioasq的前 10 个问题）。prepare.py会在所有语料下载完成、所有 FAISS 索引构建完毕之后，才执行这一 QA 数据的下载。

磁盘占用与执行建议

综合 data/README.md 与 示例主 README：

• 仅下载Textbooks+StatPearls全部文档约需3 GB磁盘空间（默认快速模式仅拉取部分分块，占用更小）；
• 下载全部四个语料的全部文档约需120 GB；
• 再加上每个语料索引的同等磁盘开销，全量场景下请按约240 GB预留空间。

实操建议：先用默认命令./prepare.sh快速验证流程；确认跑通后再按需以--index_num_chunks 0对特定语料（如textbooks）做全量建索引；pubmed与wikipedia体量巨大，务必评估好下载时长与存储容量后再执行。

小结

examples/fedrag/data为 FedRAG 提供了完整的数据层支撑：从 MedRAG/Hugging Face 与 NCBI 拉取四个领域语料，经statpearls.py等脚本分块后，由Retriever基于all-MiniLM-L6-v2嵌入模型构建IndexIVFFlat + METRIC_L2的 FAISS 索引，并最终准备 MIRAGE QA 基准数据供流水线评测使用。理解了这一层，就能独立完成 FedRAG 示例从数据到检索的全部前置准备。

【免费下载链接】flowerFlower: A Friendly Federated AI Framework项目地址: https://gitcode.com/GitHub_Trending/flo/flower