OpenScholar 仓库中的 Contriever：无监督稠密信息检索的完整实践——模型加载、检索评测与对比学习预训练

• 人工智能
• RAG
• NLP
• 微调
• 模型推理服务

本篇以 OpenScholar 仓库自带的retriever/contriever子仓库为对象，完整讲解 Contriever（对比学习预训练的无监督稠密检索器）的使用方法：如何用 HuggingFace 接口加载四个官方预训练模型并生成句子嵌入、如何复现 NaturalQuestions/TriviaQA、BEIR、Mr. TyDi 与跨语言 MKQA 的检索评测流程，以及如何从分词预处理到train.py启动 MoCo 式对比预训练。读完本文，你可以直接复制其中的命令在本仓库代码上跑通“下载数据 → 生成语料嵌入 → 建 FAISS 索引 → top-k 检索 → 计算召回”的完整链路，并理解--momentum、--temperature、--queue_size等参数在 MoCo 实现 中的真实作用。

一、Contriever 是什么：无监督稠密检索与模型家族

Contriever 源自论文Unsupervised Dense Information Retrieval with Contrastive Learning（G. Izacard 等，2021）。其核心思想是：不使用任何人工标注的 query–document 对，仅靠对比学习（contrastive learning）在 CC-Net 与英文维基百科等大规模无标注文本上预训练稠密检索模型；预训练后的模型在 BEIR 基准的 R@100 上即可与 BM25 相媲美，而在 MSMARCO 上进一步微调（finetuning）后召回率（recall@100）表现尤其突出（见 README）。

仓库提供四个官方预训练模型：

这四个模型都可以直接用 HuggingFacetransformers生态加载。运行环境上，该子仓库的 requirements.txt 固定了torch==1.11.0、transformers==4.18.0、beir==1.0.0，复现评测时建议按此版本搭建 conda 环境。

二、快速上手：加载模型并计算句子相似度

README 给出的最小示例：

from src.contriever import Contriever from transformers import AutoTokenizer contriever = Contriever.from_pretrained("facebook/contriever") tokenizer = AutoTokenizer.from_pretrained("facebook/contriever") # 加载配套分词器

对若干句子生成嵌入并用点积计算相似度：

sentences = [ "Where was Marie Curie born?", "Maria Sklodowska, later known as Marie Curie, was born on November 7, 1867.", "Born in Paris on 15 May 1859, Pierre Curie was the son of Eugène Curie, a doctor of French Catholic origin from Alsace." ] inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt") embeddings = contriever(**inputs) score01 = embeddings[0] @ embeddings[1] # 1.0473 score02 = embeddings[0] @ embeddings[2] # 1.0095

注意两点细节（原文示例中变量名为model，此处已对齐为contriever）：

1. 默认输出不做归一化。从源码看，Contriever 类 的forward只有传入normalize=True时才会执行F.normalize(emb, dim=-1)，因此默认返回的是原始向量，句子间相似度用裸点积计算（这也解释了示例中 score 约为 1.0 而非余弦值）。
2. 句向量采用平均池化。forward先把 padding 位置的last_hidden_state置零（masked_fill），再按config.pooling选择average（对有效 token 求平均）或cls（取第一个位置）作为句子表示。

四个模型的加载方式完全一致：

from src.contriever import Contriever contriever = Contriever.from_pretrained("facebook/contriever") contriever_msmarco = Contriever.from_pretrained("facebook/contriever-msmarco") mcontriever = Contriever.from_pretrained("facebook/mcontriever") mcontriever_msmarco = Contriever.from_pretrained("facebook/mcontriever-msmarco")

2.1 源码层面的模型封装

除Contriever外，contriever.py 还定义了：

• XLMRetriever（XLMRobertaModel子类，L58-L100）：结构与Contriever相同，用于retriever_model_id中含xlm的模型；
• load_retriever(model_path, pooling="average", random_init=False)（L103-L138）：检索/评测脚本统一使用的加载入口。它分两条路径：
  • 若model_path下存在checkpoint.pth（训练产出的本地检查点），则读取其中保存的opt.retriever_model_id重建模型类（含xlm关键字则用XLMRetriever），并自动剥离encoder_q./encoder.前缀后加载权重——这意味着train.py训练出的检查点可以直接被passage_retrieval.py等脚本复用；
  • 否则按 HuggingFace Hub id / 本地目录走AutoConfig+AutoTokenizer标准加载。

mcontriever/mcontriever-msmarco底层基于bert-base-multilingual-cased（mBERT），因此评测多语言数据时需要额外的文本归一化，这在 normalize_text.py 中实现：把控制字符、多种连字符/破折号、引号、省略号等映射到 mBERT 词表中存在的规范形式（例如把 en/em dash 统一为-、…替换为...）。

三、评测（一）：问答检索（NaturalQuestions / TriviaQA）

这是 README 中最完整的端到端工作流：给定问题，从维基百科段落库中检索 top-100，再用“答案是否出现在检索段落中”计算 R@5/R@20/R@100。

数据准备：NaturalQuestions 与 TriviaQA 数据需从 FiD 仓库下载（注意 README 特别说明：这里的 NQ 数据沿用原始 NQ 的答案，与 DPR 仓库做了后处理的版本略有差异，会影响分词）。检索所用的维基百科段落（DPR 的psgs_w100.tsv）下载方式：

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

第 1 步：生成段落嵌入，使用 generate_passage_embeddings.py：

python generate_passage_embeddings.py \ --model_name_or_path facebook/contriever \ --output_dir contriever_embeddings \ --passages psgs_w100.tsv \ --shard_id 0 --num_shards 1 \

从源码可以看到其可配置项：--per_gpu_batch_size（默认 512）、--passage_maxlength（默认 512）、--no_fp16（默认使用 fp16 推理）、--no_title（默认把title + " " + text拼接后编码，见 embed_passages）、--lowercase、--normalize_text，以及--shard_id/--num_shards分片（各分片输出 pickle 文件passages_XX，可多机并行生成）。

也可以直接下载官方预计算好的嵌入，跳过此步：

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

第 2 步：top-100 检索，使用 passage_retrieval.py：

python passage_retrieval.py \ --model_name_or_path facebook/contriever \ --passages psgs_w100.tsv \ --passages_embeddings "contriever_embeddings/*" \ --data nq_dir/test.json \ --output_dir contriever_nq \

其关键参数（见 参数定义）：

• --n_docs（默认 100）：每个问题检索的段落数；
• --save_or_load_index：若嵌入目录下已有index.faiss则直接反序列化，否则现场建索引并落盘；
• --n_subquantizers/--n_bits：向量量化选项。从 Indexer 看，n_subquantizers=0（默认）时使用精确的faiss.IndexFlatIP（内积），设为正数则使用faiss.IndexPQ乘积量化索引以压缩内存；
• --question_maxlength（默认 512）、--per_gpu_batch_size（默认 64）、--no_fp16、--lowercase、--normalize_text。

流程上，脚本先glob所有嵌入分片 → 分批index_data灌入 FAISS → 对问题批量编码（同样支持--lowercase/--normalize_text）→index.search_knn取 top-k → 用calculate_matches并行判定每条段落是否包含答案 → 把ctxs（id/title/text/score）与hasanswer写回输出 JSONL，并打印 R@5/10/20/100。

复现结果（README 原表）：

四、评测（二）：BEIR 基准

BEIR 上各数据集的 nDCG@10 可以用仓库内入口脚本复现。README 中写的是beireval.py，而当前仓库实际提供的文件是 eval_beir.py（二者参数一致，以仓库内文件名为准）：

python eval_beir.py --model_name_or_path contriever-msmarco --dataset scifact

eval_beir.py的关键参数：--dataset（BEIR 数据集名）、--beir_dir（数据集缓存目录，缺失时 beir_utils.evaluate_model 会按 BEIR 官方 URL 自动下载解压）、--score_function（默认dot）、--norm_query/--norm_doc（对查询/文档嵌入归一化）、--lower_case、--normalize_text、--per_gpu_batch_size（默认 128）。其内部用 BEIR 库的DenseRetrievalExactSearch+ 自实现的DenseEncoderModel（对 query/corpus 分批编码，支持多进程切分与varsize_gather_nograd聚合），文档编码规则同样是title + " " + text。

README 特别提醒：Touche-2020 数据集在 BEIR 中更新过版本，若使用当前版本，结果会与下表不同。

nDCG@10：

R@100：

* README 原表中两行均标注为 “Contriever-msmarco”；对照上方 nDCG 表的行序（第一行为未微调的 Contriever），可推断第一行对应该为无监督版 Contriever 的 R@100 结果，引用时建议保留此说明。

五、多语言评测

5.1 Mr. TyDi v1.1

mContriever 在 Mr. TyDi v1.1（11 种语言）上评测，流程为“下载 → 转 BEIR 格式 → 评测”。以 Swahili 为例：

# 1. 下载数据 wget https://git.uwaterloo.ca/jimmylin/mr.tydi/-/raw/master/data/mrtydi-v1.1-swahili.tar.gz -P mrtydi tar -xf mrtydi/mrtydi-v1.1-swahili.tar.gz -C mrtydi gzip -d mrtydi/mrtydi-v1.1-swahili/collection/docs.jsonl.gz # 2. 转换为 BEIR 格式（queries.jsonl / corpus.jsonl / qrels/{train,dev,test}.tsv） python data_scripts/convertmrtydi2beir.py mrtydi/mrtydi-v1.1-swahili mrtydi/mrtydi-v1.1-swahili # 3. 评测 python eval_beir.py --model_name_or_path facebook/mcontriever --dataset mrtydi/mrtydi-v1.1-swahili --normalize_text

其中 convertmrtydi2beir.py 把topic.tsv转成queries.jsonl（_id+text），把collection/docs.jsonl的id/contents映射为 BEIR 的_id/text（title置空），并把qrels.{split}.txt转成三列 TSV。评测时务必带--normalize_text，以与 mBERT 词表对齐（见第二节 normalize_text.py 的说明）。

MRR@100（语言列：ar, bn, en, fi, id, ja, ko, ru, sw, te, th）：

R@100：

（“+ Mr. TyDi”表示在各自语言 TyDi 数据上做了语言内微调。）

5.2 跨语言 MKQA（xMKQA）

目标是衡量“用另一种语言的问题去检索英文维基百科”的能力：给定非英语 query，看英文答案是否出现在检回的段落中（沿用 DPR 的 hasanswer 判定）。步骤：

# 1. 下载 MKQA 数据 wget https://raw.githubusercontent.com/apple/ml-mkqa/master/dataset/mkqa.jsonl.gz # 2. 预处理：按语言拆分为 xmkqa/<lang>.jsonl（25 种语言），英文答案+别名作为 answers python data_scripts/preprocess_xmkqa.py mkqa.jsonl xmkqa # 3. 生成段落嵌入（注意 mContriever 需 --lowercase --normalize_text） python generate_passage_embeddings.py \ --model_name_or_path facebook/mcontriever \ --output_dir mcontriever_embeddings \ --passages psgs_w100.tsv \ --shard_id 0 --num_shards 1 \ --lowercase --normalize_text \ # 也可直接下载官方预计算嵌入 wget https://dl.fbaipublicfiles.com/contriever/embeddings/mcontriever/wikipedia_embeddings.tar wget https://dl.fbaipublicfiles.com/contriever/embeddings/mcontriever-msmarco/wikipedia_embeddings.tar # 4. 检索并计算召回（--data 支持通配符，一次评测所有语言） python passage_retrieval.py \ --model_name_or_path facebook/mcontriever \ --passages psgs_w100.tsv \ --passages_embeddings "mcontriever_embeddings/*" \ --data "xmkqa/*.jsonl" \ --output_dir mcontriever_xmkqa \ --lowercase --normalize_text \

preprocess_xmkqa.py 的实现细节：它只保留存在非unanswerable/binary/long_answer类型答案的样本，对 25 种语言（en, ar, fi, ja, ko, ru, es, sv, he, th, da, de, fr, it, nl, pl, pt, hu, vi, ms, km, no, tr, zh-cn, zh-hk, zh-tw）各生成一条{id, lang, question, answers}记录，其中answers统一使用英文答案及其别名——这正是“跨语言检索英文语料”设定在数据层的体现。

R@100 / R@20（摘要，完整 25 语言数据见 README 原表）：

可以看到两个规律：MSMARCO 微调带来约 10+ 点的平均提升，且 Malay（ms）这类语料薄弱语言始终是最短板。

六、训练：数据预处理与对比学习

6.1 数据预处理：分词 + 切块

Contriever（英文单语）在英文 Wikipedia 与 CC-Net 上训练；mContriever（29 语言）在 CC-Net 上训练。原始文本先转成纯文本文件，然后用 data_scripts/tokenization_script.sh 完成分词与切块：

NSPLIT=128 #Must be larger than the number of processes used during training FILENAME=en_XX.txt INFILE=./${FILENAME} TOKENIZER=bert-base-uncased #TOKENIZER=bert-base-multilingual-cased SPLITDIR=./tmp-tokenization-${TOKENIZER}-${FILENAME}/ OUTDIR=./encoded-data/${TOKENIZER}/$(echo "$FILENAME" | cut -f 1 -d '.') NPROCESS=8 mkdir -p ${SPLITDIR} split -a 3 -d -n l/${NSPLIT} ${INFILE} ${SPLITDIR} for ((i=0;i<$NSPLIT;i++)); do num=$(printf "%03d\n" $i); FILE=${SPLITDIR}${num}; #we used --normalize_text as an additional option for mContriever python3 preprocess.py --tokenizer ${TOKENIZER} --datapath ${FILE} --outdir ${OUTDIR} & ... done

要点（结合 preprocess.py 源码）：

• 先把文本split成NSPLIT个子文件（必须大于训练进程数），每个子文件独立调用preprocess.py分词后保存为{文件名}.pkl（一维input_ids张量，add_special_tokens=False）；
• 训练时--loading_mode split会让各分布式 rank 各取一部分.p*张量（见 load_dataset 中np.array_split(files, world_size)[rank]），从而把大数据分散到各进程；
• mContriever 必须加--normalize_text：在分词前用第二节提到的归一化函数处理 mBERT 词表外的常见字符，脚本注释中也明确标注了这一点。

训练样本的构造在 src/data.py：从一个chunk_length（默认 256 token）的滑窗中，用randomcrop按--ratio_min~--ratio_max比例随机裁出两段作为 query/“key”文本，再分别做可选增强并添加[CLS]/[SEP]。支持四种增强（--augmentation，由 apply_augmentation 实现）：delete（按概率删词）、replace（随机词替换）、mask（置为[MASK]）、shuffle（乱序若干词），概率由--prob_augmentation控制。多个数据目录之间由MultiDataset按样本量加权随机采样（--sampling_coefficient可调节数据集间的采样偏斜）。

6.2 对比训练：train.py 与 MoCo 实现

train.py 是预训练入口，--contrastive_mode支持两种模式：

• moco（默认，MoCo 类）：维护一个动量更新的 key encoder 和一个负样本队列；
• inbatch（InBatch 类）：单编码器 + 批内互为负例（跨 GPU gather 后以对角线为正例）。

Contriever（英文）官方配置，32 GPU：

python train.py \ --retriever_model_id bert-base-uncased --pooling average \ --augmentation delete --prob_augmentation 0.1 \ --train_data "data/wiki/ data/cc-net/" --loading_mode split \ --ratio_min 0.1 --ratio_max 0.5 --chunk_length 256 \ --momentum 0.9995 --queue_size 131072 --temperature 0.05 \ --warmup_steps 20000 --total_steps 500000 --lr 0.00005 \ --scheduler linear --optim adamw --per_gpu_batch_size 64 \ --output_dir /checkpoint/gizacard/contriever/xling/contriever \

mContriever（多语言，29 语言）官方配置，32 GPU：

TDIR=encoded-data/bert-base-multilingual-cased/ TRAINDATASETS="${TDIR}fr_XX ${TDIR}en_XX ${TDIR}ar_AR ${TDIR}bn_IN ${TDIR}fi_FI ${TDIR}id_ID ${TDIR}ja_XX ${TDIR}ko_KR ${TDIR}ru_RU ${TDIR}sw_KE ${TDIR}hu_HU ${TDIR}he_IL ${TDIR}it_IT ${TDIR}km_KM ${TDIR}ms_MY ${TDIR}nl_XX ${TDIR}no_XX ${TDIR}pl_PL ${TDIR}pt_XX ${TDIR}sv_SE ${TDIR}te_IN ${TDIR}th_TH ${TDIR}tr_TR ${TDIR}vi_VN ${TDIR}zh_CN ${TDIR}zh_TW ${TDIR}es_XX ${TDIR}de_DE ${TDIR}da_DK" python train.py \ --retriever_model_id bert-base-multilingual-cased --pooling average \ --train_data ${TRAINDATASETS} --loading_mode split \ --ratio_min 0.1 --ratio_max 0.5 --chunk_length 256 \ --momentum 0.999 --queue_size 32768 --temperature 0.05 \ --warmup_steps 20000 --total_steps 500000 --lr 0.00005 \ --scheduler linear --optim adamw --per_gpu_batch_size 64 \ --output_dir /checkpoint/gizacard/contriever/xling/mcontriever \

两个差异值得注意：mContriever 不用词级增强（--augmentation none），且队列更小、动量略低（32768 / 0.999 vs 131072 / 0.9995）。完整 SLURM 作业脚本见 example_scripts/contriever.sh 与 example_scripts/mcontriever.sh（各 4/8 节点 × 8 GPU，srun提交，--main_port随机端口）。

参数名提示：README 早期示例写作--moco_queue 131072，而当前代码库 src/options.py 中对应的实际选项名为--queue_size（默认 65536），直接运行请以代码为准。

参数在 MoCo 源码中的确切含义（moco.py）：

• --temperature：logits 除数（logits = self._compute_logits(q, k) / self.temperature，L118），0.05的低温使对比信号更尖锐；
• --momentum：key encoder 的动量更新系数（_momentum_update_key_encoder：param_k = m*param_k + (1-m)*param_q）；
• --queue_size：负样本环形队列大小（register_buffer("queue", ...)，L39-L42），每个 key 经cross_entropy(logits, labels=0)与“1 正例 + 队列全部负例”做分类，随后_dequeue_and_enqueue跨 GPU gather 后把新 key 轮转写入队列（L81-L95）；
• --label_smoothing（默认 0）、--norm_query/--norm_doc：分别控制 logits 的 label smoothing 和两侧嵌入是否 L2 归一化；
• --pooling average|cls、--retriever_model_id、--chunk_length、--ratio_min/--ratio_max、--augmentation/--prob_augmentation、--sampling_coefficient则对应第二节的样本构造逻辑。

训练循环本身（train.py 主流程）：DataLoader(batch_size=per_gpu_batch_size, collate_fn=Collator)→model(**batch)返回 loss → 每eval_freq步在--eval_datasets（BEIR 格式）上算 NDCG@10 / Recall@10 / Recall@100 → 每save_freq步用 utils.save 保存checkpoint.pth（含模型、优化器、scheduler、完整opt，并维护latest软链）——这就是 3.1 节load_retriever能直接识别训练产物的原因。模型在total_steps（示例为 500000 步）内持续训练，warmup_steps之后按--scheduler linear线性衰减学习率。

MSMARCO 微调由 finetuning.py 承担：它复用train.py的选项体系（--model_path指向预训练检查点、--eval_datasets指定 BEIR 数据集做训练中评测），并引入--negative_ctxs（每个正例配多少个负例）、--negative_hard_ratio、--negative_hard_min_idx等难负例采样选项，数据来自src/finetuning_data.py；contriever-msmarco/mcontriever-msmarco即此流程的产物。

七、与 OpenScholar 检索栈的关系

OpenScholar 主仓库的 retriever/ 目录是面向 peS2O 语料的大规模检索栈（ric/main_ric.py+ Hydra 配置 + IVFPQ 索引）；retriever/contriever作为独立子仓库被纳入其中，提供了一套“开箱即用”的稠密检索器实现——无论是直接调用Contriever.from_pretrained为问答/文献检索生成句子向量，还是复用generate_passage_embeddings.py+passage_retrieval.py的“嵌入分片 + FAISS”工程模式，都可以作为 OpenScholar 检索流水线的对照与补充。

八、引用与环境说明

若在你的研究中使用了本仓库的 Contriever 实现与模型，README 建议引用：

G. Izacard, M. Caron, L. Hosseini, S. Riedel, P. Bojanowski, A. Joulin, E. Grave.Unsupervised Dense Information Retrieval with Contrastive Learning. 2021.

@misc{izacard2021contriever, title={Unsupervised Dense Information Retrieval with Contrastive Learning}, author={Gautier Izacard and Mathilde Caron and Lucas Hosseini and Sebastian Riedel and Piotr Bojanowski and Armand Joulin and Edouard Grave}, year={2021}, url = {https://arxiv.org/abs/2112.09118}, doi = {10.48550/ARXIV.2112.09118}, }

适用前提与限制小结：代码固定于torch==1.11.0 / transformers==4.18.0 / beir==1.0.0环境；passage_retrieval.py与generate_passage_embeddings.py依赖 GPU（model.cuda()、fp16 推理），大规模建索引建议加--save_or_load_index并考虑--n_subquantizers量化；BEIR 中 Touche-2020 使用新版数据时指标会与 README 表格不一致；多语言评测务必保持--lowercase --normalize_text与预计算嵌入生成时一致，否则相似度尺度会偏移。

• 人工智能
• RAG
• NLP
• 微调
• 模型推理服务