language:

• en tags:
• ColBERT
• PyLate
• sentence-transformers
• sentence-similarity
• feature-extraction
• generated_from_trainer
• dataset_size:100521
• loss:CachedContrastive base_model: lightonai/GTE-ModernColBERT-v1 datasets:
• reasonir/reasonir-data pipeline_tag: sentence-similarity library_name: PyLate license: cc-by-nc-4.0

Reason-ModernColBERT

Reason-ModernColBERT 是一个基于 reasonir-hq 数据集训练的延迟交互模型。在评估推理密集型检索性能的 BRIGHT 基准测试 中，它展现出极具竞争力的表现，超越了所有现有规模达 7B 的模型（是其大小的 45 倍以上），甚至在 Stack Exchange 分割测试中平均比 ReasonIR-8B（基于相同数据训练的 8B 模型）提升了超过 2.5 个 NDCG@10 点。我们将这一强劲表现归功于延迟交互机制，详见 评估部分。

许可协议

由于 ReasonIR 数据 采用 cc-by-nc-4.0 许可发布，我们无法将此模型以 Apache 2.0 许可发布。不过，ReasonIR 的作者 提供了生成该数据的代码。任何希望复现数据的人都可以通过运行不到 2 小时的微调（使用 此模板），在 Apache 2.0 许可下轻松复现此模型。

基于 lightonai/GTE-ModernColBERT-v1 的 PyLate 模型

这是一个从 lightonai/GTE-ModernColBERT-v1 微调而来的 PyLate 模型，使用 reasonir-hq 数据集训练。它将句子和段落映射为 128 维密集向量序列，可通过 MaxSim 操作符用于语义文本相似性计算。

模型详情

模型描述

• 模型类型: PyLate 模型
• 基础模型: lightonai/GTE-ModernColBERT-v1
• 文档长度: 8192 个标记
• 查询长度: 128 个标记
• 输出维度: 128 个标记
• 相似性函数: MaxSim
• 训练数据集:
  • reasonir-hq
• 语言: en

模型来源

• 文档: PyLate 文档
• 代码库: GitHub 上的 PyLate
• Hugging Face: Hugging Face 上的 PyLate 模型

完整模型架构

ColBERT(
  (0): Transformer({'max_seq_length': 127, 'do_lower_case': False}) with Transformer model: ModernBertModel 
  (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity'})
)

使用方式

首先安装 PyLate 库：

pip install -U pylate

检索

PyLate 提供了简洁的接口，用于使用 ColBERT 模型索引和检索文档。索引利用 Voyager HNSW 索引高效处理文档嵌入并实现快速检索。

文档索引

首先加载 ColBERT 模型并初始化 Voyager 索引，然后编码并索引文档：

from pylate import indexes, models, retrieve

# 步骤 1: 加载 ColBERT 模型
model = models.ColBERT(
    model_name_or_path=pylate_model_id,
)

# 步骤 2: 初始化 Voyager 索引
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # 这将覆盖现有索引（如果有）
)

# 步骤 3: 编码文档
documents_ids = ["1", "2", "3"]
documents = ["文档 1 文本", "文档 2 文本", "文档 3 文本"]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # 确保设置为 False，表示这些是文档而非查询
    show_progress_bar=True,
)

# 步骤 4: 通过提供嵌入和对应 ID 将文档嵌入添加到索引
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

注意，您无需每次都重新创建索引和编码文档。一旦创建索引并添加文档后，您可以通过加载索引重复使用：

# 要加载索引，只需实例化正确的文件夹/名称且不覆盖
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
)

查询的 top-k 文档检索

文档索引完成后，您可以检索给定查询集的最相关 top-k 文档。 为此，使用您要搜索的索引初始化 ColBERT 检索器，编码查询，然后检索 top-k 文档以获取最匹配的 ID 和相关性分数：

# 步骤 1: 初始化 ColBERT 检索器
retriever = retrieve.ColBERT(index=index)

# 步骤 2: 编码查询
queries_embeddings = model.encode(
    ["查询文档 3", "查询文档 1"],
    batch_size=32,
    is_query=True,  # 确保设置为 True，表示这些是查询
    show_progress_bar=True,
)

# 步骤 3: 检索 top-k 文档
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=10,  # 为每个查询检索前 10 个匹配项
)

重排序

如果您只想使用 ColBERT 模型在初步检索管道的基础上进行重排序，而无需构建索引，可以简单地使用 rank 函数并传递查询和待重排序的文档：

from pylate import rank, models

queries = [
    "查询 A",
    "查询 B",
]

documents = [
    ["文档 A", "文档 B"],
    ["文档 1", "文档 C", "文档 B"],
]

documents_ids = [
    [1, 2],
    [1, 3, 2],
]

model = models.ColBERT(
    model_name_or_path=pylate_model_id,
)

queries_embeddings = model.encode(
    queries,
    is_query=True,
)

documents_embeddings = model.encode(
    documents,
    is_query=False,
)

reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)

评估

BRIGHT 基准测试

BRIGHT 基准测试 旨在评估推理密集型检索性能。Reason-ModernColBERT 超越了所有现有规模达 7B 的模型（是其大小的 45 倍以上），甚至在 Stack Exchange 分割测试中平均比 ReasonIR-8B（基于相同数据训练的 8B 模型）提升了超过 2.5 个 NDCG@10 点。我们将这一强劲表现归功于延迟交互机制，相比其他模型通常执行的密集（单向量）检索，如下一节所示。

与密集模型的比较

一个合理的说法是，Reason-ModernColBERT 的性能主要归功于 ReasonIR 数据。尽管 ReasonIR-8B 和 Reason-ModernColBERT 之间的差异已经暗示这不仅仅是数据的问题，我们还是进行了一个小实验，通过使用 Sentence Transformers 训练一个密集（单向量）模型，与使用 PyLate 训练的多向量模型在相同设置下进行比较。这个实验显示出性能上的巨大差距。 显然，需要更严格的实验来得出结论（例如，两个模型都可以进一步调优，训练也可以增强（例如，我们目前没有从其他 GPU 收集负样本，因为 ST 目前不支持）），但差距看起来确实很大，并且与 Reason-ModernColBERT 在比 ReasonIR-8B 小 50 倍以上的情况下仍具有竞争力的情况非常吻合。