<p align="center"> <img align="center" src="https://github.com/stanford-futuredata/ColBERT/blob/main/docs/images/colbertofficial.png?raw=true" width="430px" /> </p> <p align="left">

ColBERT (v2)

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

<img align="center" src="https://colab.research.google.com/assets/colab-badge.svg" />

<p align="center"> <img align="center" src="https://github.com/stanford-futuredata/ColBERT/blob/main/docs/images/ColBERT-Framework-MaxSim-W370px.png?raw=true" /> </p> <p align="center"> <b>图 1:</b> ColBERT 的延迟交互，高效地计算查询与段落之间的细粒度相似度。 </p>

如图 1 所示，ColBERT 依赖于细粒度的上下文延迟交互：它将每个段落编码为一个矩阵，即 token 级别的嵌入向量（图中以蓝色显示）。然后在搜索时，它将每个查询嵌入到另一个矩阵中（图中以绿色显示），并使用可扩展的向量相似度（MaxSim）算子高效地找到与查询上下文匹配的段落。

这些丰富的交互使 ColBERT 能够超越单向量表示模型的质量，同时能够高效地扩展到大型语料库。您可以在以下论文中了解更多：

• ColBERT：通过 BERT 上的上下文延迟交互实现高效且有效的段落检索（SIGIR'20）。
• 基于相关性引导的监督实现 OpenQA 与 ColBERT（TACL'21）。
• Baleen：通过压缩检索实现大规模多跳推理（NeurIPS'21）。
• ColBERTv2：通过轻量级延迟交互实现高效且有效的检索（NAACL'22）。
• PLAID：延迟交互检索的高效引擎（CIKM'22）。

🚨 公告

• （1/29/23）我们合并了新的索引更新器功能，并支持更多 Hugging Face 模型！这些功能仍在测试阶段，请在使用过程中提供反馈。
• （1/24/23）如果您正在寻找用于组合 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 检查点。 该检查点已在 MS MARCO 段落排序任务上进行了训练。您也可以可选地训练您自己的 ColBERT 模型。

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

步骤 3：使用您的查询搜索集合。 有了模型和索引后，您可以对集合发出查询以检索每个查询的前 k 个段落。

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

API 使用教程 Notebook

新增功能：我们有一个实验性的 Google Colab Notebook，您可以使用免费的 GPU。使用免费的 Colab T4 GPU 为 10,000 条数据建立索引需要六分钟。

这个 Jupyter Notebook docs/intro.ipynb 展示了如何使用新版 Python API 使用 ColBERT 的关键功能。

它包括如何下载在 MS MARCO 段落排序任务上训练的 ColBERTv2 模型检查点，以及如何下载我们的新 LoTTE 基准数据集。

数据

该仓库直接使用简单的制表符分隔文件格式来存储查询、段落和前 k 名排序列表。

• 查询：每行格式为 qid \t 查询文本。
• 集合：每行格式为 pid \t 段落文本。
• 前 k 名排序：每行格式为 qid \t pid \t 排名。

这与 MS MARCO 段落排序数据集的数据格式直接兼容。您将需要训练三元组（triples.train.small.tar.gz）、dev 集查询的官方前 1000 名排序列表（top1000.dev）以及 dev 集相关段落（qrels.dev.small.tsv）。要为完整集合建立索引，您还需要段落列表（collection.tar.gz）。

索引

为了快速检索，索引会预计算段落的 ColBERT 表示。

示例用法：

from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Indexer

if __name__=='__main__':
    with Run().context(RunConfig(nranks=1, experiment="msmarco")):

        config = ColBERTConfig(
            nbits=2,
            root="/path/to/experiments",
        )
        indexer = Indexer(checkpoint="/path/to/checkpoint", config=config)
        indexer.index(name="msmarco.nbits=2", collection="/path/to/MSMARCO/collection.tsv")

检索

我们通常建议您使用 ColBERT 进行端到端检索，它直接从完整集合中找出前 k 个段落：

from colbert.data import Queries
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher

if __name__=='__main__':
    with Run().context(RunConfig(nranks=1, experiment="msmarco")):

        config = ColBERTConfig(
            root="/path/to/experiments",
        )
        searcher = Searcher(index="msmarco.nbits=2", config=config)
        queries = Queries("/path/to/MSMARCO/queries.dev.small.tsv")
        ranking = searcher.search_all(queries, k=100)
        ranking.save("msmarco.nbits=2.ranking.tsv")

您可以选择指定 ncells、centroid_score_threshold 和 ndocs 搜索超参数，以在速度和结果质量之间进行权衡。不同的 k 值的默认值列在 colbert/searcher.py 中。

我们可以使用以下命令评估 MSMARCO 排序结果：

python -m utility.evaluate.msmarco_passages --ranking "/path/to/msmarco.nbits=2.ranking.tsv" --qrels "/path/to/MSMARCO/qrels.dev.small.tsv"

训练

我们提供了一个预训练模型检查点，但我们也在此详细说明如何从头开始训练。 请注意，此示例演示了 ColBERTv1 风格的训练，但提供的检查点是用 ColBERTv2 训练的。

训练需要一个 JSONL 三元组文件，每行包含 [qid, pid+, pid-] 列表。查询 ID 和段落 ID 分别对应指定的 queries.tsv 和 collection.tsv 文件。

示例用法（在 4 个 GPU 上训练）：

from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Trainer

if __name__=='__main__':
    with Run().context(RunConfig(nranks=4, experiment="msmarco")):

        config = ColBERTConfig(
            bsize=32,
            root="/path/to/experiments",
        )
        trainer = Trainer(
            triples="/path/to/MSMARCO/triples.train.small.tsv",
            queries="/path/to/MSMARCO/queries.train.small.tsv",
            collection="/path/to/MSMARCO/collection.tsv",
            config=config,
        )

        checkpoint_path = trainer.train()

        print(f"检查点已保存至 {checkpoint_path}...")

运行轻量级 ColBERTv2 服务器

我们提供了一个脚本，用于运行轻量级服务器，该服务器以 JSON 格式为给定搜索查询返回按排名顺序排列的 k 个（最多 100 个）结果。此脚本可用于驱动 DSP 程序。

要运行服务器，请在 .env 文件中更新环境变量 INDEX_ROOT 和 INDEX_NAME，以指向相应的 ColBERT 索引。然后运行以下命令：

python server.py

示例查询：

http://localhost:8893/api/search?query=Who won the 2022 FIFA world cup&k=25

分支

支持的分支

• main：包含 ColBERTv2 + PLAID 的稳定分支。
• colbertv1：ColBERTv1 的遗留分支。

已弃用的分支

• new_api：ColBERTv2 基础实现。
• cpu_inference：支持 CPU 搜索的 ColBERTv2 实现。
• fast_search：带有 PLAID 的 ColBERTv2 实现。
• binarization：使用基线二值化压缩策略的 ColBERT（与 ColBERTv2 的残差压缩相反，我们发现残差压缩更加鲁棒）。

致谢

ColBERT 标志由 Chuyi Zhang 设计。