ColBERT (v2)

ColBERT 是一个快速且准确的检索模型，能够在几十毫秒内实现基于 BERT 的大规模文本集合搜索。

图 1： ColBERT 的晚期交互机制，高效计算查询与段落之间的细粒度相似度。

如图 1 所示，ColBERT 依赖于细粒度的上下文晚期交互：它将每个段落编码为一个由 token 级 embedding 组成的矩阵（上图蓝色部分）。在搜索时，它将每个查询嵌入为另一个矩阵（绿色部分），并通过可扩展的向量相似度（MaxSim）算子高效地找到与查询上下文匹配的段落。

这种丰富的交互使 ColBERT 能够超越_单向量_表示模型的质量，同时高效扩展到大规模语料库。更多详情请参阅我们的论文：

• ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT (SIGIR'20)
• Relevance-guided Supervision for OpenQA with ColBERT (TACL'21)
• Baleen: Robust Multi-Hop Reasoning at Scale via Condensed Retrieval (NeurIPS'21)
• ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction (NAACL'22)
• PLAID: An Efficient Engine for Late Interaction Retrieval (CIKM'22)

🚨 公告

• (2023年1月29日) 我们合并了一个新的索引更新器功能，并增加了对更多 Hugging Face 模型的支持！这些功能目前处于测试阶段，欢迎试用并提供反馈。
• (2023年1月24日) 如果您在寻找用于组合 ColBERTv2 和 LLM 的 DSP 框架，请访问：https://github.com/stanfordnlp/dsp

ColBERTv1

SIGIR'20 论文中的 ColBERTv1 代码位于 colbertv1 分支。其他分支的信息请参见此处。

安装

ColBERT 需要 Python 3.7+ 和 Pytorch 1.9+，并使用 Hugging Face Transformers 库。

我们强烈建议使用以下命令创建 conda 环境。（如果您没有 conda，请按照官方 conda 安装指南 进行安装。）

我们还提供了一个专门用于纯 CPU 环境的新环境文件（conda_env_cpu.yml），但请注意，如果您在包含 GPU 的机器上测试 CPU 执行，可能需要在命令中指定 CUDA_VISIBLE_DEVICES=""。此外，训练和索引需要 GPU。

conda env create -f conda_env[_cpu].yml
conda activate colbert

如果遇到任何问题，请提交新 issue，我们会及时帮助您！

概述

在数据集上使用 ColBERT 通常包括以下步骤。

步骤 0：预处理您的语料库。 在最简单的情况下，ColBERT 使用制表符分隔（TSV）文件：一个文件（例如 collection.tsv）包含所有段落，另一个文件（例如 queries.tsv）包含用于搜索语料库的一组查询。

步骤 1：下载预训练的 ColBERTv2 checkpoint。 该 checkpoint 已在 MS MARCO 段落排序任务上训练完成。您也可以_选择性地_训练自己的 ColBERT 模型。

步骤 2：索引您的语料库。 一旦拥有训练好的 ColBERT 模型，您需要索引您的语料库以实现快速检索。此步骤将所有段落编码为矩阵，存储在磁盘上，并构建用于高效搜索的数据结构。

步骤 3：使用查询搜索语料库。 基于模型和索引，您可以对语料库发起查询，为每个查询检索 top-k 段落。

下面，我们通过 MS MARCO 段落排序任务的示例运行来演示这些步骤。

API 使用笔记本

新功能：我们提供了一个实验性笔记本，可在 Google Colab 上使用免费 GPU 运行。在免费的 Colab T4 GPU 上索引 10,000 个段落只需六分钟。

该 Jupyter 笔记本 docs/intro.ipynb 笔记本 演示了使用 ColBERT 的关键功能。