OpenScholar 检索器实战:peS2o 语料库的索引构建、离线检索与在线 API 服务

发布时间:2026/10/14 1:52:30
OpenScholar 检索器实战:peS2o 语料库的索引构建、离线检索与在线 API 服务 人工智能RAGNLP微调模型推理服务【免费下载链接】OpenScholarThis repository includes the official implementation of OpenScholar: Synthesizing Scientific Literature with Retrieval-augmented LMs.项目地址https://gitcode.com/gh_mirrors/op/OpenScholar点击查看免费下载本文基于 OpenScholar 仓库中 retriever/README.md 的官方说明,完整讲解 peS2o 文献语料库检索器的三条核心操作链路:构建分片向量索引、离线 top-k 检索、以及基于 Flask 的在线检索 API。读完本文,你将能够理解ric/main_ric.py统一入口背后的任务调度机制、faiss 索引与分片数据的组织方式,并能独立配置 Hydrа 参数把检索结果接入 OpenScholar 的 RAG 推理主流程。一、检索器在 OpenScholar 中的定位与代码结构OpenScholar 是一个检索增强语言模型(Retrieval-Augmented LM):先检索相关论文段落,再基于召回内容生成带引用的回答。整个仓库分为三部分:src/:OpenScholar 主推理代码;training/:基于修改版torchtune训练 Llama 3.1 8B 的代码;retriever/:对 peS2o 数据做离线检索,并为在线检索提供服务的代码库,即本文主体。根目录 README.md 说明了 peS2o v2/v3 两个语料库(分块文本 索引)已公开发布,并提示:由于索引规模庞大(基于 4500 万篇论文的 2 亿 向量嵌入),peS2o 检索器需要大量 CPU 内存,官方主实验中采用离线方式先检索初始段落。从源码结构看,retriever/内部组织如下:路径作用ric/main_ric.py统一入口,根据tasks开关调度 embedding / index / search / 合并 / 困惑度评测ric/conf/pes2o_v3.yamlpeS2o v3 域的配置(与default.yaml合并)src/data.pyJSONL 数据加载、字节级分片、256 词切块src/embed.py段落向量化(Contriever 双编码器)src/index.pyfaiss 索引构建/反序列化、BM25 索引封装src/search.py稠密/稀疏 top-k 检索、多分片结果合并api/serve_pes2o.pyFlask 在线检索服务api/pes2o_ds.py在线服务用的数据存取(Datastore)封装contriever/修改版 Contriever 子模块(提供load_retriever等)入口文件 ric/main_ric.py 用自定义的 Hydrа 装饰器(见 src/hydra_runner.py)加载conf/下的配置,再按四个布尔开关顺序执行:hydra_runner(config_pathconf, config_namedefault) def main(cfg) - None: if cfg.tasks.datastore.get(embedding, False): generate_passage_embeddings(cfg) # 段落向量化 if cfg.tasks.datastore.get(index, False): build_index(cfg) # 构建索引 if cfg.tasks.eval.get(search, False): search_topk(cfg) # 离线 top-k 检索 if cfg.tasks.eval.get(merge_search, False): post_hoc_merge_topk_multi_domain(cfg) ...因此 README 中所有命令的形态都是python ric/main_ric.py --config-name... keyvalue,即通过命令行 override 开关和参数,而不需要改任何配置文件。二、环境准备按 retriever/README.md 的 Setup Environment 一节,进入检索器目录(原文档基于上游retrieval-scaling仓库,本仓库中对应目录为retriever/),用 conda 环境文件建环境:cd retrieval-scaling # 本仓库中即进入 retriever/ 目录 # create a new conda environment and install dependencies conda env create -f environment.yml conda activate scaling环境定义在 retriever/environment.yml,关键依赖包括:python3.11、torch、faiss-gpu1.8.0(向量索引);hydra-core、omegaconf(配置系统);transformers、sentence_transformers(编码器模型加载);pyserini(BM25 稀疏检索)、datasketch(去重)、tqdm。后续所有命令都要设置PYTHONPATH.并从retriever/目录执行,因为代码内部使用from src.xxx import ...与contriever.src.xxx这类相对导入。三、Step 1:准备 JSONL 数据数据加载器假设输入为 JSONL 格式。需要下载peS2o v3 数据(Hugging Face 上allenai/peS2o数据集的data/v3部分),将其转换为 JSONL 格式(每行一个含id与text字段的对象),并把转换后的 JSONL 文件统一保存到一个目录,后续通过datastore.raw_data_path指向该目录。raw_data_path既可以是单个 JSONL 文件,也可以是一个 JSONL 文件目录。真正读取时,src/data.py 中的fast_load_jsonl_shard做了两件关键的事:按字节大小分片:先把目录下所有文件的总大小均分为num_shards份,每个分片进程只seek到自己那一段字节范围再逐行解析——这对 2 亿级向量的语料库避免了整库载入内存;按词切块:每行文本经split_data_into_chunks(见 src/data.py)按chunk_size(默认256 词)滑窗切块,keep_last_chunktrue时保留最后不足一块的残块。每个段落被写为如下结构并缓存为raw_passages-{shard}-of-{num_shards}.pkl(或.jsonl):{text: chunk, origin_id: ex[id], id: idx, shard_id: shard_index, num_shards: num_shards}注意id是分片内从 0 递增的局部编号;跨分片的全局 id 偏移在构建索引与检索时另行累加(src/index.py 的get_index_passages_and_id_map)。四、Step 2:构建索引(Embedding Index)retriever/README.md 给出的完整构建脚本如下,官方推荐使用 SLURM 数组任务并行:datastore_raw_data_path$PES2O_V3_JSONL_DIR # directory that contains the converted jsonl data num_shards16 for SLURM_ARRAY_TASK_ID in {0..15}; do # we recommend to parallelize this with slurm jobs PYTHONPATH. python ric/main_ric.py \ --config-namepes2o_v3 \ tasks.datastore.embeddingtrue \ tasks.datastore.indextrue \ datastore.raw_data_path$datastore_raw_data_path \ datastore.embedding.num_shards$num_shards \ datastore.embedding.shard_ids[$SLURM_ARRAY_TASK_ID] \ hydra.job_logging.handlers.file.filenameembedding.log done参数说明:参数含义--config-namepes2o_v3使用 ric/conf/pes2o_v3.yaml,覆盖default.yaml中的模型与路径tasks.datastore.embeddingtrue/tasks.datastore.indextrue依次执行段落向量化与索引构建(见第一节调用链)datastore.raw_data_path指向转换后的 JSONL 目录(配置中默认null,必须 override)datastore.embedding.num_shards16数据切为 16 片datastore.embedding.shard_ids[N]当前进程只处理第 N 片,循环 0..15 即可全量并行4.1 向量化阶段(embedding)generate_passage_embeddings实现在 src/embed.py:按模型名分发加载器:模型路径含contriever时走 contriever/src/contriever.py 的load_retriever。pes2o_v3.yaml中model.datastore_encoder默认为akariasai/pes2o_contriever,即针对 peS2o 语料微调过的双编码器;除no_fp16为真外以fp16上卡前向,per_gpu_batch_size默认 512,passage_maxlength与chunk_size同为 256;每个分片产出embedding_dir下的passages_{shard:02d}.pkl(内容为ids embeddings的 pickle)。use_saved_if_existstrue时已存在的分片直接跳过,支持断点续跑。4.2 索引阶段(index)build_index根据model.sparse_retriever是否为空选择 BM25 或稠密索引,peS2o 流程走稠密路径 src/index.py 的build_dense_index。核心是Indexer类(src/index.py):def __init__(self, vector_sz, n_subquantizers0, n_bits8): if n_subquantizers 0: self.index faiss.IndexPQ(vector_sz, n_subquantizers, n_bits, faiss.METRIC_INNER_PRODUCT) else: self.index faiss.IndexFlatIP(vector_sz)pes2o_v3.yaml中n_subquantizers: 0,因此实际使用faiss 内积扁平索引IndexFlatIP(768)——保证检索精度,但要求全部向量常驻内存,这也解释了根 README 强调的大内存需求。构建时逐分片读入 pkl、按indexing_batch_size(默认 100 万)批量add,并用index_shard_ids的累加偏移把局部 id 映射为全局 id;完成后序列化为index.faiss与index_meta.faiss(id 映射表)两个文件。index_shard_ids支持两种模式(src/index.py):[0, 1, 2]表示多个分片建一个索引;[[0], [1, 2]]表示为psg-0与psg-12各建一个索引(即多索引模式)。peS2o 场景采用后者——16 个分片各自一个索引,检索时再合并。pes2o_v3.yaml中与索引构建相关的关键默认值:datastore: domain: pes2o_v3 chunk_size: 256 # 词数 datastore_root_dir: /gscratch/zlab/rulins/data/scaling_out # 内部路径,建议 override embedding: num_shards: 16 per_gpu_batch_size: 512 no_fp16: False index: projection_size: 768 # 向量维度 n_subquantizers: 0 # 0 faiss 扁平内积索引 indexing_batch_size: 1000000 save_or_load_index: True # 已存在则复用注意datastore_root_dir等默认值指向上游集群的内部绝对路径(如/gscratch/...),本地运行时应通过datastore.datastore_root_dir...覆写;save_or_load_index: true且索引已存在时直接deserialize,不会重复构建。五、离线检索(Offline Search)5.1 准备查询查询文件需为 JSONL,每行至少含query键,例如{query: Text used as a retrieval query, other_key: other_value}。源码侧对应 src/data.py 的prepare_lm_eval_data:当tasks.eval.task_namelm-eval时,直接取ex[query]作为raw_query。5.2 检索命令retriever/README.md 给出的完整脚本:EVAL_DOMAIN$EXP_ID # self-defined exp id to distinguish output RAW_QUERY$PATH_TO_QUERY_JSONL_FILE DS_NAMEpes2o_v3 NUM_SHARDS16 N_DOCS100 index_list[[0] for (( i1; i$((NUM_SHARDS - 1)); i )); do index_list,[$i] done index_list] echo INDEX_IDS:$index_list PYTHONPATH. python ric/main_ric.py \ --config-name pes2o_v3 \ tasks.eval.task_namelm-eval \ tasks.eval.searchtrue \ datastore.embedding.num_shards$NUM_SHARDS \ datastore.embedding.shard_ids[] \ datastore.index.index_shard_ids$index_list \ evaluation.domain$EVAL_DOMAIN \ evaluation.data.eval_data$RAW_QUERY \ evaluation.search.n_docs$N_DOCS \ evaluation.eval_output_dir$EVAL_OUTPUT_DIR # where the retrieved documents will be saved要点:index_shard_ids[[0],[1],...,[15]]让检索器逐索引加载 16 个分片索引分别召回,再由merge_multi_index_results(默认 true)合并;n_docs100为每个查询保留的段落数(配置默认仅 1,必须 override);eval_output_dir决定结果落盘位置,配置默认值同样指向内部绝对路径,必须覆写。5.3 检索实现与输出格式search_topk在 src/search.py 中按sparse_retriever分派;peS2o 走search_dense_topk(src/search.py),流程为:用model.query_encoder(akariasai/pes2o_contriever)加载查询编码器(fp16),embed_queries按per_gpu_batch_size64批量编码,question_maxlength512;对每组index_shard_ids:Indexer(...).deserialize_from加载index.faiss,get_index_passages_and_id_map加载对应分片段落并按全局偏移建 id→段落 映射,随后index.search_knn(query_embeddings, n_docs)取 top-k(src/search.py);add_passages(src/search.py)把召回结果写回每条查询的ctxs字段,每个段落包含:id、source(域)、origin_id、retrieval text、retrieval next text(下一分块文本,便于上下文延续)、retrieval score;多索引模式下post_hoc_merge_topk(src/search.py)把各索引的ctxs拼接,按retrieval score降序排序并截断到n_docs,写出合并文件。输出文件路径由eval_output_dir 分片后缀 原始文件名_retrieved_results.jsonl组成(见 src/search.py),例如每个单索引结果在0/、1/… 子目录下,合并结果在0-1-...-15/子目录下。这份ctxs格式正是 OpenScholar 推理入口的输入契约:run.py 的process_input_data会把retrieval text规范化为text并去重,再交给 src/open_scholar.py 的OpenScholar管线(重排、生成、引用)。因此离线检索结果可直接配合根 README 的推理命令使用:python run.py \ --input_file YOUR_INPUT_FILE \ --model_name OpenScholar/Llama-3.1_OpenScholar-8B \ --use_contexts \ --output_file OUTPUT_FILE_PATH \ --top_n 10 --llama3 --zero_shot如需重排,追加--ranking_ce --reranker OpenScholar/OpenScholar_Reranker。六、在线 API 服务(Online API Serving)除离线批量检索外,retriever/api/提供了一个轻量在线服务,共两步。Step 1. 运行服务脚本cd retrieval-scaling/api/ # 本仓库中即 retriever/api/ python serve_pes2o.py服务实现在 api/serve_pes2o.py:启动时先hydra.initialize(config_pathconf)并 composepes2o配置(api/conf/pes2o.yaml),再通过pes2o_ds.get_datastore构造Pes2oDatastoreAPI;Flask 应用监听0.0.0.0:5005(api/serve_pes2o.py);所有查询由SearchQueue串行处理:同一时刻只有一个搜索持有锁,其余请求进入队列等待(api/serve_pes2o.py),避免多路请求同时读大索引造成内存峰值;除/search外,还提供/current_search(查看当前正在处理的查询)与/queue_size(查看排队长度)两个辅助接口。Pes2oDatastoreAPI(api/pes2o_ds.py)在初始化时一次性加载 faiss 索引、全部段落及 peS2o 原始 id 映射;search方法编码查询后search_knn,返回三组字段:pes2o IDs、scores、passages。Step 2. 发送请求curl -X POST localhost:FORWARD_PORT/search -H Content-Type: application/json -d {query: example query, domains: pes2o}从路由代码看(api/serve_pes2o.py),请求体实际读取三个字段:query、domains、n_docs(n_docs直接取自request.json[n_docs],建议显式给出)。成功时返回 200 及如下 JSON:{ message: Search completed for example query from pes2o, query: example query, n_docs: 3, results: {pes2o IDs: [...], scores: [...], passages: [...]} }若服务运行在远端节点,FORWARD_PORT为本地转发端口:按 api/README.md 的说明,通过 SSH 端口转发(ssh -NL 0.0.0.0:FORWARD_PORT:job node:SERVE_PORT)把远端服务端口映射到本地后,再用上面的 curl 访问localhost即可。使用在线服务前必须注意的路径适配api/目录的代码带有上游集群的硬编码路径,本地部署前需要自行适配(仓库为只读,请在自己的工作副本中修改):api/conf/pes2o.yaml中datastore.datastore_root_dir、datastore.embedding.num_shards等默认指向内部集群路径且为 8 分片;若服务 peS2o v3 16 分片索引,应改为--config-name覆盖或在配置中同步num_shards与index_shard_ids;api/pes2o_ds.py 中load_pes2o_id_mapping读取的updated_pes2o_id_mapping.pkl是硬编码的绝对路径,需替换为你本地的 peS2o id 映射文件;api/serve_pes2o.py 中查询日志路径self.query_log也是硬编码内部路径,需要改到本地可写位置。七、小结与适用前提适用前提:本流程面向 peS2o v2/v3 大规模文献语料。由于索引达 2 亿 向量、需大量 CPU 内存与 GPU 显存(向量化阶段),官方主实验采用离线检索 后续 Semantic Scholar / 网页 API 补充检索的混合方案(见 src/use_search_apis.py);可复用性:ric/main_ric.py的tasks开关 Hydrа override 机制意味着同一套代码可对任意 JSONL 语料建索引与检索,参考配置 ric/conf/example_config.yaml(含fineweb_edu_1m等示例)与 ric/conf/default.yaml;结果衔接:离线检索产出的ctxs(retrieval text/retrieval score/origin_id等字段)与 run.py 的输入规范化逻辑一一对应,可直接作为--use_contexts管线的上下文输入,实现peS2o 离线召回 → 重排 → Llama 3.1 OpenScholar-8B 生成的完整链路。更早期的索引/检索说明(与 README 基本一致的上游版本)见 retriever/build_pes2o_index.md,可作为命令参数的交叉参考。赞分享人工智能RAGNLP微调模型推理服务【免费下载链接】OpenScholarThis repository includes the official implementation of OpenScholar: Synthesizing Scientific Literature with Retrieval-augmented LMs.项目地址https://gitcode.com/gh_mirrors/op/OpenScholar点击查看免费下载相关推荐openJiuwen 检索系统 SDK 接口指南从离线索引构建到在线混合检索openJiuwen 检索系统 SDK 接口指南从离线索引构建到在线混合检索 导读 本文是 openJiuwen 生态技能分发与检索能力SDK 包名 ope后端前端AI 技能AI 插件openJiuwen 检索系统在线检索架构auto 统一检索路线与 Progressive 树检索实战openJiuwen 检索系统在线检索架构auto 统一检索路线与 Progressive 树检索实战 本文深入解读 openJiuwen 技能分发仓库sk后端前端AI 技能AI 插件openJiuwen Skill Hub 检索系统架构解析从离线索引构建到在线渐进式检索的完整链路openJiuwen Skill Hub 检索系统架构解析从离线索引构建到在线渐进式检索的完整链路 导读 本文以 marketplace/retrieval后端前端AI 技能AI 插件创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考